LCOV - code coverage report
Current view: top level - compute_tools/src/bin - fast_import.rs (source / functions) Coverage Total Hit
Test: 1b0a6a0c05cee5a7de360813c8034804e105ce1c.info Lines: 0.0 % 476 0
Test Date: 2025-03-12 00:01:28 Functions: 0.0 % 46 0

            Line data    Source code
       1              : //! This program dumps a remote Postgres database into a local Postgres database
       2              : //! and uploads the resulting PGDATA into object storage for import into a Timeline.
       3              : //!
       4              : //! # Context, Architecture, Design
       5              : //!
       6              : //! See cloud.git Fast Imports RFC (<https://github.com/neondatabase/cloud/pull/19799>)
       7              : //! for the full picture.
       8              : //! The RFC describing the storage pieces of importing the PGDATA dump into a Timeline
       9              : //! is publicly accessible at <https://github.com/neondatabase/neon/pull/9538>.
      10              : //!
      11              : //! # This is a Prototype!
      12              : //!
      13              : //! This program is part of a prototype feature and not yet used in production.
      14              : //!
      15              : //! The cloud.git RFC contains lots of suggestions for improving e2e throughput
      16              : //! of this step of the timeline import process.
      17              : //!
      18              : //! # Local Testing
      19              : //!
      20              : //! - Comment out most of the pgxns in compute-node.Dockerfile to speed up the build.
      21              : //! - Build the image with the following command:
      22              : //!
      23              : //! ```bash
      24              : //! docker buildx build --platform linux/amd64 --build-arg DEBIAN_VERSION=bullseye --build-arg GIT_VERSION=local --build-arg PG_VERSION=v14 --build-arg BUILD_TAG="$(date --iso-8601=s -u)" -t localhost:3030/localregistry/compute-node-v14:latest -f compute/compute-node.Dockerfile .
      25              : //! docker push localhost:3030/localregistry/compute-node-v14:latest
      26              : //! ```
      27              : 
      28              : use anyhow::{Context, bail};
      29              : use aws_config::BehaviorVersion;
      30              : use camino::{Utf8Path, Utf8PathBuf};
      31              : use clap::{Parser, Subcommand};
      32              : use compute_tools::extension_server::{PostgresMajorVersion, get_pg_version};
      33              : use nix::unistd::Pid;
      34              : use tracing::{Instrument, error, info, info_span, warn};
      35              : use utils::fs_ext::is_directory_empty;
      36              : 
      37              : #[path = "fast_import/aws_s3_sync.rs"]
      38              : mod aws_s3_sync;
      39              : #[path = "fast_import/child_stdio_to_log.rs"]
      40              : mod child_stdio_to_log;
      41              : #[path = "fast_import/s3_uri.rs"]
      42              : mod s3_uri;
      43              : 
      44              : const PG_WAIT_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(600);
      45              : const PG_WAIT_RETRY_INTERVAL: std::time::Duration = std::time::Duration::from_millis(300);
      46              : 
      47              : #[derive(Subcommand, Debug)]
      48              : enum Command {
      49              :     /// Runs local postgres (neon binary), restores into it,
      50              :     /// uploads pgdata to s3 to be consumed by pageservers
      51              :     Pgdata {
      52              :         /// Raw connection string to the source database. Used only in tests,
      53              :         /// real scenario uses encrypted connection string in spec.json from s3.
      54              :         #[clap(long)]
      55              :         source_connection_string: Option<String>,
      56              :         /// If specified, will not shut down the local postgres after the import. Used in local testing
      57              :         #[clap(short, long)]
      58            0 :         interactive: bool,
      59              :         /// Port to run postgres on. Default is 5432.
      60            0 :         #[clap(long, default_value_t = 5432)]
      61            0 :         pg_port: u16, // port to run postgres on, 5432 is default
      62              : 
      63              :         /// Number of CPUs in the system. This is used to configure # of
      64              :         /// parallel worker processes, for index creation.
      65              :         #[clap(long, env = "NEON_IMPORTER_NUM_CPUS")]
      66              :         num_cpus: Option<usize>,
      67              : 
      68              :         /// Amount of RAM in the system. This is used to configure shared_buffers
      69              :         /// and maintenance_work_mem.
      70              :         #[clap(long, env = "NEON_IMPORTER_MEMORY_MB")]
      71              :         memory_mb: Option<usize>,
      72              :     },
      73              : 
      74              :     /// Runs pg_dump-pg_restore from source to destination without running local postgres.
      75              :     DumpRestore {
      76              :         /// Raw connection string to the source database. Used only in tests,
      77              :         /// real scenario uses encrypted connection string in spec.json from s3.
      78              :         #[clap(long)]
      79              :         source_connection_string: Option<String>,
      80              :         /// Raw connection string to the destination database. Used only in tests,
      81              :         /// real scenario uses encrypted connection string in spec.json from s3.
      82              :         #[clap(long)]
      83              :         destination_connection_string: Option<String>,
      84              :     },
      85              : }
      86              : 
      87              : #[derive(clap::Parser)]
      88              : struct Args {
      89              :     #[clap(long, env = "NEON_IMPORTER_WORKDIR")]
      90            0 :     working_directory: Utf8PathBuf,
      91              :     #[clap(long, env = "NEON_IMPORTER_S3_PREFIX")]
      92              :     s3_prefix: Option<s3_uri::S3Uri>,
      93              :     #[clap(long, env = "NEON_IMPORTER_PG_BIN_DIR")]
      94            0 :     pg_bin_dir: Utf8PathBuf,
      95              :     #[clap(long, env = "NEON_IMPORTER_PG_LIB_DIR")]
      96            0 :     pg_lib_dir: Utf8PathBuf,
      97              : 
      98              :     #[clap(subcommand)]
      99              :     command: Command,
     100              : }
     101              : 
     102              : #[serde_with::serde_as]
     103            0 : #[derive(serde::Deserialize)]
     104              : struct Spec {
     105              :     encryption_secret: EncryptionSecret,
     106              :     #[serde_as(as = "serde_with::base64::Base64")]
     107              :     source_connstring_ciphertext_base64: Vec<u8>,
     108              :     #[serde_as(as = "Option<serde_with::base64::Base64>")]
     109              :     destination_connstring_ciphertext_base64: Option<Vec<u8>>,
     110              : }
     111              : 
     112            0 : #[derive(serde::Deserialize)]
     113              : enum EncryptionSecret {
     114              :     #[allow(clippy::upper_case_acronyms)]
     115              :     KMS { key_id: String },
     116              : }
     117              : 
     118              : // copied from pageserver_api::config::defaults::DEFAULT_LOCALE to avoid dependency just for a constant
     119              : const DEFAULT_LOCALE: &str = if cfg!(target_os = "macos") {
     120              :     "C"
     121              : } else {
     122              :     "C.UTF-8"
     123              : };
     124              : 
     125            0 : async fn decode_connstring(
     126            0 :     kms_client: &aws_sdk_kms::Client,
     127            0 :     key_id: &String,
     128            0 :     connstring_ciphertext_base64: Vec<u8>,
     129            0 : ) -> Result<String, anyhow::Error> {
     130            0 :     let mut output = kms_client
     131            0 :         .decrypt()
     132            0 :         .key_id(key_id)
     133            0 :         .ciphertext_blob(aws_sdk_s3::primitives::Blob::new(
     134            0 :             connstring_ciphertext_base64,
     135            0 :         ))
     136            0 :         .send()
     137            0 :         .await
     138            0 :         .context("decrypt connection string")?;
     139              : 
     140            0 :     let plaintext = output
     141            0 :         .plaintext
     142            0 :         .take()
     143            0 :         .context("get plaintext connection string")?;
     144              : 
     145            0 :     String::from_utf8(plaintext.into_inner()).context("parse connection string as utf8")
     146            0 : }
     147              : 
     148              : struct PostgresProcess {
     149              :     pgdata_dir: Utf8PathBuf,
     150              :     pg_bin_dir: Utf8PathBuf,
     151              :     pgbin: Utf8PathBuf,
     152              :     pg_lib_dir: Utf8PathBuf,
     153              :     postgres_proc: Option<tokio::process::Child>,
     154              : }
     155              : 
     156              : impl PostgresProcess {
     157            0 :     fn new(pgdata_dir: Utf8PathBuf, pg_bin_dir: Utf8PathBuf, pg_lib_dir: Utf8PathBuf) -> Self {
     158            0 :         Self {
     159            0 :             pgdata_dir,
     160            0 :             pgbin: pg_bin_dir.join("postgres"),
     161            0 :             pg_bin_dir,
     162            0 :             pg_lib_dir,
     163            0 :             postgres_proc: None,
     164            0 :         }
     165            0 :     }
     166              : 
     167            0 :     async fn prepare(&self, initdb_user: &str) -> Result<(), anyhow::Error> {
     168            0 :         tokio::fs::create_dir(&self.pgdata_dir)
     169            0 :             .await
     170            0 :             .context("create pgdata directory")?;
     171              : 
     172            0 :         let pg_version = match get_pg_version(self.pgbin.as_ref()) {
     173            0 :             PostgresMajorVersion::V14 => 14,
     174            0 :             PostgresMajorVersion::V15 => 15,
     175            0 :             PostgresMajorVersion::V16 => 16,
     176            0 :             PostgresMajorVersion::V17 => 17,
     177              :         };
     178            0 :         postgres_initdb::do_run_initdb(postgres_initdb::RunInitdbArgs {
     179            0 :             superuser: initdb_user,
     180            0 :             locale: DEFAULT_LOCALE, // XXX: this shouldn't be hard-coded,
     181            0 :             pg_version,
     182            0 :             initdb_bin: self.pg_bin_dir.join("initdb").as_ref(),
     183            0 :             library_search_path: &self.pg_lib_dir, // TODO: is this right? Prob works in compute image, not sure about neon_local.
     184            0 :             pgdata: &self.pgdata_dir,
     185            0 :         })
     186            0 :         .await
     187            0 :         .context("initdb")
     188            0 :     }
     189              : 
     190            0 :     async fn start(
     191            0 :         &mut self,
     192            0 :         initdb_user: &str,
     193            0 :         port: u16,
     194            0 :         nproc: usize,
     195            0 :         memory_mb: usize,
     196            0 :     ) -> Result<&tokio::process::Child, anyhow::Error> {
     197            0 :         self.prepare(initdb_user).await?;
     198              : 
     199              :         // Somewhat arbitrarily, use 10 % of memory for shared buffer cache, 70% for
     200              :         // maintenance_work_mem (i.e. for sorting during index creation), and leave the rest
     201              :         // available for misc other stuff that PostgreSQL uses memory for.
     202            0 :         let shared_buffers_mb = ((memory_mb as f32) * 0.10) as usize;
     203            0 :         let maintenance_work_mem_mb = ((memory_mb as f32) * 0.70) as usize;
     204              : 
     205              :         //
     206              :         // Launch postgres process
     207              :         //
     208            0 :         let mut proc = tokio::process::Command::new(&self.pgbin)
     209            0 :             .arg("-D")
     210            0 :             .arg(&self.pgdata_dir)
     211            0 :             .args(["-p", &format!("{port}")])
     212            0 :             .args(["-c", "wal_level=minimal"])
     213            0 :             .args(["-c", &format!("shared_buffers={shared_buffers_mb}MB")])
     214            0 :             .args(["-c", "max_wal_senders=0"])
     215            0 :             .args(["-c", "fsync=off"])
     216            0 :             .args(["-c", "full_page_writes=off"])
     217            0 :             .args(["-c", "synchronous_commit=off"])
     218            0 :             .args([
     219            0 :                 "-c",
     220            0 :                 &format!("maintenance_work_mem={maintenance_work_mem_mb}MB"),
     221            0 :             ])
     222            0 :             .args(["-c", &format!("max_parallel_maintenance_workers={nproc}")])
     223            0 :             .args(["-c", &format!("max_parallel_workers={nproc}")])
     224            0 :             .args(["-c", &format!("max_parallel_workers_per_gather={nproc}")])
     225            0 :             .args(["-c", &format!("max_worker_processes={nproc}")])
     226            0 :             .args(["-c", "effective_io_concurrency=100"])
     227            0 :             .env_clear()
     228            0 :             .env("LD_LIBRARY_PATH", &self.pg_lib_dir)
     229            0 :             .env(
     230            0 :                 "ASAN_OPTIONS",
     231            0 :                 std::env::var("ASAN_OPTIONS").unwrap_or_default(),
     232            0 :             )
     233            0 :             .env(
     234            0 :                 "UBSAN_OPTIONS",
     235            0 :                 std::env::var("UBSAN_OPTIONS").unwrap_or_default(),
     236            0 :             )
     237            0 :             .stdout(std::process::Stdio::piped())
     238            0 :             .stderr(std::process::Stdio::piped())
     239            0 :             .spawn()
     240            0 :             .context("spawn postgres")?;
     241              : 
     242            0 :         info!("spawned postgres, waiting for it to become ready");
     243              :         tokio::spawn(
     244            0 :             child_stdio_to_log::relay_process_output(proc.stdout.take(), proc.stderr.take())
     245            0 :                 .instrument(info_span!("postgres")),
     246              :         );
     247              : 
     248            0 :         self.postgres_proc = Some(proc);
     249            0 :         Ok(self.postgres_proc.as_ref().unwrap())
     250            0 :     }
     251              : 
     252            0 :     async fn shutdown(&mut self) -> Result<(), anyhow::Error> {
     253            0 :         let proc: &mut tokio::process::Child = self.postgres_proc.as_mut().unwrap();
     254            0 :         info!("shutdown postgres");
     255            0 :         nix::sys::signal::kill(
     256            0 :             Pid::from_raw(i32::try_from(proc.id().unwrap()).expect("convert child pid to i32")),
     257            0 :             nix::sys::signal::SIGTERM,
     258            0 :         )
     259            0 :         .context("signal postgres to shut down")?;
     260            0 :         proc.wait()
     261            0 :             .await
     262            0 :             .context("wait for postgres to shut down")
     263            0 :             .map(|_| ())
     264            0 :     }
     265              : }
     266              : 
     267            0 : async fn wait_until_ready(connstring: String, create_dbname: String) {
     268            0 :     // Create neondb database in the running postgres
     269            0 :     let start_time = std::time::Instant::now();
     270              : 
     271              :     loop {
     272            0 :         if start_time.elapsed() > PG_WAIT_TIMEOUT {
     273            0 :             error!(
     274            0 :                 "timeout exceeded: failed to poll postgres and create database within 10 minutes"
     275              :             );
     276            0 :             std::process::exit(1);
     277            0 :         }
     278            0 : 
     279            0 :         match tokio_postgres::connect(
     280            0 :             &connstring.replace("dbname=neondb", "dbname=postgres"),
     281            0 :             tokio_postgres::NoTls,
     282            0 :         )
     283            0 :         .await
     284              :         {
     285            0 :             Ok((client, connection)) => {
     286            0 :                 // Spawn the connection handling task to maintain the connection
     287            0 :                 tokio::spawn(async move {
     288            0 :                     if let Err(e) = connection.await {
     289            0 :                         warn!("connection error: {}", e);
     290            0 :                     }
     291            0 :                 });
     292            0 : 
     293            0 :                 match client
     294            0 :                     .simple_query(format!("CREATE DATABASE {create_dbname};").as_str())
     295            0 :                     .await
     296              :                 {
     297              :                     Ok(_) => {
     298            0 :                         info!("created {} database", create_dbname);
     299            0 :                         break;
     300              :                     }
     301            0 :                     Err(e) => {
     302            0 :                         warn!(
     303            0 :                             "failed to create database: {}, retying in {}s",
     304            0 :                             e,
     305            0 :                             PG_WAIT_RETRY_INTERVAL.as_secs_f32()
     306              :                         );
     307            0 :                         tokio::time::sleep(PG_WAIT_RETRY_INTERVAL).await;
     308            0 :                         continue;
     309              :                     }
     310              :                 }
     311              :             }
     312              :             Err(_) => {
     313            0 :                 info!(
     314            0 :                     "postgres not ready yet, retrying in {}s",
     315            0 :                     PG_WAIT_RETRY_INTERVAL.as_secs_f32()
     316              :                 );
     317            0 :                 tokio::time::sleep(PG_WAIT_RETRY_INTERVAL).await;
     318            0 :                 continue;
     319              :             }
     320              :         }
     321              :     }
     322            0 : }
     323              : 
     324            0 : async fn run_dump_restore(
     325            0 :     workdir: Utf8PathBuf,
     326            0 :     pg_bin_dir: Utf8PathBuf,
     327            0 :     pg_lib_dir: Utf8PathBuf,
     328            0 :     source_connstring: String,
     329            0 :     destination_connstring: String,
     330            0 : ) -> Result<(), anyhow::Error> {
     331            0 :     let dumpdir = workdir.join("dumpdir");
     332            0 : 
     333            0 :     let common_args = [
     334            0 :         // schema mapping (prob suffices to specify them on one side)
     335            0 :         "--no-owner".to_string(),
     336            0 :         "--no-privileges".to_string(),
     337            0 :         "--no-publications".to_string(),
     338            0 :         "--no-security-labels".to_string(),
     339            0 :         "--no-subscriptions".to_string(),
     340            0 :         "--no-tablespaces".to_string(),
     341            0 :         // format
     342            0 :         "--format".to_string(),
     343            0 :         "directory".to_string(),
     344            0 :         // concurrency
     345            0 :         "--jobs".to_string(),
     346            0 :         num_cpus::get().to_string(),
     347            0 :         // progress updates
     348            0 :         "--verbose".to_string(),
     349            0 :     ];
     350            0 : 
     351            0 :     info!("dump into the working directory");
     352              :     {
     353            0 :         let mut pg_dump = tokio::process::Command::new(pg_bin_dir.join("pg_dump"))
     354            0 :             .args(&common_args)
     355            0 :             .arg("-f")
     356            0 :             .arg(&dumpdir)
     357            0 :             .arg("--no-sync")
     358            0 :             // POSITIONAL args
     359            0 :             // source db (db name included in connection string)
     360            0 :             .arg(&source_connstring)
     361            0 :             // how we run it
     362            0 :             .env_clear()
     363            0 :             .env("LD_LIBRARY_PATH", &pg_lib_dir)
     364            0 :             .env(
     365            0 :                 "ASAN_OPTIONS",
     366            0 :                 std::env::var("ASAN_OPTIONS").unwrap_or_default(),
     367            0 :             )
     368            0 :             .env(
     369            0 :                 "UBSAN_OPTIONS",
     370            0 :                 std::env::var("UBSAN_OPTIONS").unwrap_or_default(),
     371            0 :             )
     372            0 :             .kill_on_drop(true)
     373            0 :             .stdout(std::process::Stdio::piped())
     374            0 :             .stderr(std::process::Stdio::piped())
     375            0 :             .spawn()
     376            0 :             .context("spawn pg_dump")?;
     377              : 
     378            0 :         info!(pid=%pg_dump.id().unwrap(), "spawned pg_dump");
     379              : 
     380              :         tokio::spawn(
     381            0 :             child_stdio_to_log::relay_process_output(pg_dump.stdout.take(), pg_dump.stderr.take())
     382            0 :                 .instrument(info_span!("pg_dump")),
     383              :         );
     384              : 
     385            0 :         let st = pg_dump.wait().await.context("wait for pg_dump")?;
     386            0 :         info!(status=?st, "pg_dump exited");
     387            0 :         if !st.success() {
     388            0 :             error!(status=%st, "pg_dump failed, restore will likely fail as well");
     389            0 :             bail!("pg_dump failed");
     390            0 :         }
     391              :     }
     392              : 
     393              :     // TODO: maybe do it in a streaming way, plenty of internal research done on this already
     394              :     // TODO: do the unlogged table trick
     395              :     {
     396            0 :         let mut pg_restore = tokio::process::Command::new(pg_bin_dir.join("pg_restore"))
     397            0 :             .args(&common_args)
     398            0 :             .arg("-d")
     399            0 :             .arg(&destination_connstring)
     400            0 :             // POSITIONAL args
     401            0 :             .arg(&dumpdir)
     402            0 :             // how we run it
     403            0 :             .env_clear()
     404            0 :             .env("LD_LIBRARY_PATH", &pg_lib_dir)
     405            0 :             .env(
     406            0 :                 "ASAN_OPTIONS",
     407            0 :                 std::env::var("ASAN_OPTIONS").unwrap_or_default(),
     408            0 :             )
     409            0 :             .env(
     410            0 :                 "UBSAN_OPTIONS",
     411            0 :                 std::env::var("UBSAN_OPTIONS").unwrap_or_default(),
     412            0 :             )
     413            0 :             .kill_on_drop(true)
     414            0 :             .stdout(std::process::Stdio::piped())
     415            0 :             .stderr(std::process::Stdio::piped())
     416            0 :             .spawn()
     417            0 :             .context("spawn pg_restore")?;
     418              : 
     419            0 :         info!(pid=%pg_restore.id().unwrap(), "spawned pg_restore");
     420              :         tokio::spawn(
     421            0 :             child_stdio_to_log::relay_process_output(
     422            0 :                 pg_restore.stdout.take(),
     423            0 :                 pg_restore.stderr.take(),
     424            0 :             )
     425            0 :             .instrument(info_span!("pg_restore")),
     426              :         );
     427            0 :         let st = pg_restore.wait().await.context("wait for pg_restore")?;
     428            0 :         info!(status=?st, "pg_restore exited");
     429            0 :         if !st.success() {
     430            0 :             error!(status=%st, "pg_restore failed, restore will likely fail as well");
     431            0 :             bail!("pg_restore failed");
     432            0 :         }
     433            0 :     }
     434            0 : 
     435            0 :     Ok(())
     436            0 : }
     437              : 
     438              : #[allow(clippy::too_many_arguments)]
     439            0 : async fn cmd_pgdata(
     440            0 :     s3_client: Option<aws_sdk_s3::Client>,
     441            0 :     kms_client: Option<aws_sdk_kms::Client>,
     442            0 :     maybe_s3_prefix: Option<s3_uri::S3Uri>,
     443            0 :     maybe_spec: Option<Spec>,
     444            0 :     source_connection_string: Option<String>,
     445            0 :     interactive: bool,
     446            0 :     pg_port: u16,
     447            0 :     workdir: Utf8PathBuf,
     448            0 :     pg_bin_dir: Utf8PathBuf,
     449            0 :     pg_lib_dir: Utf8PathBuf,
     450            0 :     num_cpus: Option<usize>,
     451            0 :     memory_mb: Option<usize>,
     452            0 : ) -> Result<(), anyhow::Error> {
     453            0 :     if maybe_spec.is_none() && source_connection_string.is_none() {
     454            0 :         bail!("spec must be provided for pgdata command");
     455            0 :     }
     456            0 :     if maybe_spec.is_some() && source_connection_string.is_some() {
     457            0 :         bail!("only one of spec or source_connection_string can be provided");
     458            0 :     }
     459              : 
     460            0 :     let source_connection_string = if let Some(spec) = maybe_spec {
     461            0 :         match spec.encryption_secret {
     462            0 :             EncryptionSecret::KMS { key_id } => {
     463            0 :                 decode_connstring(
     464            0 :                     kms_client.as_ref().unwrap(),
     465            0 :                     &key_id,
     466            0 :                     spec.source_connstring_ciphertext_base64,
     467            0 :                 )
     468            0 :                 .await?
     469              :             }
     470              :         }
     471              :     } else {
     472            0 :         source_connection_string.unwrap()
     473              :     };
     474              : 
     475            0 :     let superuser = "cloud_admin";
     476            0 :     let destination_connstring = format!(
     477            0 :         "host=localhost port={} user={} dbname=neondb",
     478            0 :         pg_port, superuser
     479            0 :     );
     480            0 : 
     481            0 :     let pgdata_dir = workdir.join("pgdata");
     482            0 :     let mut proc = PostgresProcess::new(pgdata_dir.clone(), pg_bin_dir.clone(), pg_lib_dir.clone());
     483            0 :     let nproc = num_cpus.unwrap_or_else(num_cpus::get);
     484            0 :     let memory_mb = memory_mb.unwrap_or(256);
     485            0 :     proc.start(superuser, pg_port, nproc, memory_mb).await?;
     486            0 :     wait_until_ready(destination_connstring.clone(), "neondb".to_string()).await;
     487              : 
     488            0 :     run_dump_restore(
     489            0 :         workdir.clone(),
     490            0 :         pg_bin_dir,
     491            0 :         pg_lib_dir,
     492            0 :         source_connection_string,
     493            0 :         destination_connstring,
     494            0 :     )
     495            0 :     .await?;
     496              : 
     497              :     // If interactive mode, wait for Ctrl+C
     498            0 :     if interactive {
     499            0 :         info!("Running in interactive mode. Press Ctrl+C to shut down.");
     500            0 :         tokio::signal::ctrl_c().await.context("wait for ctrl-c")?;
     501            0 :     }
     502              : 
     503            0 :     proc.shutdown().await?;
     504              : 
     505              :     // Only sync if s3_prefix was specified
     506            0 :     if let Some(s3_prefix) = maybe_s3_prefix {
     507            0 :         info!("upload pgdata");
     508            0 :         aws_s3_sync::upload_dir_recursive(
     509            0 :             s3_client.as_ref().unwrap(),
     510            0 :             Utf8Path::new(&pgdata_dir),
     511            0 :             &s3_prefix.append("/pgdata/"),
     512            0 :         )
     513            0 :         .await
     514            0 :         .context("sync dump directory to destination")?;
     515              : 
     516            0 :         info!("write status");
     517              :         {
     518            0 :             let status_dir = workdir.join("status");
     519            0 :             std::fs::create_dir(&status_dir).context("create status directory")?;
     520            0 :             let status_file = status_dir.join("pgdata");
     521            0 :             std::fs::write(&status_file, serde_json::json!({"done": true}).to_string())
     522            0 :                 .context("write status file")?;
     523            0 :             aws_s3_sync::upload_dir_recursive(
     524            0 :                 s3_client.as_ref().unwrap(),
     525            0 :                 &status_dir,
     526            0 :                 &s3_prefix.append("/status/"),
     527            0 :             )
     528            0 :             .await
     529            0 :             .context("sync status directory to destination")?;
     530              :         }
     531            0 :     }
     532              : 
     533            0 :     Ok(())
     534            0 : }
     535              : 
     536            0 : async fn cmd_dumprestore(
     537            0 :     kms_client: Option<aws_sdk_kms::Client>,
     538            0 :     maybe_spec: Option<Spec>,
     539            0 :     source_connection_string: Option<String>,
     540            0 :     destination_connection_string: Option<String>,
     541            0 :     workdir: Utf8PathBuf,
     542            0 :     pg_bin_dir: Utf8PathBuf,
     543            0 :     pg_lib_dir: Utf8PathBuf,
     544            0 : ) -> Result<(), anyhow::Error> {
     545            0 :     let (source_connstring, destination_connstring) = if let Some(spec) = maybe_spec {
     546            0 :         match spec.encryption_secret {
     547            0 :             EncryptionSecret::KMS { key_id } => {
     548            0 :                 let source = decode_connstring(
     549            0 :                     kms_client.as_ref().unwrap(),
     550            0 :                     &key_id,
     551            0 :                     spec.source_connstring_ciphertext_base64,
     552            0 :                 )
     553            0 :                 .await?;
     554              : 
     555            0 :                 let dest = if let Some(dest_ciphertext) =
     556            0 :                     spec.destination_connstring_ciphertext_base64
     557              :                 {
     558            0 :                     decode_connstring(kms_client.as_ref().unwrap(), &key_id, dest_ciphertext)
     559            0 :                         .await?
     560              :                 } else {
     561            0 :                     bail!(
     562            0 :                         "destination connection string must be provided in spec for dump_restore command"
     563            0 :                     );
     564              :                 };
     565              : 
     566            0 :                 (source, dest)
     567              :             }
     568              :         }
     569              :     } else {
     570              :         (
     571            0 :             source_connection_string.unwrap(),
     572            0 :             if let Some(val) = destination_connection_string {
     573            0 :                 val
     574              :             } else {
     575            0 :                 bail!("destination connection string must be provided for dump_restore command");
     576              :             },
     577              :         )
     578              :     };
     579              : 
     580            0 :     run_dump_restore(
     581            0 :         workdir,
     582            0 :         pg_bin_dir,
     583            0 :         pg_lib_dir,
     584            0 :         source_connstring,
     585            0 :         destination_connstring,
     586            0 :     )
     587            0 :     .await
     588            0 : }
     589              : 
     590              : #[tokio::main]
     591            0 : pub(crate) async fn main() -> anyhow::Result<()> {
     592            0 :     utils::logging::init(
     593            0 :         utils::logging::LogFormat::Json,
     594            0 :         utils::logging::TracingErrorLayerEnablement::EnableWithRustLogFilter,
     595            0 :         utils::logging::Output::Stdout,
     596            0 :     )?;
     597            0 : 
     598            0 :     info!("starting");
     599            0 : 
     600            0 :     let args = Args::parse();
     601            0 : 
     602            0 :     // Initialize AWS clients only if s3_prefix is specified
     603            0 :     let (s3_client, kms_client) = if args.s3_prefix.is_some() {
     604            0 :         let config = aws_config::load_defaults(BehaviorVersion::v2024_03_28()).await;
     605            0 :         let s3_client = aws_sdk_s3::Client::new(&config);
     606            0 :         let kms = aws_sdk_kms::Client::new(&config);
     607            0 :         (Some(s3_client), Some(kms))
     608            0 :     } else {
     609            0 :         (None, None)
     610            0 :     };
     611            0 : 
     612            0 :     let spec: Option<Spec> = if let Some(s3_prefix) = &args.s3_prefix {
     613            0 :         let spec_key = s3_prefix.append("/spec.json");
     614            0 :         let object = s3_client
     615            0 :             .as_ref()
     616            0 :             .unwrap()
     617            0 :             .get_object()
     618            0 :             .bucket(&spec_key.bucket)
     619            0 :             .key(spec_key.key)
     620            0 :             .send()
     621            0 :             .await
     622            0 :             .context("get spec from s3")?
     623            0 :             .body
     624            0 :             .collect()
     625            0 :             .await
     626            0 :             .context("download spec body")?;
     627            0 :         serde_json::from_slice(&object.into_bytes()).context("parse spec as json")?
     628            0 :     } else {
     629            0 :         None
     630            0 :     };
     631            0 : 
     632            0 :     match tokio::fs::create_dir(&args.working_directory).await {
     633            0 :         Ok(()) => {}
     634            0 :         Err(e) if e.kind() == std::io::ErrorKind::AlreadyExists => {
     635            0 :             if !is_directory_empty(&args.working_directory)
     636            0 :                 .await
     637            0 :                 .context("check if working directory is empty")?
     638            0 :             {
     639            0 :                 bail!("working directory is not empty");
     640            0 :             } else {
     641            0 :                 // ok
     642            0 :             }
     643            0 :         }
     644            0 :         Err(e) => return Err(anyhow::Error::new(e).context("create working directory")),
     645            0 :     }
     646            0 : 
     647            0 :     match args.command {
     648            0 :         Command::Pgdata {
     649            0 :             source_connection_string,
     650            0 :             interactive,
     651            0 :             pg_port,
     652            0 :             num_cpus,
     653            0 :             memory_mb,
     654            0 :         } => {
     655            0 :             cmd_pgdata(
     656            0 :                 s3_client,
     657            0 :                 kms_client,
     658            0 :                 args.s3_prefix,
     659            0 :                 spec,
     660            0 :                 source_connection_string,
     661            0 :                 interactive,
     662            0 :                 pg_port,
     663            0 :                 args.working_directory,
     664            0 :                 args.pg_bin_dir,
     665            0 :                 args.pg_lib_dir,
     666            0 :                 num_cpus,
     667            0 :                 memory_mb,
     668            0 :             )
     669            0 :             .await?;
     670            0 :         }
     671            0 :         Command::DumpRestore {
     672            0 :             source_connection_string,
     673            0 :             destination_connection_string,
     674            0 :         } => {
     675            0 :             cmd_dumprestore(
     676            0 :                 kms_client,
     677            0 :                 spec,
     678            0 :                 source_connection_string,
     679            0 :                 destination_connection_string,
     680            0 :                 args.working_directory,
     681            0 :                 args.pg_bin_dir,
     682            0 :                 args.pg_lib_dir,
     683            0 :             )
     684            0 :             .await?;
     685            0 :         }
     686            0 :     }
     687            0 : 
     688            0 :     Ok(())
     689            0 : }
        

Generated by: LCOV version 2.1-beta