1use 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
14pub struct SpannerDriver {
19 runtime: SharedRuntime,
20}
21
22impl SpannerDriver {
23 pub fn try_new() -> Result<Self> {
25 Ok(Self {
26 runtime: new_runtime()?,
27 })
28 }
29}
30
31impl Default for SpannerDriver {
32 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
61pub 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 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
205fn 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}