From d2dd38c2cf2674224e29631720e4b4c5b1b8d510 Mon Sep 17 00:00:00 2001
From: "Carol (Nichols || Goulding)" <carol.nichols@gmail.com>
Date: Wed, 21 Sep 2022 14:22:52 -0400
Subject: [PATCH] feat: Set up CLI for compactor generate

---
 influxdb_iox/src/commands/compactor.rs        | 72 ++++++++++++++++++-
 .../tests/end_to_end_cases/compactor.rs       | 42 +++++++++++
 influxdb_iox/tests/end_to_end_cases/mod.rs    |  1 +
 3 files changed, 114 insertions(+), 1 deletion(-)
 create mode 100644 influxdb_iox/tests/end_to_end_cases/compactor.rs

diff --git a/influxdb_iox/src/commands/compactor.rs b/influxdb_iox/src/commands/compactor.rs
index afcd4b4f23..ac26070782 100644
--- a/influxdb_iox/src/commands/compactor.rs
+++ b/influxdb_iox/src/commands/compactor.rs
@@ -1,3 +1,4 @@
+use clap::ValueEnum;
 use clap_blocks::{
     catalog_dsn::CatalogDsnConfig,
     compactor::CompactorOnceConfig,
@@ -9,7 +10,7 @@ use ioxd_compactor::build_compactor_from_config;
 use object_store::DynObjectStore;
 use object_store_metrics::ObjectStoreMetrics;
 use snafu::prelude::*;
-use std::sync::Arc;
+use std::{num::NonZeroUsize, sync::Arc};
 
 #[derive(Debug, clap::Parser)]
 pub struct Config {
@@ -39,6 +40,74 @@ pub enum Command {
         )]
         query_exec_thread_count: usize,
     },
+
+    /// Generate Parquet files and catalog entries with different characteristics for the purposes
+    /// of investigating how the compactor handles them.
+    Generate {
+        #[clap(flatten)]
+        object_store_config: ObjectStoreConfig,
+
+        #[clap(flatten)]
+        catalog_dsn: CatalogDsnConfig,
+
+        /// The type of compaction to be done on the files. If `hot` is specified, the generated
+        /// files will have compaction level 0. If `cold` is specified, the generated files will
+        /// have compaction level 1 and will be marked that they were created at least 8 hours ago.
+        #[clap(
+            arg_enum,
+            value_parser,
+            long = "--compaction-type",
+            env = "INFLUXDB_IOX_COMPACTOR_GENERATE_TYPE",
+            default_value = "hot",
+            action
+        )]
+        compaction_type: CompactionType,
+
+        /// The number of IOx partitions to generate files for. Each partition will have the number
+        /// of files specified by `--num-files` generated.
+        #[clap(
+            long = "--num-partitions",
+            env = "INFLUXDB_IOX_COMPACTOR_GENERATE_NUM_PARTITIONS",
+            default_value = "1",
+            action
+        )]
+        num_partitions: NonZeroUsize,
+
+        /// The number of parquet files to generate per partition.
+        #[clap(
+            long = "--num-files",
+            env = "INFLUXDB_IOX_COMPACTOR_GENERATE_NUM_FILES",
+            default_value = "1",
+            action
+        )]
+        num_files: NonZeroUsize,
+
+        /// The number of columns to generate in each file. One column will always be the
+        /// timestamp. Additional columns will be given a type in I64, F64, U64, String, Bool, and
+        /// Tag in equal proportion.
+        #[clap(
+            long = "--num-cols",
+            env = "INFLUXDB_IOX_COMPACTOR_GENERATE_NUM_COLS",
+            default_value = "7",
+            action
+        )]
+        num_columns: NonZeroUsize,
+
+        /// The number of rows to generate in each file.
+        #[clap(
+            long = "--num-rows",
+            env = "INFLUXDB_IOX_COMPACTOR_GENERATE_NUM_ROWS",
+            default_value = "1",
+            action
+        )]
+        num_rows: NonZeroUsize,
+    },
+}
+
+#[derive(Debug, Copy, Clone, PartialEq, Eq, PartialOrd, Ord, ValueEnum)]
+pub enum CompactionType {
+    Hot,
+    Cold,
 }
 
 pub async fn command(config: Config) -> Result<()> {
@@ -82,6 +151,7 @@ pub async fn command(config: Config) -> Result<()> {
 
             compactor::handler::run_compactor_once(compactor).await;
         }
+        Command::Generate { .. } => {}
     }
 
     Ok(())
diff --git a/influxdb_iox/tests/end_to_end_cases/compactor.rs b/influxdb_iox/tests/end_to_end_cases/compactor.rs
new file mode 100644
index 0000000000..0116d0fc05
--- /dev/null
+++ b/influxdb_iox/tests/end_to_end_cases/compactor.rs
@@ -0,0 +1,42 @@
+use assert_cmd::Command;
+use predicates::prelude::*;
+use test_helpers_end_to_end::maybe_skip_integration;
+
+#[tokio::test]
+async fn compactor_generate_has_defaults() {
+    let database_url = maybe_skip_integration!();
+
+    Command::cargo_bin("influxdb_iox")
+        .unwrap()
+        .arg("compactor")
+        .arg("generate")
+        .arg("--catalog-dsn")
+        .arg(database_url)
+        .assert()
+        .success();
+}
+
+#[tokio::test]
+async fn compactor_generate_zeroes_are_invalid() {
+    let database_url = maybe_skip_integration!();
+
+    Command::cargo_bin("influxdb_iox")
+        .unwrap()
+        .arg("compactor")
+        .arg("generate")
+        .arg("--catalog-dsn")
+        .arg(database_url)
+        .arg("--num-partitions")
+        .arg("0")
+        .arg("--num-files")
+        .arg("0")
+        .arg("--num-cols")
+        .arg("0")
+        .arg("--num-rows")
+        .arg("0")
+        .assert()
+        .failure()
+        .stderr(predicate::str::contains(
+            "number would be zero for non-zero type",
+        ));
+}
diff --git a/influxdb_iox/tests/end_to_end_cases/mod.rs b/influxdb_iox/tests/end_to_end_cases/mod.rs
index 8628901fcb..99a35bec97 100644
--- a/influxdb_iox/tests/end_to_end_cases/mod.rs
+++ b/influxdb_iox/tests/end_to_end_cases/mod.rs
@@ -3,6 +3,7 @@ mod all_in_one;
 // loading shared libraries: libjemalloc.so.2: cannot open shared object file: No such file or directory"
 #[cfg(not(feature = "heappy"))]
 mod cli;
+mod compactor;
 mod debug;
 mod error;
 mod ingester;