From fe975acc715ccce0df7e67accbee1fe2a5635544 Mon Sep 17 00:00:00 2001
From: Heikki Linnakangas <heikki@neon.tech>
Date: Thu, 12 Sep 2024 13:28:12 +0300
Subject: [PATCH] Add --tenant-id and --timeline-id options

---
 pageserver/src/bin/import.rs          | 14 ++++++++++++--
 pageserver/src/pg_import.rs           | 18 +++++++++---------
 test_runner/regress/test_pg_import.py |  1 +
 3 files changed, 22 insertions(+), 11 deletions(-)

diff --git a/pageserver/src/bin/import.rs b/pageserver/src/bin/import.rs
index a4d96bcc9c..7873f4d5d6 100644
--- a/pageserver/src/bin/import.rs
+++ b/pageserver/src/bin/import.rs
@@ -2,8 +2,11 @@ use anyhow;
 use camino::Utf8PathBuf;
 use clap::Parser;
 use pageserver::{pg_import, virtual_file::{self, api::IoEngineKind}};
+use utils::id::{TenantId, TimelineId};
 use utils::logging::{self, LogFormat, TracingErrorLayerEnablement};
 
+use std::str::FromStr;
+
 //project_git_version!(GIT_VERSION);
 
 #[derive(Parser)]
@@ -18,6 +21,11 @@ struct CliOpts {
 
     /// Path to local dir where the layer files will be stored
     dest_path: Utf8PathBuf,
+
+    #[arg(long, default_value_t = TenantId::from_str("42424242424242424242424242424242").unwrap())]
+    tenant_id: TenantId,
+    #[arg(long, default_value_t = TimelineId::from_str("42424242424242424242424242424242").unwrap())]
+    timeline_id: TimelineId,
 }
 
 fn main() -> anyhow::Result<()> {
@@ -45,7 +53,9 @@ fn main() -> anyhow::Result<()> {
 }
 
 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?;
+    let mut import = pg_import::PgImportEnv::init(&cli.dest_path, cli.tenant_id, cli.timeline_id).await?;
+
+    import.import_datadir(&cli.pgdata).await?;
+
     Ok(())
 }
diff --git a/pageserver/src/pg_import.rs b/pageserver/src/pg_import.rs
index cd11bfd8f7..1cb1ba5654 100644
--- a/pageserver/src/pg_import.rs
+++ b/pageserver/src/pg_import.rs
@@ -1,4 +1,4 @@
-use std::{fs::metadata, path::Path, str::FromStr};
+use std::{fs::metadata, path::Path};
 
 use anyhow::{bail, ensure, Context};
 use bytes::Bytes;
@@ -27,20 +27,18 @@ pub struct PgImportEnv {
 
 impl PgImportEnv {
 
-    pub async fn init() -> anyhow::Result<PgImportEnv> {
+    pub async fn init(dstdir: &Utf8Path, tenant_id: TenantId, timeline_id: TimelineId) -> anyhow::Result<PgImportEnv> {
         let ctx: RequestContext = RequestContext::new(TaskKind::DebugTool, DownloadBehavior::Error);
         let config = toml_edit::Document::new();
         let conf = PageServerConf::parse_and_validate(
             NodeId(42), 
             &config,
-            &Utf8PathBuf::from("layers")
+            dstdir
         )?;
         let conf = Box::leak(Box::new(conf));
 
-        let tni = TenantId::from_str("42424242424242424242424242424242")?;
-        let tli = TimelineId::from_str("42424242424242424242424242424242")?;
         let tsi = TenantShardId {
-            tenant_id: tni,
+            tenant_id,
             shard_number: ShardNumber(0),
             shard_count: ShardCount(0),
         };
@@ -48,16 +46,18 @@ impl PgImportEnv {
         Ok(PgImportEnv {
             ctx,
             conf, 
-            tli,
+            tli: timeline_id,
             tsi,
         })
     }
 
-    pub async fn import_datadir(&mut self, pgdata_path: &Utf8Path, _tenant_path: &Utf8Path) -> anyhow::Result<()> {
+    pub async fn import_datadir(&mut self, pgdata_path: &Utf8Path) -> anyhow::Result<()> {
 
         let pgdata_lsn = import_datadir::get_lsn_from_controlfile(&pgdata_path)?.align();
 
-        println!("Importing {pgdata_path} to {_tenant_path} as lsn {pgdata_lsn}...");
+        let timeline_path = self.conf.timeline_path(&self.tsi, &self.tli);
+
+        println!("Importing {pgdata_path} to {timeline_path} as lsn {pgdata_lsn}...");
 
         let range = Key::MIN..Key::NON_L0_MAX;
         let mut one_big_layer = ImageLayerWriter::new(
diff --git a/test_runner/regress/test_pg_import.py b/test_runner/regress/test_pg_import.py
index 85c826b5ae..1a55513442 100644
--- a/test_runner/regress/test_pg_import.py
+++ b/test_runner/regress/test_pg_import.py
@@ -51,6 +51,7 @@ def test_pg_import(test_output_dir, pg_bin, vanilla_pg, neon_env_builder):
     timeline_id = TimelineId.generate()
 
     tline_path = env.pageserver_remote_storage.timeline_path(tenant_id, timeline_id)
+    tline_path.mkdir(parents=True)
 
     cli = ImportCli(env)
     cli.run_import(vanilla_pg.pgdatadir, tline_path)