From 2f0a127e0cf93896c3cb49dcbaac81bb006f6393 Mon Sep 17 00:00:00 2001 From: Gleb Novikov Date: Tue, 31 Dec 2024 16:34:33 +0000 Subject: [PATCH] Create neondb database and restore into it --- compute_tools/src/bin/fast_import.rs | 29 ++++++++++++++++++++++++---- 1 file changed, 25 insertions(+), 4 deletions(-) diff --git a/compute_tools/src/bin/fast_import.rs b/compute_tools/src/bin/fast_import.rs index c012826fad..d80910c4cc 100644 --- a/compute_tools/src/bin/fast_import.rs +++ b/compute_tools/src/bin/fast_import.rs @@ -217,16 +217,37 @@ pub(crate) async fn main() -> anyhow::Result<()> { ) .instrument(info_span!("postgres")), ); + + // Create neondb database in the running postgres let restore_pg_connstring = format!("host=localhost port=5432 user={superuser} dbname=postgres"); loop { - let res = tokio_postgres::connect(&restore_pg_connstring, tokio_postgres::NoTls).await; - if res.is_ok() { - info!("postgres is ready, could connect to it"); - break; + match tokio_postgres::connect(&restore_pg_connstring, tokio_postgres::NoTls).await { + Ok((client, connection)) => { + // Spawn the connection handling task to maintain the connection + tokio::spawn(async move { + if let Err(e) = connection.await { + eprintln!("connection error: {}", e); + } + }); + + match client.simple_query("CREATE DATABASE neondb;").await { + Ok(_) => { + info!("created neondb database"); + break; + } + Err(e) => { + info!("failed to create database: {}", e); + break; + } + } + } + Err(_) => continue } } + let restore_pg_connstring = restore_pg_connstring.replace("dbname=postgres", "dbname=neondb"); + let dumpdir = working_directory.join("dumpdir"); let common_args = [