From 241724f3fcd7b22d0159d2e6fcf76a2f28391396 Mon Sep 17 00:00:00 2001
From: Heikki Linnakangas <heikki@neon.tech>
Date: Thu, 12 Sep 2024 12:31:07 +0300
Subject: [PATCH] CLI args parsing

---
 pageserver/src/bin/import.rs | 32 ++++++++++++++++++++++++++++----
 pageserver/src/pg_import.rs  | 22 ++++++----------------
 2 files changed, 34 insertions(+), 20 deletions(-)

diff --git a/pageserver/src/bin/import.rs b/pageserver/src/bin/import.rs
index e710cd09d5..a4d96bcc9c 100644
--- a/pageserver/src/bin/import.rs
+++ b/pageserver/src/bin/import.rs
@@ -1,10 +1,26 @@
 use anyhow;
+use camino::Utf8PathBuf;
+use clap::Parser;
 use pageserver::{pg_import, virtual_file::{self, api::IoEngineKind}};
 use utils::logging::{self, LogFormat, TracingErrorLayerEnablement};
 
-fn main() -> anyhow::Result<()> {
-    println!("Hey!");
+//project_git_version!(GIT_VERSION);
 
+#[derive(Parser)]
+#[command(
+    //version = GIT_VERSION,
+    about = "Utility to import a Postgres data directory directly into image layers",
+    //long_about = "..."
+)]
+struct CliOpts {
+    /// Input Postgres data directory
+    pgdata: Utf8PathBuf,
+
+    /// Path to local dir where the layer files will be stored
+    dest_path: Utf8PathBuf,
+}
+
+fn main() -> anyhow::Result<()> {
     logging::init(
         LogFormat::Plain,
         TracingErrorLayerEnablement::EnableWithRustLogFilter,
@@ -21,7 +37,15 @@ fn main() -> anyhow::Result<()> {
         .enable_all()
         .build()?;
 
-    rt.block_on(pg_import::do_import())?;
-    
+    let cli = CliOpts::parse();
+
+    rt.block_on(async_main(cli))?;
+
+    Ok(())
+}
+
+async fn async_main(cli: CliOpts) -> anyhow::Result<()> {
+    let mut import = pg_import::PgImportEnv::init().await?;
+    import.import_datadir(&cli.pgdata, &cli.dest_path).await?;
     Ok(())
 }
diff --git a/pageserver/src/pg_import.rs b/pageserver/src/pg_import.rs
index 140a7dc423..4b205e4250 100644
--- a/pageserver/src/pg_import.rs
+++ b/pageserver/src/pg_import.rs
@@ -18,16 +18,6 @@ use tokio::io::AsyncReadExt;
 
 use pageserver_api::key::Key;
 
-
-pub async fn do_import() -> anyhow::Result<()> {
-
-    let mut import = PgImportEnv::init().await?;
-
-    import.import_datadir(&Utf8PathBuf::from("pg_in"), &Utf8PathBuf::from("pg_out")).await?;
-
-    Ok(())
-}
-
 pub struct PgImportEnv {
     ctx: RequestContext,
     conf: &'static PageServerConf,
@@ -37,7 +27,7 @@ pub struct PgImportEnv {
 
 impl PgImportEnv {
 
-    async fn init() -> anyhow::Result<PgImportEnv> {
+    pub async fn init() -> anyhow::Result<PgImportEnv> {
         let ctx: RequestContext = RequestContext::new(TaskKind::DebugTool, DownloadBehavior::Error);
         let config = toml_edit::Document::new();
         let conf = PageServerConf::parse_and_validate(
@@ -63,7 +53,7 @@ impl PgImportEnv {
         })
     }
 
-    async fn import_datadir(&mut self, pgdata_path: &Utf8Path, _tenant_path: &Utf8Path) -> anyhow::Result<()> {
+    pub async fn import_datadir(&mut self, pgdata_path: &Utf8Path, _tenant_path: &Utf8Path) -> anyhow::Result<()> {
 
         let pgdata_lsn = import_datadir::get_lsn_from_controlfile(&pgdata_path)?.align();
 
@@ -90,11 +80,11 @@ impl PgImportEnv {
             .sorted()
             .for_each(|dboid| {
                 let path = pgdata_path.join("base").join(dboid.to_string());
-                self.import_db(&mut one_big_layer, &path, pg_constants::DEFAULTTABLESPACE_OID);
+                self.import_db(&mut one_big_layer, &path, pg_constants::DEFAULTTABLESPACE_OID).await;
             });
 
         // global catalogs now
-        self.import_db(&mut one_big_layer, &pgdata_path.join("global"), postgres_ffi::pg_constants::GLOBALTABLESPACE_OID);
+        self.import_db(&mut one_big_layer, &pgdata_path.join("global"), postgres_ffi::pg_constants::GLOBALTABLESPACE_OID).await?;
 
         
         one_big_layer.finish_layer(&self.ctx).await?;
@@ -136,7 +126,7 @@ impl PgImportEnv {
     ) -> anyhow::Result<()> {
 
         let mut reader = tokio::fs::File::open(path).await?;
-        let len = metadata(path)?.len();
+        let len = std::fs::metadata(path)?.len();
 
         let mut buf: [u8; 8192] = [0u8; 8192];
 
@@ -168,7 +158,7 @@ impl PgImportEnv {
                         break;
                     }
                     _ => {
-                        bail!("error reading file {}: {:#}", path.display(), err);
+                        bail!("error reading file {}: {:#}", path.as_display(), err);
                     }
                 },
             };