Add omdb support for clickhouse policy

Uses new internal nexus API endpoints.

clients/nexus-client/src/lib.rs

@@ -30,6 +30,8 @@ progenitor::generate_api!(
         // (e.g., diff'ing) that's implemented on our local type.
         Blueprint = nexus_types::deployment::Blueprint,
         Certificate = omicron_common::api::internal::nexus::Certificate,
+        ClickhouseMode = nexus_types::deployment::ClickhouseMode,
+        ClickhousePolicy = nexus_types::deployment::ClickhousePolicy,
         DatasetKind = omicron_common::api::internal::shared::DatasetKind,
         DnsConfigParams = nexus_types::internal_api::params::DnsConfigParams,
         DnsConfigZone = nexus_types::internal_api::params::DnsConfigZone,

dev-tools/omdb/Cargo.toml

@@ -26,6 +26,7 @@ futures.workspace = true
 gateway-client.workspace = true
 gateway-messages.workspace = true
 gateway-test-utils.workspace = true
+http.workspace = true
 humantime.workspace = true
 internal-dns-resolver.workspace = true
 internal-dns-types.workspace = true
@@ -34,6 +35,7 @@ nexus-client.workspace = true
 nexus-config.workspace = true
 nexus-db-model.workspace = true
 nexus-db-queries.workspace = true
+nexus-inventory.workspace = true
 nexus-reconfigurator-preparation.workspace = true
 nexus-saga-recovery.workspace = true
 nexus-types.workspace = true

dev-tools/omdb/src/bin/omdb/nexus.rs

@@ -22,6 +22,7 @@ use clap::Subcommand;
 use clap::ValueEnum;
 use futures::future::try_join;
 use futures::TryStreamExt;
+use http::StatusCode;
 use internal_dns_types::names::ServiceName;
 use itertools::Itertools;
 use nexus_client::types::ActivationReason;
@@ -34,8 +35,11 @@ use nexus_client::types::SagaState;
 use nexus_client::types::SledSelector;
 use nexus_client::types::UninitializedSledId;
 use nexus_db_queries::db::lookup::LookupPath;
+use nexus_inventory::now_db_precision;
 use nexus_saga_recovery::LastPass;
 use nexus_types::deployment::Blueprint;
+use nexus_types::deployment::ClickhouseMode;
+use nexus_types::deployment::ClickhousePolicy;
 use nexus_types::internal_api::background::AbandonedVmmReaperStatus;
 use nexus_types::internal_api::background::InstanceReincarnationStatus;
 use nexus_types::internal_api::background::InstanceUpdaterStatus;
@@ -101,6 +105,8 @@ enum NexusCommands {
     BackgroundTasks(BackgroundTasksArgs),
     /// interact with blueprints
     Blueprints(BlueprintsArgs),
+    /// Interact with clickhouse policy
+    ClickhousePolicy(ClickhousePolicyArgs),
     /// view sagas, create and complete demo sagas
     Sagas(SagasArgs),
     /// interact with sleds
@@ -299,6 +305,42 @@ struct BlueprintImportArgs {
     input: Utf8PathBuf,
 }
 
+#[derive(Debug, Args)]
+struct ClickhousePolicyArgs {
+    #[command(subcommand)]
+    command: ClickhousePolicyCommands,
+}
+
+#[derive(Debug, Subcommand)]
+enum ClickhousePolicyCommands {
+    /// Get the current policy
+    Get,
+    /// Set the new policy
+    Set(ClickhousePolicySetArgs),
+}
+
+#[derive(Debug, Args)]
+struct ClickhousePolicySetArgs {
+    mode: ClickhousePolicyMode,
+
+    /// The number of servers in a clickhouse cluster
+    #[arg(long, default_value_t = 3)]
+    target_servers: u8,
+    /// The number of keepers in a clickhouse cluster
+    #[arg(long, default_value_t = 5)]
+    target_keepers: u8,
+}
+
+#[derive(Debug, Clone, Copy, ValueEnum)]
+enum ClickhousePolicyMode {
+    /// Run only a single node clickhouse instance
+    SingleNodeOnly,
+    // Run only a clickhouse cluster
+    ClusterOnly,
+    // Run both single-node and clustered clickhouse deployments
+    Both,
+}
+
 #[derive(Debug, Args)]
 struct SagasArgs {
     #[command(subcommand)]
@@ -494,6 +536,18 @@ impl NexusArgs {
                 cmd_nexus_blueprints_import(&client, token, args).await
             }
 
+            NexusCommands::ClickhousePolicy(ClickhousePolicyArgs {
+                command,
+            }) => match command {
+                ClickhousePolicyCommands::Get => {
+                    cmd_nexus_clickhouse_policy_get(&client).await
+                }
+                ClickhousePolicyCommands::Set(args) => {
+                    let token = omdb.check_allow_destructive()?;
+                    cmd_nexus_clickhouse_policy_set(&client, args, token).await
+                }
+            },
+
             NexusCommands::Sagas(SagasArgs { command }) => {
                 if self.nexus_internal_url.is_none() {
                     eprintln!(
@@ -2376,6 +2430,100 @@ async fn cmd_nexus_blueprints_import(
     Ok(())
 }
 
+async fn cmd_nexus_clickhouse_policy_get(
+    client: &nexus_client::Client,
+) -> Result<(), anyhow::Error> {
+    let res = client.clickhouse_policy_get().await;
+
+    match res {
+        Err(err) => {
+            if err.status() == Some(StatusCode::NOT_FOUND) {
+                println!(
+                    "No clickhouse policy: \
+                    Defaulting to single-node deployment"
+                );
+            } else {
+                eprintln!("error: {:#}", err);
+            }
+        }
+        Ok(policy) => {
+            println!("Clickhouse Policy: ");
+            println!("    version: {}", policy.version);
+            println!("    creation time: {}", policy.time_created);
+            match policy.mode {
+                ClickhouseMode::SingleNodeOnly => {
+                    println!("    mode: single-node-only");
+                }
+                ClickhouseMode::ClusterOnly {
+                    target_servers,
+                    target_keepers,
+                } => {
+                    println!("    mode: cluster-only");
+                    println!("    target-servers: {}", target_servers);
+                    println!("    target-keepers: {}", target_keepers);
+                }
+                ClickhouseMode::Both { target_servers, target_keepers } => {
+                    println!("    mode: both single-node and cluster");
+                    println!("    target-servers: {}", target_servers);
+                    println!("    target-keepers: {}", target_keepers);
+                }
+            }
+        }
+    }
+
+    Ok(())
+}
+
+async fn cmd_nexus_clickhouse_policy_set(
+    client: &nexus_client::Client,
+    args: &ClickhousePolicySetArgs,
+    _destruction_token: DestructiveOperationToken,
+) -> Result<(), anyhow::Error> {
+    let mode = match args.mode {
+        ClickhousePolicyMode::SingleNodeOnly => ClickhouseMode::SingleNodeOnly,
+        ClickhousePolicyMode::ClusterOnly => ClickhouseMode::ClusterOnly {
+            target_servers: args.target_servers,
+            target_keepers: args.target_keepers,
+        },
+        ClickhousePolicyMode::Both => ClickhouseMode::Both {
+            target_servers: args.target_servers,
+            target_keepers: args.target_keepers,
+        },
+    };
+
+    let res = client.clickhouse_policy_get().await;
+    let new_policy = match res {
+        Err(err) => {
+            if err.status() == Some(StatusCode::NOT_FOUND) {
+                ClickhousePolicy {
+                    version: 1,
+                    mode,
+                    time_created: now_db_precision(),
+                }
+            } else {
+                eprintln!("error: {:#}", err);
+                return Err(err).context("retrieving clickhouse policy");
+            }
+        }
+        Ok(policy) => ClickhousePolicy {
+            version: policy.version + 1,
+            mode,
+            time_created: now_db_precision(),
+        },
+    };
+
+    client.clickhouse_policy_set(&new_policy).await.with_context(|| {
+        format!("inserting new context at version {}", new_policy.version)
+    })?;
+
+    println!(
+        "Successfully inserted new policy at version {}",
+        new_policy.version
+    );
+
+    Ok(())
+}
+
 /// Runs `omdb nexus sagas list`
 async fn cmd_nexus_sagas_list(
     client: &nexus_client::Client,

nexus/internal-api/src/lib.rs

@@ -12,6 +12,7 @@ use dropshot::{
 use nexus_types::{
     deployment::{
         Blueprint, BlueprintMetadata, BlueprintTarget, BlueprintTargetSet,
+        ClickhousePolicy,
     },
     external_api::{
         params::{PhysicalDiskPath, SledSelector, UninitializedSledId},
@@ -530,6 +531,25 @@ pub trait NexusInternalApi {
         path_params: Path<ProbePathParam>,
         query_params: Query<PaginatedById>,
     ) -> Result<HttpResponseOk<Vec<ProbeInfo>>, HttpError>;
+
+    /// Get the current clickhouse policy
+    #[endpoint {
+     method = GET,
+     path = "/clickhouse/policy"
+     }]
+    async fn clickhouse_policy_get(
+        rqctx: RequestContext<Self::Context>,
+    ) -> Result<HttpResponseOk<ClickhousePolicy>, HttpError>;
+
+    /// Set the new clickhouse policy
+    #[endpoint {
+        method = POST,
+        path = "/clickhouse/policy"
+    }]
+    async fn clickhouse_policy_set(
+        rqctx: RequestContext<Self::Context>,
+        policy: TypedBody<ClickhousePolicy>,
+    ) -> Result<HttpResponseUpdatedNoContent, HttpError>;
 }
 
 /// Path parameters for Sled Agent requests (internal API)

nexus/src/internal_api/http_entrypoints.rs

@@ -24,6 +24,7 @@ use nexus_types::deployment::Blueprint;
 use nexus_types::deployment::BlueprintMetadata;
 use nexus_types::deployment::BlueprintTarget;
 use nexus_types::deployment::BlueprintTargetSet;
+use nexus_types::deployment::ClickhousePolicy;
 use nexus_types::external_api::params::PhysicalDiskPath;
 use nexus_types::external_api::params::SledSelector;
 use nexus_types::external_api::params::UninitializedSledId;
@@ -939,4 +940,51 @@ impl NexusInternalApi for NexusInternalApiImpl {
             .instrument_dropshot_handler(&rqctx, handler)
             .await
     }
+
+    async fn clickhouse_policy_get(
+        rqctx: RequestContext<Self::Context>,
+    ) -> Result<HttpResponseOk<ClickhousePolicy>, HttpError> {
+        let apictx = &rqctx.context().context;
+        let handler = async {
+            let nexus = &apictx.nexus;
+            let opctx =
+                crate::context::op_context_for_internal_api(&rqctx).await;
+            match nexus.datastore().clickhouse_policy_get_latest(&opctx).await?
+            {
+                Some(policy) => Ok(HttpResponseOk(policy)),
+                None => Err(HttpError::for_not_found(
+                    None,
+                    "No clickhouse policy in database".into(),
+                )),
+            }
+        };
+        apictx
+            .internal_latencies
+            .instrument_dropshot_handler(&rqctx, handler)
+            .await
+    }
+
+    async fn clickhouse_policy_set(
+        rqctx: RequestContext<Self::Context>,
+        policy: TypedBody<ClickhousePolicy>,
+    ) -> Result<HttpResponseUpdatedNoContent, HttpError> {
+        let apictx = &rqctx.context().context;
+        let nexus = &apictx.nexus;
+        let handler = async {
+            let opctx =
+                crate::context::op_context_for_internal_api(&rqctx).await;
+            nexus
+                .datastore()
+                .clickhouse_policy_insert_latest_version(
+                    &opctx,
+                    &policy.into_inner(),
+                )
+                .await?;
+            Ok(HttpResponseUpdatedNoContent())
+        };
+        apictx
+            .internal_latencies
+            .instrument_dropshot_handler(&rqctx, handler)
+            .await
+    }
 }

nexus/types/src/deployment/planning_input.rs

@@ -887,15 +887,16 @@ pub struct Policy {
     pub clickhouse_policy: Option<ClickhousePolicy>,
 }
 
-#[derive(Debug, Clone, Serialize, Deserialize)]
+#[derive(Debug, Clone, Serialize, Deserialize, JsonSchema)]
 pub struct ClickhousePolicy {
     pub version: u32,
     pub mode: ClickhouseMode,
     pub time_created: DateTime<Utc>,
 }
 
 /// How to deploy clickhouse nodes
-#[derive(Debug, Clone, Serialize, Deserialize)]
+#[derive(Debug, Clone, Serialize, Deserialize, JsonSchema)]
+#[serde(rename_all = "snake_case", tag = "type", content = "value")]
 pub enum ClickhouseMode {
     SingleNodeOnly,
     ClusterOnly { target_servers: u8, target_keepers: u8 },