Improve dash query prompt placeholder

ulagbulag · Oct 21, 2023 · 83a9be3 · 83a9be3
1 parent 040f021
commit 83a9be3
Show file tree

Hide file tree

Showing 2 changed files with 35 additions and 24 deletions.
diff --git a/dash/query/cli/src/main.rs b/dash/query/cli/src/main.rs
@@ -1,8 +1,8 @@
 use std::{cell::RefCell, process::exit};
 
-use anyhow::{anyhow, Result};
+use anyhow::{anyhow, bail, Result};
 use clap::{ArgAction, Parser};
-use dash_query_provider::{QueryClient, QueryClientArgs};
+use dash_query_provider::{Name, QueryClient, QueryClientArgs};
 use futures::TryStreamExt;
 use inquire::{autocompletion::Replacement, Autocomplete, CustomUserError, Text};
 use serde_json::Value;
@@ -34,16 +34,22 @@ async fn try_main() -> Result<()> {
     let QueryArgs { client, sql } = QueryArgs::parse();
     let client = QueryClient::try_new(&client).await?;
 
+    let tables = client.list_table_names();
+    let table_sample = match tables.first() {
+        Some(table) => table.clone(),
+        None => bail!("None of the tables are detected!"),
+    };
+
     match sql {
         Some(sql) => run_query(&client, &sql).await,
         None => {
-            try_main_interactive(client).await;
+            try_main_interactive(client, table_sample).await;
             Ok(())
         }
     }
 }
 
-async fn try_main_interactive(client: QueryClient) {
+async fn try_main_interactive(client: QueryClient, table_sample: Name) {
     #[derive(Clone, Default)]
     struct History(RefCell<Vec<String>>);
 
@@ -67,11 +73,13 @@ async fn try_main_interactive(client: QueryClient) {
         }
     }
 
+    let placeholder = format!("SELECT * FROM {table_sample};");
+
     let history = History::default();
     loop {
         let sql = match Text::new("sql> ")
             .with_autocomplete(history.clone())
-            .with_placeholder("SELECT * FROM model;")
+            .with_placeholder(&placeholder)
             .prompt()
         {
             Ok(sql) => {

diff --git a/dash/query/provider/src/lib.rs b/dash/query/provider/src/lib.rs
@@ -4,13 +4,11 @@ use dash_api::{
     model_storage_binding::{ModelStorageBindingCrd, ModelStorageBindingState},
     storage::ModelStorageKindSpec,
 };
-use dash_pipe_provider::{
-    storage::{
-        lakehouse::{decoder::TryIntoTableDecoder, StorageContext},
-        StorageS3Args, Stream,
-    },
-    Name,
+use dash_pipe_provider::storage::{
+    lakehouse::{decoder::TryIntoTableDecoder, StorageContext},
+    StorageS3Args, Stream,
 };
+pub use dash_pipe_provider::Name;
 use dash_provider::storage::ObjectStorageRef;
 use datafusion::prelude::DataFrame;
 use futures::future::try_join_all;
@@ -28,23 +26,28 @@ pub struct QueryClientArgs {
 
 pub struct QueryClient {
     ctx: StorageContext,
+    tables: Vec<Name>,
 }
 
 impl QueryClient {
     pub async fn try_new(args: &QueryClientArgs) -> Result<Self> {
-        Ok(Self {
-            ctx: {
-                let ctx = StorageContext::default();
-                for (model, args) in load_models(args.namespace.as_deref()).await? {
-                    let (_, _, has_inited) =
-                        ctx.register_table_with_name(args, &model, None).await?;
-                    if !has_inited {
-                        warn!("Model {model:?} is not inited yet; skipping...");
-                    }
-                }
-                ctx
-            },
-        })
+        let ctx = StorageContext::default();
+        let mut tables = vec![];
+
+        for (model, args) in load_models(args.namespace.as_deref()).await? {
+            let (table, _, has_inited) = ctx.register_table_with_name(args, &model, None).await?;
+            if has_inited {
+                tables.push(table);
+            } else {
+                warn!("Model {model:?} is not inited yet; skipping...");
+            }
+        }
+
+        Ok(Self { ctx, tables })
+    }
+
+    pub fn list_table_names(&self) -> &[Name] {
+        &self.tables
     }
 
     pub async fn sql(&self, sql: &str) -> Result<DataFrame> {