sea_orm_spanner/
database.rs

1use {
2    crate::{error::SpannerDbErr, proxy::SpannerProxy},
3    gcloud_gax::conn::Environment,
4    gcloud_googleapis::spanner::admin::{
5        database::v1::{CreateDatabaseRequest, DatabaseDialect as GrpcDatabaseDialect},
6        instance::v1::{CreateInstanceRequest, Instance},
7    },
8    gcloud_spanner::{
9        admin::{client::Client as AdminClient, AdminClientConfig},
10        client::{Client, ClientConfig},
11    },
12    sea_orm::{Database, DatabaseConnection, DbErr},
13    std::sync::Arc,
14};
15
16/// Database dialect for Spanner
17#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
18pub enum DatabaseDialect {
19    #[default]
20    GoogleStandardSql,
21    PostgreSql,
22}
23
24impl From<DatabaseDialect> for i32 {
25    fn from(dialect: DatabaseDialect) -> Self {
26        match dialect {
27            DatabaseDialect::GoogleStandardSql => GrpcDatabaseDialect::GoogleStandardSql.into(),
28            DatabaseDialect::PostgreSql => GrpcDatabaseDialect::Postgresql.into(),
29        }
30    }
31}
32
33/// Configuration for Spanner instance creation
34#[derive(Debug, Clone, Default)]
35pub struct InstanceConfig {
36    /// Display name for the instance
37    pub display_name: Option<String>,
38    /// Instance configuration (e.g., "regional-us-central1")
39    /// For emulator, this can be empty
40    pub config: Option<String>,
41    /// Number of nodes (for production instances)
42    pub node_count: Option<i32>,
43    /// Processing units (alternative to node_count)
44    pub processing_units: Option<i32>,
45}
46
47/// Options for creating Spanner instance and database
48#[derive(Debug, Clone)]
49pub struct CreateOptions {
50    /// Create instance if it doesn't exist (default: false)
51    ///
52    /// **Warning**: Instance creation can take several minutes and requires
53    /// appropriate IAM permissions. Usually only needed for emulator/testing.
54    pub create_instance_if_not_exists: bool,
55
56    /// Create database if it doesn't exist (default: true)
57    pub create_database_if_not_exists: bool,
58
59    /// Configuration for instance creation
60    pub instance_config: InstanceConfig,
61
62    /// Database dialect (default: GoogleStandardSql)
63    pub database_dialect: DatabaseDialect,
64}
65
66impl Default for CreateOptions {
67    fn default() -> Self {
68        Self {
69            create_instance_if_not_exists: false,
70            create_database_if_not_exists: true,
71            instance_config: InstanceConfig::default(),
72            database_dialect: DatabaseDialect::default(),
73        }
74    }
75}
76
77impl CreateOptions {
78    pub fn new() -> Self {
79        Self::default()
80    }
81
82    pub fn with_instance_creation(mut self) -> Self {
83        self.create_instance_if_not_exists = true;
84        self
85    }
86
87    pub fn with_dialect(mut self, dialect: DatabaseDialect) -> Self {
88        self.database_dialect = dialect;
89        self
90    }
91
92    pub fn with_instance_config(mut self, config: InstanceConfig) -> Self {
93        self.instance_config = config;
94        self
95    }
96}
97
98/// Parsed components of a Spanner database path
99#[derive(Debug, Clone)]
100pub struct DatabasePath {
101    pub project: String,
102    pub instance: String,
103    pub database: String,
104}
105
106impl DatabasePath {
107    /// Parse a database path string into components
108    ///
109    /// Expected format: `projects/{project}/instances/{instance}/databases/{database}`
110    pub fn parse(path: &str) -> Result<Self, DbErr> {
111        let parts: Vec<&str> = path.split('/').collect();
112
113        if parts.len() != 6
114            || parts[0] != "projects"
115            || parts[2] != "instances"
116            || parts[4] != "databases"
117        {
118            return Err(DbErr::Custom(format!(
119                "Invalid database path format. Expected: projects/{{project}}/instances/{{instance}}/databases/{{database}}, got: {}",
120                path
121            )));
122        }
123
124        Ok(Self {
125            project: parts[1].to_string(),
126            instance: parts[3].to_string(),
127            database: parts[5].to_string(),
128        })
129    }
130
131    /// Get the full database path
132    pub fn full_path(&self) -> String {
133        format!(
134            "projects/{}/instances/{}/databases/{}",
135            self.project, self.instance, self.database
136        )
137    }
138
139    /// Get the project path
140    pub fn project_path(&self) -> String {
141        format!("projects/{}", self.project)
142    }
143
144    /// Get the instance path
145    pub fn instance_path(&self) -> String {
146        format!("projects/{}/instances/{}", self.project, self.instance)
147    }
148}
149
150/// Install the rustls crypto provider for GCP TLS connections.
151///
152/// This is called automatically when connecting to GCP (non-emulator).
153/// Safe to call multiple times — subsequent calls are no-ops.
154pub fn ensure_tls() {
155    let _ = rustls::crypto::aws_lc_rs::default_provider().install_default();
156}
157
158pub struct SpannerDatabase;
159
160impl SpannerDatabase {
161    /// Connect to an existing Spanner database
162    ///
163    /// Automatically detects the environment:
164    /// - If `SPANNER_EMULATOR_HOST` is set, connects without authentication
165    /// - Otherwise, uses Application Default Credentials (ADC) for authentication
166    ///
167    /// ADC discovers credentials from (in order):
168    /// 1. `GOOGLE_APPLICATION_CREDENTIALS` environment variable (service account JSON)
169    /// 2. `gcloud auth application-default login` (local development)
170    /// 3. GCE/GKE metadata server (when running on Google Cloud)
171    pub async fn connect(database: &str) -> Result<DatabaseConnection, DbErr> {
172        let config = if std::env::var("SPANNER_EMULATOR_HOST").is_ok() {
173            ClientConfig::default()
174        } else {
175            ensure_tls();
176            ClientConfig::default().with_auth().await.map_err(|e| {
177                SpannerDbErr::Connection(format!("Failed to authenticate with GCP: {}", e))
178            })?
179        };
180        Self::connect_with_config(database, config).await
181    }
182
183    /// Connect to an existing Spanner database with custom configuration
184    pub async fn connect_with_config(
185        database: &str,
186        config: ClientConfig,
187    ) -> Result<DatabaseConnection, DbErr> {
188        let client = Client::new(database, config)
189            .await
190            .map_err(|e| SpannerDbErr::Connection(e.to_string()))?;
191
192        let proxy = SpannerProxy::new(Arc::new(client));
193        Database::connect_proxy(sea_orm::DbBackend::MySql, Arc::new(Box::new(proxy))).await
194    }
195
196    /// Connect to Spanner using the local emulator
197    pub async fn connect_with_emulator(database: &str) -> Result<DatabaseConnection, DbErr> {
198        Self::connect_with_emulator_host(database, "localhost:9010").await
199    }
200
201    /// Connect to Spanner using a custom emulator host
202    pub async fn connect_with_emulator_host(
203        database: &str,
204        emulator_host: &str,
205    ) -> Result<DatabaseConnection, DbErr> {
206        let config = ClientConfig {
207            environment: Environment::Emulator(emulator_host.to_string()),
208            ..Default::default()
209        };
210        Self::connect_with_config(database, config).await
211    }
212
213    /// Connect to Spanner emulator, creating instance and/or database if they don't exist
214    ///
215    /// **Note**: This function only works with the Spanner emulator (localhost:9010).
216    /// It will fail if `SPANNER_EMULATOR_HOST` is not set.
217    ///
218    /// # Example
219    ///
220    /// ```rust,ignore
221    /// use sea_orm_spanner::{SpannerDatabase, CreateOptions};
222    ///
223    /// // Create database if not exists (default behavior)
224    /// let db = SpannerDatabase::connect_or_create_with_emulator(
225    ///     "projects/test/instances/test/databases/test",
226    ///     CreateOptions::default(),
227    /// ).await?;
228    ///
229    /// // Also create instance if not exists
230    /// let db = SpannerDatabase::connect_or_create_with_emulator(
231    ///     "projects/test/instances/test/databases/test",
232    ///     CreateOptions::new().with_instance_creation(),
233    /// ).await?;
234    /// ```
235    pub async fn connect_or_create_with_emulator(
236        database: &str,
237        options: CreateOptions,
238    ) -> Result<DatabaseConnection, DbErr> {
239        Self::connect_or_create_with_emulator_host(database, "localhost:9010", options).await
240    }
241
242    /// Connect to Spanner emulator at custom host, creating instance/database if needed
243    ///
244    /// **Note**: This function only works with the Spanner emulator.
245    pub async fn connect_or_create_with_emulator_host(
246        database: &str,
247        emulator_host: &str,
248        options: CreateOptions,
249    ) -> Result<DatabaseConnection, DbErr> {
250        let path = DatabasePath::parse(database)?;
251
252        if options.create_instance_if_not_exists {
253            ensure_instance(&path, &options.instance_config, emulator_host).await?;
254        }
255
256        if options.create_database_if_not_exists {
257            ensure_database(&path, options.database_dialect, emulator_host).await?;
258        }
259
260        let config = ClientConfig {
261            environment: Environment::Emulator(emulator_host.to_string()),
262            ..Default::default()
263        };
264        Self::connect_with_config(database, config).await
265    }
266}
267
268pub async fn ensure_instance(
269    path: &DatabasePath,
270    config: &InstanceConfig,
271    emulator_host: &str,
272) -> Result<bool, DbErr> {
273    let admin_config = AdminClientConfig {
274        environment: Environment::Emulator(emulator_host.to_string()),
275        ..Default::default()
276    };
277    let admin_client = AdminClient::new(admin_config)
278        .await
279        .map_err(|e| SpannerDbErr::Connection(format!("Failed to create admin client: {}", e)))?;
280
281    let display_name = config
282        .display_name
283        .clone()
284        .unwrap_or_else(|| format!("{} Instance", path.instance));
285
286    let instance_config = config.config.clone().unwrap_or_default();
287
288    let mut instance = Instance {
289        name: path.instance_path(),
290        config: instance_config,
291        display_name,
292        ..Default::default()
293    };
294
295    if let Some(node_count) = config.node_count {
296        instance.node_count = node_count;
297    }
298    if let Some(processing_units) = config.processing_units {
299        instance.processing_units = processing_units;
300    }
301
302    let result = admin_client
303        .instance()
304        .create_instance(
305            CreateInstanceRequest {
306                parent: path.project_path(),
307                instance_id: path.instance.clone(),
308                instance: Some(instance),
309            },
310            None,
311        )
312        .await;
313
314    match result {
315        Ok(mut op) => {
316            op.wait(None).await.map_err(|e| {
317                SpannerDbErr::Connection(format!("Instance creation failed: {}", e))
318            })?;
319            Ok(true)
320        }
321        Err(e) => {
322            let err_str = e.to_string();
323            if err_str.contains("AlreadyExists") || err_str.contains("already exists") {
324                Ok(false)
325            } else {
326                Err(SpannerDbErr::Connection(format!("Failed to create instance: {}", e)).into())
327            }
328        }
329    }
330}
331
332pub async fn ensure_database(
333    path: &DatabasePath,
334    dialect: DatabaseDialect,
335    emulator_host: &str,
336) -> Result<bool, DbErr> {
337    let admin_config = AdminClientConfig {
338        environment: Environment::Emulator(emulator_host.to_string()),
339        ..Default::default()
340    };
341    let admin_client = AdminClient::new(admin_config)
342        .await
343        .map_err(|e| SpannerDbErr::Connection(format!("Failed to create admin client: {}", e)))?;
344
345    let result = admin_client
346        .database()
347        .create_database(
348            CreateDatabaseRequest {
349                parent: path.instance_path(),
350                create_statement: format!("CREATE DATABASE `{}`", path.database),
351                extra_statements: vec![],
352                encryption_config: None,
353                database_dialect: dialect.into(),
354                proto_descriptors: vec![],
355            },
356            None,
357        )
358        .await;
359
360    match result {
361        Ok(mut op) => {
362            op.wait(None).await.map_err(|e| {
363                SpannerDbErr::Connection(format!("Database creation failed: {}", e))
364            })?;
365            Ok(true)
366        }
367        Err(e) => {
368            let err_str = e.to_string();
369            if err_str.contains("AlreadyExists") || err_str.contains("already exists") {
370                Ok(false)
371            } else {
372                Err(SpannerDbErr::Connection(format!("Failed to create database: {}", e)).into())
373            }
374        }
375    }
376}
377
378#[cfg(test)]
379mod tests {
380    use super::*;
381
382    #[test]
383    fn test_database_path_parse() {
384        let path = DatabasePath::parse("projects/my-project/instances/my-instance/databases/my-db")
385            .expect("Should parse valid path");
386
387        assert_eq!(path.project, "my-project");
388        assert_eq!(path.instance, "my-instance");
389        assert_eq!(path.database, "my-db");
390        assert_eq!(path.project_path(), "projects/my-project");
391        assert_eq!(
392            path.instance_path(),
393            "projects/my-project/instances/my-instance"
394        );
395        assert_eq!(
396            path.full_path(),
397            "projects/my-project/instances/my-instance/databases/my-db"
398        );
399    }
400
401    #[test]
402    fn test_database_path_parse_invalid() {
403        assert!(DatabasePath::parse("invalid/path").is_err());
404        assert!(DatabasePath::parse("projects/p/instances/i").is_err());
405        assert!(DatabasePath::parse("").is_err());
406    }
407
408    #[test]
409    fn test_create_options_default() {
410        let options = CreateOptions::default();
411        assert!(!options.create_instance_if_not_exists);
412        assert!(options.create_database_if_not_exists);
413        assert_eq!(options.database_dialect, DatabaseDialect::GoogleStandardSql);
414    }
415}
sea_orm_spanner/database.rs

sea_orm_spanner/
database.rs
