Skip to main content

adbc_spanner/
driver.rs

1//! The [`SpannerDriver`] and [`SpannerDatabase`] — the two top levels of the ADBC hierarchy.
2
3use adbc_core::error::{Result, Status};
4use adbc_core::options::{OptionDatabase, OptionValue};
5use adbc_core::{Database, Driver, Optionable};
6use google_cloud_auth::credentials::anonymous::Builder as AnonymousCredentials;
7use google_cloud_spanner::client::{DatabaseClient, Spanner};
8
9use crate::connection::SpannerConnection;
10use crate::error::{err, from_spanner, invalid_argument, invalid_state};
11use crate::runtime::{new_runtime, SharedRuntime};
12use crate::{OPTION_DATABASE, OPTION_EMULATOR, OPTION_ENDPOINT};
13
14/// The Spanner ADBC driver — the entry point for creating [`SpannerDatabase`] instances.
15///
16/// The driver owns the shared Tokio runtime used to drive the asynchronous Spanner client, so a
17/// single driver instance should be reused for the lifetime of the application.
18pub struct SpannerDriver {
19    runtime: SharedRuntime,
20}
21
22impl SpannerDriver {
23    /// Create a new driver, initialising its Tokio runtime.
24    pub fn try_new() -> Result<Self> {
25        Ok(Self {
26            runtime: new_runtime()?,
27        })
28    }
29}
30
31impl Default for SpannerDriver {
32    /// Create a driver with a fresh runtime.
33    ///
34    /// Required by the C FFI driver exporter, which cannot surface a fallible constructor. Panics
35    /// only if the Tokio runtime cannot be created (catastrophic OS resource exhaustion); prefer
36    /// [`SpannerDriver::try_new`] in Rust code.
37    fn default() -> Self {
38        Self::try_new().expect("failed to initialize the Spanner ADBC driver Tokio runtime")
39    }
40}
41
42impl Driver for SpannerDriver {
43    type DatabaseType = SpannerDatabase;
44
45    fn new_database(&mut self) -> Result<Self::DatabaseType> {
46        Ok(SpannerDatabase::new(self.runtime.clone()))
47    }
48
49    fn new_database_with_opts(
50        &mut self,
51        opts: impl IntoIterator<Item = (OptionDatabase, OptionValue)>,
52    ) -> Result<Self::DatabaseType> {
53        let mut database = SpannerDatabase::new(self.runtime.clone());
54        for (key, value) in opts {
55            database.set_option(key, value)?;
56        }
57        Ok(database)
58    }
59}
60
61/// A configured, but not yet connected, Spanner database.
62///
63/// Holds the connection parameters (the database path and, optionally, an emulator endpoint) and
64/// mints [`SpannerConnection`]s from them.
65pub struct SpannerDatabase {
66    runtime: SharedRuntime,
67    database: Option<String>,
68    endpoint: Option<String>,
69    emulator: bool,
70}
71
72impl SpannerDatabase {
73    pub(crate) fn new(runtime: SharedRuntime) -> Self {
74        Self {
75            runtime,
76            database: None,
77            endpoint: None,
78            emulator: false,
79        }
80    }
81
82    /// Resolve the effective configuration and build a connected [`DatabaseClient`].
83    ///
84    /// Emulator handling: if `SPANNER_EMULATOR_HOST` is set it supplies the endpoint (unless one was
85    /// given explicitly) and forces anonymous credentials.
86    pub(crate) fn connect(&self) -> Result<DatabaseClient> {
87        let database = self.database.clone().ok_or_else(|| {
88            invalid_state(
89                "Spanner database path is not set; provide the `uri` or \
90                 `adbc.spanner.database` option (projects/<p>/instances/<i>/databases/<d>)",
91            )
92        })?;
93
94        let mut endpoint = self.endpoint.clone();
95        let mut emulator = self.emulator;
96        if let Ok(host) = std::env::var("SPANNER_EMULATOR_HOST") {
97            if !host.is_empty() {
98                if endpoint.is_none() {
99                    endpoint = Some(ensure_scheme(&host));
100                }
101                emulator = true;
102            }
103        }
104
105        self.runtime.block_on(async move {
106            let mut builder = Spanner::builder();
107            if let Some(endpoint) = endpoint {
108                builder = builder.with_endpoint(endpoint);
109            }
110            if emulator {
111                builder = builder.with_credentials(AnonymousCredentials::new().build());
112            }
113            let spanner = builder.build().await.map_err(from_spanner)?;
114            spanner
115                .database_client(database)
116                .build()
117                .await
118                .map_err(from_spanner)
119        })
120    }
121}
122
123impl Optionable for SpannerDatabase {
124    type Option = OptionDatabase;
125
126    fn set_option(&mut self, key: Self::Option, value: OptionValue) -> Result<()> {
127        match &key {
128            OptionDatabase::Uri => self.database = Some(string_value(&key, value)?),
129            OptionDatabase::Other(name) if name == OPTION_DATABASE => {
130                self.database = Some(string_value(&key, value)?)
131            }
132            OptionDatabase::Other(name) if name == OPTION_ENDPOINT => {
133                self.endpoint = Some(string_value(&key, value)?)
134            }
135            OptionDatabase::Other(name) if name == OPTION_EMULATOR => {
136                self.emulator = bool_value(&key, value)?
137            }
138            other => {
139                return Err(invalid_argument(format!(
140                    "unsupported Spanner database option: {}",
141                    option_name(other)
142                )))
143            }
144        }
145        Ok(())
146    }
147
148    fn get_option_string(&self, key: Self::Option) -> Result<String> {
149        let value = match &key {
150            OptionDatabase::Uri => self.database.clone(),
151            OptionDatabase::Other(name) if name == OPTION_DATABASE => self.database.clone(),
152            OptionDatabase::Other(name) if name == OPTION_ENDPOINT => self.endpoint.clone(),
153            OptionDatabase::Other(name) if name == OPTION_EMULATOR => {
154                Some(self.emulator.to_string())
155            }
156            _ => None,
157        };
158        value.ok_or_else(|| {
159            err(
160                format!("option {} is not set", option_name(&key)),
161                Status::NotFound,
162            )
163        })
164    }
165
166    fn get_option_bytes(&self, key: Self::Option) -> Result<Vec<u8>> {
167        Ok(self.get_option_string(key)?.into_bytes())
168    }
169
170    fn get_option_int(&self, key: Self::Option) -> Result<i64> {
171        Err(err(
172            format!("option {} is not an integer", option_name(&key)),
173            Status::NotFound,
174        ))
175    }
176
177    fn get_option_double(&self, key: Self::Option) -> Result<f64> {
178        Err(err(
179            format!("option {} is not a double", option_name(&key)),
180            Status::NotFound,
181        ))
182    }
183}
184
185impl Database for SpannerDatabase {
186    type ConnectionType = SpannerConnection;
187
188    fn new_connection(&self) -> Result<Self::ConnectionType> {
189        let client = self.connect()?;
190        Ok(SpannerConnection::new(self.runtime.clone(), client))
191    }
192
193    fn new_connection_with_opts(
194        &self,
195        opts: impl IntoIterator<Item = (adbc_core::options::OptionConnection, OptionValue)>,
196    ) -> Result<Self::ConnectionType> {
197        let mut connection = self.new_connection()?;
198        for (key, value) in opts {
199            connection.set_option(key, value)?;
200        }
201        Ok(connection)
202    }
203}
204
205/// Prefix a bare `host:port` emulator address with an `http://` scheme, as expected by the gRPC
206/// transport.
207fn ensure_scheme(host: &str) -> String {
208    if host.starts_with("http://") || host.starts_with("https://") {
209        host.to_string()
210    } else {
211        format!("http://{host}")
212    }
213}
214
215fn option_name(key: &OptionDatabase) -> String {
216    key.as_ref().to_string()
217}
218
219fn string_value(key: &OptionDatabase, value: OptionValue) -> Result<String> {
220    match value {
221        OptionValue::String(s) => Ok(s),
222        _ => Err(invalid_argument(format!(
223            "option {} requires a string value",
224            option_name(key)
225        ))),
226    }
227}
228
229fn bool_value(key: &OptionDatabase, value: OptionValue) -> Result<bool> {
230    match value {
231        OptionValue::String(s) => match s.to_ascii_lowercase().as_str() {
232            "true" | "1" | "yes" => Ok(true),
233            "false" | "0" | "no" => Ok(false),
234            other => Err(invalid_argument(format!(
235                "option {} expects a boolean, got {other:?}",
236                option_name(key)
237            ))),
238        },
239        OptionValue::Int(i) => Ok(i != 0),
240        _ => Err(invalid_argument(format!(
241            "option {} requires a boolean value",
242            option_name(key)
243        ))),
244    }
245}
246
247#[cfg(test)]
248mod tests {
249    use super::*;
250    use adbc_core::error::Status;
251
252    fn new_database() -> SpannerDatabase {
253        SpannerDatabase::new(new_runtime().unwrap())
254    }
255
256    #[test]
257    fn ensure_scheme_adds_http_prefix() {
258        assert_eq!(ensure_scheme("localhost:9010"), "http://localhost:9010");
259        assert_eq!(ensure_scheme("http://host:1"), "http://host:1");
260        assert_eq!(ensure_scheme("https://host:1"), "https://host:1");
261    }
262
263    #[test]
264    fn database_options_round_trip() {
265        let mut db = new_database();
266        db.set_option(
267            OptionDatabase::Uri,
268            OptionValue::String("projects/p/instances/i/databases/d".into()),
269        )
270        .unwrap();
271        db.set_option(
272            OptionDatabase::Other(OPTION_ENDPOINT.into()),
273            OptionValue::String("http://localhost:9010".into()),
274        )
275        .unwrap();
276        db.set_option(
277            OptionDatabase::Other(OPTION_EMULATOR.into()),
278            OptionValue::String("true".into()),
279        )
280        .unwrap();
281
282        assert_eq!(
283            db.get_option_string(OptionDatabase::Uri).unwrap(),
284            "projects/p/instances/i/databases/d"
285        );
286        assert_eq!(
287            db.get_option_string(OptionDatabase::Other(OPTION_ENDPOINT.into()))
288                .unwrap(),
289            "http://localhost:9010"
290        );
291        assert!(db.emulator);
292    }
293
294    #[test]
295    fn the_database_option_is_an_alias_for_uri() {
296        let mut db = new_database();
297        db.set_option(
298            OptionDatabase::Other(OPTION_DATABASE.into()),
299            OptionValue::String("projects/p/instances/i/databases/d".into()),
300        )
301        .unwrap();
302        assert_eq!(
303            db.get_option_string(OptionDatabase::Uri).unwrap(),
304            "projects/p/instances/i/databases/d"
305        );
306    }
307
308    #[test]
309    fn connecting_without_a_database_path_is_an_error() {
310        let db = new_database();
311        let error = db.connect().unwrap_err();
312        assert_eq!(error.status, Status::InvalidState);
313    }
314
315    #[test]
316    fn a_non_string_uri_is_rejected() {
317        let mut db = new_database();
318        let error = db
319            .set_option(OptionDatabase::Uri, OptionValue::Int(42))
320            .unwrap_err();
321        assert_eq!(error.status, Status::InvalidArguments);
322    }
323}