POC: datafusion-cli instrumented object store

datafusion-cli/Cargo.toml

@@ -39,6 +39,7 @@ arrow = { workspace = true }
 async-trait = { workspace = true }
 aws-config = "1.8.6"
 aws-credential-types = "1.2.6"
+chrono = { workspace = true }
 clap = { version = "4.5.47", features = ["derive", "cargo"] }
 datafusion = { workspace = true, features = [
     "avro",

datafusion-cli/src/command.rs

@@ -46,6 +46,7 @@ pub enum Command {
     SearchFunctions(String),
     QuietMode(Option<bool>),
     OutputFormat(Option<String>),
+    ObjectStoreProfileMode(Option<String>),
 }
 
 pub enum OutputFormat {
@@ -124,6 +125,29 @@ impl Command {
             Self::OutputFormat(_) => exec_err!(
                 "Unexpected change output format, this should be handled outside"
             ),
+            Self::ObjectStoreProfileMode(mode) => {
+                if let Some(mode) = mode {
+                    print_options.object_store_profile_mode = mode
+                        .parse()
+                        .map_err(|_| DataFusionError::Execution(
+                                format!("Failed to parse input: {mode}. Valid options are disabled, summary, trace")
+                                ))?;
+                    print_options
+                        .instrumented_registry
+                        .set_instrument_mode(print_options.object_store_profile_mode);
+                    println!(
+                        "ObjectStore Profile mode set to {:?}",
+                        print_options.object_store_profile_mode
+                    );
+                } else {
+                    println!(
+                        "ObjectStore Profile mode is {:?}",
+                        print_options.object_store_profile_mode
+                    );
+                }
+
+                Ok(())
+            }
         }
     }
 
@@ -142,11 +166,15 @@ impl Command {
             Self::OutputFormat(_) => {
                 ("\\pset [NAME [VALUE]]", "set table output option\n(format)")
             }
+            Self::ObjectStoreProfileMode(_) => (
+                "\\object_store_profiling (disabled|summary|trace)",
+                "print or set object store profile mode",
+            ),
         }
     }
 }
 
-const ALL_COMMANDS: [Command; 9] = [
+const ALL_COMMANDS: [Command; 10] = [
     Command::ListTables,
     Command::DescribeTableStmt(String::new()),
     Command::Quit,
@@ -156,6 +184,7 @@ const ALL_COMMANDS: [Command; 9] = [
     Command::SearchFunctions(String::new()),
     Command::QuietMode(None),
     Command::OutputFormat(None),
+    Command::ObjectStoreProfileMode(None),
 ];
 
 fn all_commands_info() -> RecordBatch {
@@ -206,6 +235,10 @@ impl FromStr for Command {
                 Self::OutputFormat(Some(subcommand.to_string()))
             }
             ("pset", None) => Self::OutputFormat(None),
+            ("object_store_profiling", Some(mode)) => {
+                Self::ObjectStoreProfileMode(Some(mode.to_string()))
+            }
+            ("object_store_profiling", None) => Self::ObjectStoreProfileMode(None),
             _ => return Err(()),
         })
     }

datafusion-cli/src/main.rs

@@ -27,10 +27,14 @@ use datafusion::execution::context::SessionConfig;
 use datafusion::execution::memory_pool::{
     FairSpillPool, GreedyMemoryPool, MemoryPool, TrackConsumersPool,
 };
+use datafusion::execution::object_store::DefaultObjectStoreRegistry;
 use datafusion::execution::runtime_env::RuntimeEnvBuilder;
 use datafusion::prelude::SessionContext;
 use datafusion_cli::catalog::DynamicObjectStoreCatalog;
 use datafusion_cli::functions::{MetadataCacheFunc, ParquetMetadataFunc};
+use datafusion_cli::object_storage::instrumented::{
+    InstrumentedObjectStoreMode, InstrumentedObjectStoreRegistry,
+};
 use datafusion_cli::{
     exec,
     pool_type::PoolType,
@@ -144,6 +148,13 @@ struct Args {
         value_parser(extract_disk_limit)
     )]
     disk_limit: Option<usize>,
+
+    #[clap(
+        long,
+        help = "Specify the default object_store_profiling mode, defaults to 'disabled'.\n[possible values: disabled, summary, trace]",
+        default_value_t = InstrumentedObjectStoreMode::Disabled
+    )]
+    object_store_profiling: InstrumentedObjectStoreMode,
 }
 
 #[tokio::main]
@@ -205,6 +216,12 @@ async fn main_inner() -> Result<()> {
         rt_builder = rt_builder.with_disk_manager_builder(builder);
     }
 
+    let instrumented_registry = Arc::new(InstrumentedObjectStoreRegistry::new(
+        Arc::new(DefaultObjectStoreRegistry::new()),
+        args.object_store_profiling,
+    ));
+    rt_builder = rt_builder.with_object_store_registry(instrumented_registry.clone());
+
     let runtime_env = rt_builder.build_arc()?;
 
     // enable dynamic file query
@@ -232,6 +249,8 @@ async fn main_inner() -> Result<()> {
         quiet: args.quiet,
         maxrows: args.maxrows,
         color: args.color,
+        object_store_profile_mode: args.object_store_profiling,
+        instrumented_registry: Arc::clone(&instrumented_registry),
     };
 
     let commands = args.command;