adbc_spanner/
statement.rs

1//! The [`SpannerStatement`] — an ADBC statement that runs SQL against Spanner and returns Arrow.
2//!
3//! A statement holds a SQL string set via [`Statement::set_sql_query`]. Calling
4//! [`Statement::execute`] runs it as a query in a single-use read-only transaction and streams the
5//! result back as an Arrow [`RecordBatch`]. Calling [`Statement::execute_update`] runs it as DML
6//! inside a read/write transaction and returns the number of affected rows.
7
8use adbc_core::error::{Error, Result, Status};
9use adbc_core::options::{OptionStatement, OptionValue};
10use adbc_core::{Optionable, PartitionedResult, Statement};
11use arrow_array::{RecordBatch, RecordBatchIterator, RecordBatchReader};
12use arrow_schema::Schema;
13use google_cloud_spanner::client::DatabaseClient;
14use google_cloud_spanner::statement::Statement as SpannerSql;
15use google_cloud_spanner::transaction::ReadWriteTransaction;
16
17use crate::conversion::result_set_to_batch;
18use crate::error::{err, from_spanner, invalid_state, not_implemented};
19use crate::runtime::SharedRuntime;
20
21/// An ADBC statement bound to a Spanner [`DatabaseClient`].
22pub struct SpannerStatement {
23    runtime: SharedRuntime,
24    client: DatabaseClient,
25    read_only: bool,
26    sql: Option<String>,
27}
28
29impl SpannerStatement {
30    pub(crate) fn new(runtime: SharedRuntime, client: DatabaseClient, read_only: bool) -> Self {
31        Self {
32            runtime,
33            client,
34            read_only,
35            sql: None,
36        }
37    }
38
39    fn sql(&self) -> Result<String> {
40        self.sql
41            .clone()
42            .ok_or_else(|| invalid_state("no SQL query set on statement; call set_sql_query first"))
43    }
44}
45
46impl Optionable for SpannerStatement {
47    type Option = OptionStatement;
48
49    fn set_option(&mut self, key: Self::Option, _value: OptionValue) -> Result<()> {
50        Err(not_implemented(&format!(
51            "statement option {}",
52            key.as_ref()
53        )))
54    }
55
56    fn get_option_string(&self, key: Self::Option) -> Result<String> {
57        Err(err(
58            format!("option {} is not set", key.as_ref()),
59            Status::NotFound,
60        ))
61    }
62
63    fn get_option_bytes(&self, key: Self::Option) -> Result<Vec<u8>> {
64        Err(err(
65            format!("option {} is not set", key.as_ref()),
66            Status::NotFound,
67        ))
68    }
69
70    fn get_option_int(&self, key: Self::Option) -> Result<i64> {
71        Err(err(
72            format!("option {} is not set", key.as_ref()),
73            Status::NotFound,
74        ))
75    }
76
77    fn get_option_double(&self, key: Self::Option) -> Result<f64> {
78        Err(err(
79            format!("option {} is not set", key.as_ref()),
80            Status::NotFound,
81        ))
82    }
83}
84
85impl Statement for SpannerStatement {
86    fn bind(&mut self, _batch: RecordBatch) -> Result<()> {
87        Err(not_implemented("Statement::bind (parameter binding)"))
88    }
89
90    fn bind_stream(&mut self, _reader: Box<dyn RecordBatchReader + Send>) -> Result<()> {
91        Err(not_implemented("Statement::bind_stream (bulk ingest)"))
92    }
93
94    fn execute(&mut self) -> Result<Box<dyn RecordBatchReader + Send + 'static>> {
95        let sql = self.sql()?;
96        let client = self.client.clone();
97        let (schema, batch) = self.runtime.block_on(async move {
98            let transaction = client.single_use().build();
99            let statement = SpannerSql::builder(sql).build();
100            let result_set = transaction
101                .execute_query(statement)
102                .await
103                .map_err(from_spanner)?;
104            result_set_to_batch(result_set).await
105        })?;
106        Ok(Box::new(RecordBatchIterator::new(vec![Ok(batch)], schema)))
107    }
108
109    fn execute_update(&mut self) -> Result<Option<i64>> {
110        if self.read_only {
111            return Err(invalid_state(
112                "cannot execute DML: the connection is read-only",
113            ));
114        }
115        let sql = self.sql()?;
116        let client = self.client.clone();
117        let affected = self.runtime.block_on(async move {
118            let runner = client
119                .read_write_transaction()
120                .build()
121                .await
122                .map_err(from_spanner)?;
123            // The runner may retry the closure if Spanner aborts the transaction, so rebuild the
124            // statement from the (cloned) SQL on each attempt.
125            let outcome = runner
126                .run(move |transaction: ReadWriteTransaction| {
127                    let sql = sql.clone();
128                    async move {
129                        let statement = SpannerSql::builder(sql).build();
130                        transaction.execute_update(statement).await
131                    }
132                })
133                .await
134                .map_err(from_spanner)?;
135            Ok::<i64, Error>(outcome.result)
136        })?;
137        Ok(Some(affected))
138    }
139
140    fn execute_schema(&mut self) -> Result<Schema> {
141        Err(not_implemented("Statement::execute_schema"))
142    }
143
144    fn execute_partitions(&mut self) -> Result<PartitionedResult> {
145        Err(not_implemented("Statement::execute_partitions"))
146    }
147
148    fn get_parameter_schema(&self) -> Result<Schema> {
149        Err(not_implemented("Statement::get_parameter_schema"))
150    }
151
152    fn prepare(&mut self) -> Result<()> {
153        // Spanner prepares/plans statements server-side on execution, so this is a no-op.
154        Ok(())
155    }
156
157    fn set_sql_query(&mut self, query: impl AsRef<str>) -> Result<()> {
158        self.sql = Some(query.as_ref().to_string());
159        Ok(())
160    }
161
162    fn set_substrait_plan(&mut self, _plan: impl AsRef<[u8]>) -> Result<()> {
163        Err(not_implemented("Statement::set_substrait_plan"))
164    }
165
166    fn cancel(&mut self) -> Result<()> {
167        Err(not_implemented("Statement::cancel"))
168    }
169}
adbc_spanner/statement.rs

adbc_spanner/
statement.rs
