Merge pull request #4670 from driftluo/fix-pg-init

fix: fix pg init
nervosnetwork · Oct 8, 2024 · 28fd446 · 28fd446
2 parents 762fd06 + 6677845
commit 28fd446
Showing 1 changed file with 25 additions and 3 deletions.
diff --git a/util/rich-indexer/src/store.rs b/util/rich-indexer/src/store.rs
@@ -70,6 +70,7 @@ impl SQLXPool {
                 pool
             }
             DBDriver::Postgres => {
+                self.postgres_init(db_config).await?;
                 let uri = build_url_for_postgres(db_config);
                 let connection_options =
                     AnyConnectOptions::from_str(&uri)?.log_statements(LevelFilter::Trace);
@@ -79,9 +80,6 @@ impl SQLXPool {
                     .set(pool.clone())
                     .map_err(|_| anyhow!("set pool failed"))?;
 
-                SQLXPool::new_query(r#"CREATE DATABASE IF NOT EXISTS "postgres""#)
-                    .execute(&pool)
-                    .await?;
                 self.create_tables_for_postgres().await?;
 
                 self.db_driver = DBDriver::Postgres;
@@ -207,6 +205,30 @@ impl SQLXPool {
         }
         Ok(())
     }
+
+    pub async fn postgres_init(&mut self, db_config: &RichIndexerConfig) -> Result<()> {
+        // Connect to the "postgres" database first
+        let mut temp_config = db_config.clone();
+        temp_config.db_name = "postgres".to_string();
+        let uri = build_url_for_postgres(&temp_config);
+        let connection_options =
+            AnyConnectOptions::from_str(&uri)?.log_statements(LevelFilter::Trace);
+        let tmp_pool_options = AnyPoolOptions::new();
+        let pool = tmp_pool_options.connect_with(connection_options).await?;
+        // Check if database exists
+        let query =
+            SQLXPool::new_query(r#"SELECT EXISTS (SELECT FROM pg_database WHERE datname = $1)"#)
+                .bind(db_config.db_name.as_str());
+        let row = query.fetch_one(&pool).await?;
+        // If database does not exist, create it
+        if !row.get::<bool, _>(0) {
+            let query = format!(r#"CREATE DATABASE "{}""#, db_config.db_name);
+            SQLXPool::new_query(&query).execute(&pool).await?;
+            Ok(())
+        } else {
+            Ok(())
+        }
+    }
 }
 
 fn build_url_for_sqlite(db_config: &RichIndexerConfig) -> String {