From fb6a942665d99e2231230bdda49cccceac490dd4 Mon Sep 17 00:00:00 2001
From: Alek Westover <alek.westover@gmail.com>
Date: Tue, 6 Jun 2023 15:39:05 -0400
Subject: [PATCH] drafted some modifications to compute_ctl to add support for
 downloading pg extensions. not tested yet.

---
 compute_tools/src/bin/compute_ctl.rs |  9 +++
 compute_tools/src/pg_extensions.rs   | 86 ++++++++++++++++++++++++++++
 2 files changed, 95 insertions(+)
 create mode 100644 compute_tools/src/pg_extensions.rs

diff --git a/compute_tools/src/bin/compute_ctl.rs b/compute_tools/src/bin/compute_ctl.rs
index c6cfde1d1a..cd5afa89f4 100644
--- a/compute_tools/src/bin/compute_ctl.rs
+++ b/compute_tools/src/bin/compute_ctl.rs
@@ -53,11 +53,14 @@ use compute_tools::logger::*;
 use compute_tools::monitor::launch_monitor;
 use compute_tools::params::*;
 use compute_tools::spec::*;
+use compute_tools::pg_extensions::*;
 
 fn main() -> Result<()> {
     init_tracing_and_logging(DEFAULT_LOG_LEVEL)?;
 
     let matches = cli().get_matches();
+    let config = get_S3_config(matches);
+    download_extension(matches, ExtensionType::Shared);
 
     let http_port = *matches
         .get_one::<u16>("http-port")
@@ -202,6 +205,9 @@ fn main() -> Result<()> {
     // We got all we need, update the state.
     let mut state = compute.state.lock().unwrap();
 
+    // Now we have the spec, and also the tenant id, so we can download the user's personal extensions
+    // download_extension(matches, ExtensionType::Tenant(FIXME tenant_id.into()));
+
     // Record for how long we slept waiting for the spec.
     state.metrics.wait_for_spec_ms = Utc::now()
         .signed_duration_since(state.start_time)
@@ -221,6 +227,9 @@ fn main() -> Result<()> {
     let _configurator_handle =
         launch_configurator(&compute).expect("cannot launch configurator thread");
 
+    // Now we are ready to download library extensions
+    // download_extension(matches, ExtensionType::Library(FIXME library_name.into()));
+    
     // Start Postgres
     let mut delay_exit = false;
     let mut exit_code = None;
diff --git a/compute_tools/src/pg_extensions.rs b/compute_tools/src/pg_extensions.rs
new file mode 100644
index 0000000000..d6a4cc2e3c
--- /dev/null
+++ b/compute_tools/src/pg_extensions.rs
@@ -0,0 +1,86 @@
+// This is some code for downloading postgres extensions from AWS S3
+use std::{env, ops::ControlFlow, path::Path, str::FromStr};
+use clap::{Arg, ArgAction, Command};
+
+fn get_pg_config(argument: &str) -> String {
+    // NOTE: this function panics if it runs into any issues;
+    // If this is not desired, should FIXME this
+    let config_output = std::process::Command::new("pg_config")
+        .arg(argument)
+        .output()
+        .expect("pg_config should be installed");
+    assert!(config_output.status.success());
+
+    let stdout = std::str::from_utf8(&config_output.stdout).unwrap();
+    stdout.trim().to_string()
+}
+
+fn download_helper(config: RemoteStorageConfig, from_path: &str, to_path: &str) -> anyhow::Result<()> {
+    let mut remote_storage = GenericRemoteStorage::from_config(config)?;
+    let data = remote_storage.download(remote_path);
+    // TODO: write "data" to "to_path"
+    println("received data");
+    println!("{:?}", data);
+    Ok(());
+}
+
+pub enum ExtensionType {
+    Shared,
+    Tenant(String),
+    Library(String)
+}
+
+// TODO: should I make this async?
+pub fn download_extension(config: RemoteStorageConfig, ext_type: ExtensionType) -> anyhow::Result<()>{
+    let sharedir = get_pg_config("--sharedir");
+    let sharedir = format!("{}/extension", sharedir);
+    let libdir = get_pg_config("--libdir");
+
+    match ext_type {
+        Shared => {
+            // 1. Download control files from s3-bucket/public/*.control to SHAREDIR/extension
+            // We can do this step even before we have spec,
+            // because public extensions are common for all projects.
+            let from_path = "s3-bucket/public/*.control"
+            download_helper(config, from_path, sharedir);
+        }
+        Tenant(tenant_id) => {
+            // 2. After we have spec, before project start
+            // Download control files from s3-bucket/[tenant-id]/*.control to SHAREDIR/extension
+            let from_path = format!("s3-bucket/{tenant_id}/*.control");
+            download_helper(config, from_path, sharedir);
+        }
+        Library(library_name) => {
+            // 3. After we have spec, before postgres start
+            // Download preload_shared_libraries from s3-bucket/public/[library-name].control into LIBDIR/
+            let from_path = format!("s3-bucket/public/{library_name}.control");
+            download_helper(config, from_path, libdir);
+        }
+    }
+}
+
+pub get_S3_config(arg_matches: ArgMatches) -> anyhow::Result<RemoteStorageConfig> {
+    let workdir = arg_matches
+        .get_one::<String>("workdir")
+        .map(Path::new)
+        .unwrap_or_else(|| Path::new(".neon"))
+        .canonicalize()
+        .with_context(|| format!("Error opening workdir '{}'", workdir.display()))?;
+
+    // TODO: is this the correct file location?
+    let cfg_file_path = workdir.join("pageserver.toml");
+    let conf = match initialize_config(&cfg_file_path, arg_matches, &workdir)? {
+        ControlFlow::Continue(conf) => conf,
+        ControlFlow::Break(()) => {
+            info!("Pageserver config init successful");
+            return Ok(());
+        }
+    };
+    if let Some(config) = &conf.remote_storage_config {
+        return config;
+    } else { 
+        // No remote storage configured.
+        return Ok(None);
+    };
+}
+