adbc_spanner/
connection.rs

1//! The [`SpannerConnection`] — an ADBC connection backed by a Spanner [`DatabaseClient`].
2//!
3//! This driver operates in autocommit mode: every statement runs in its own Spanner transaction
4//! (a single-use read-only transaction for queries, a read/write transaction for DML). Manual
5//! multi-statement transactions are not supported yet, so [`Connection::commit`] and
6//! [`Connection::rollback`] return an error, and disabling autocommit is rejected.
7
8use std::collections::HashSet;
9use std::sync::Arc;
10
11use adbc_core::error::{Result, Status};
12use adbc_core::options::{InfoCode, ObjectDepth, OptionConnection, OptionValue};
13use adbc_core::{Connection, Optionable};
14use arrow_array::{ArrayRef, RecordBatch, RecordBatchIterator, RecordBatchReader, StringArray};
15use arrow_schema::{DataType, Field, Schema};
16use google_cloud_spanner::client::DatabaseClient;
17
18use crate::error::{err, invalid_argument, invalid_state, not_implemented};
19use crate::runtime::SharedRuntime;
20use crate::statement::SpannerStatement;
21
22/// An ADBC connection to a Spanner database.
23pub struct SpannerConnection {
24    runtime: SharedRuntime,
25    client: DatabaseClient,
26    read_only: bool,
27}
28
29impl SpannerConnection {
30    pub(crate) fn new(runtime: SharedRuntime, client: DatabaseClient) -> Self {
31        Self {
32            runtime,
33            client,
34            read_only: false,
35        }
36    }
37}
38
39impl Optionable for SpannerConnection {
40    type Option = OptionConnection;
41
42    fn set_option(&mut self, key: Self::Option, value: OptionValue) -> Result<()> {
43        match &key {
44            OptionConnection::AutoCommit => {
45                if !parse_bool(value)? {
46                    return Err(not_implemented("disabling autocommit"));
47                }
48            }
49            OptionConnection::ReadOnly => self.read_only = parse_bool(value)?,
50            other => {
51                return Err(invalid_argument(format!(
52                    "unsupported Spanner connection option: {}",
53                    connection_option_name(other)
54                )))
55            }
56        }
57        Ok(())
58    }
59
60    fn get_option_string(&self, key: Self::Option) -> Result<String> {
61        match &key {
62            OptionConnection::AutoCommit => Ok(true.to_string()),
63            OptionConnection::ReadOnly => Ok(self.read_only.to_string()),
64            other => Err(err(
65                format!("option {} is not set", connection_option_name(other)),
66                Status::NotFound,
67            )),
68        }
69    }
70
71    fn get_option_bytes(&self, key: Self::Option) -> Result<Vec<u8>> {
72        Ok(self.get_option_string(key)?.into_bytes())
73    }
74
75    fn get_option_int(&self, key: Self::Option) -> Result<i64> {
76        Err(err(
77            format!("option {} is not an integer", connection_option_name(&key)),
78            Status::NotFound,
79        ))
80    }
81
82    fn get_option_double(&self, key: Self::Option) -> Result<f64> {
83        Err(err(
84            format!("option {} is not a double", connection_option_name(&key)),
85            Status::NotFound,
86        ))
87    }
88}
89
90impl Connection for SpannerConnection {
91    type StatementType = SpannerStatement;
92
93    fn new_statement(&mut self) -> Result<Self::StatementType> {
94        Ok(SpannerStatement::new(
95            self.runtime.clone(),
96            self.client.clone(),
97            self.read_only,
98        ))
99    }
100
101    fn cancel(&mut self) -> Result<()> {
102        Err(not_implemented("Connection::cancel"))
103    }
104
105    fn get_info(
106        &self,
107        _codes: Option<HashSet<InfoCode>>,
108    ) -> Result<Box<dyn RecordBatchReader + Send + 'static>> {
109        Err(not_implemented("Connection::get_info"))
110    }
111
112    fn get_objects(
113        &self,
114        _depth: ObjectDepth,
115        _catalog: Option<&str>,
116        _db_schema: Option<&str>,
117        _table_name: Option<&str>,
118        _table_type: Option<Vec<&str>>,
119        _column_name: Option<&str>,
120    ) -> Result<Box<dyn RecordBatchReader + Send + 'static>> {
121        Err(not_implemented("Connection::get_objects"))
122    }
123
124    fn get_table_schema(
125        &self,
126        _catalog: Option<&str>,
127        _db_schema: Option<&str>,
128        _table_name: &str,
129    ) -> Result<Schema> {
130        Err(not_implemented("Connection::get_table_schema"))
131    }
132
133    /// Return the table types supported by Spanner as a single-column (`table_type: utf8`) batch,
134    /// per the ADBC specification.
135    fn get_table_types(&self) -> Result<Box<dyn RecordBatchReader + Send + 'static>> {
136        let schema = Arc::new(Schema::new(vec![Field::new(
137            "table_type",
138            DataType::Utf8,
139            false,
140        )]));
141        let array = Arc::new(StringArray::from(vec!["TABLE", "VIEW"])) as ArrayRef;
142        let batch = RecordBatch::try_new(schema.clone(), vec![array]).map_err(|e| {
143            err(
144                format!("failed to build table types batch: {e}"),
145                Status::Internal,
146            )
147        })?;
148        Ok(Box::new(RecordBatchIterator::new(vec![Ok(batch)], schema)))
149    }
150
151    fn get_statistic_names(&self) -> Result<Box<dyn RecordBatchReader + Send + 'static>> {
152        Err(not_implemented("Connection::get_statistic_names"))
153    }
154
155    fn get_statistics(
156        &self,
157        _catalog: Option<&str>,
158        _db_schema: Option<&str>,
159        _table_name: Option<&str>,
160        _approximate: bool,
161    ) -> Result<Box<dyn RecordBatchReader + Send + 'static>> {
162        Err(not_implemented("Connection::get_statistics"))
163    }
164
165    fn commit(&mut self) -> Result<()> {
166        Err(invalid_state(
167            "connection is in autocommit mode; explicit commit is not supported",
168        ))
169    }
170
171    fn rollback(&mut self) -> Result<()> {
172        Err(invalid_state(
173            "connection is in autocommit mode; explicit rollback is not supported",
174        ))
175    }
176
177    fn read_partition(
178        &self,
179        _partition: impl AsRef<[u8]>,
180    ) -> Result<Box<dyn RecordBatchReader + Send + 'static>> {
181        Err(not_implemented("Connection::read_partition"))
182    }
183}
184
185fn parse_bool(value: OptionValue) -> Result<bool> {
186    match value {
187        OptionValue::String(s) => match s.to_ascii_lowercase().as_str() {
188            "true" | "1" | "yes" => Ok(true),
189            "false" | "0" | "no" => Ok(false),
190            other => Err(invalid_argument(format!(
191                "expected a boolean, got {other:?}"
192            ))),
193        },
194        OptionValue::Int(i) => Ok(i != 0),
195        _ => Err(invalid_argument("expected a boolean option value")),
196    }
197}
198
199fn connection_option_name(key: &OptionConnection) -> String {
200    key.as_ref().to_string()
201}
adbc_spanner/connection.rs

adbc_spanner/
connection.rs
