Schemas and Databases search through the query engine (#839)

* Done

* Schemas ready

* Extended databases view and volumes view

* Databases ready

* cargo clippy + fmt

* cargo clippy + fmt

* test fix

* Requested changes

* Requested changes

Author: DanCodedThis

Cargo.lock

@@ -2,6 +2,7 @@ use crate::error::ErrorResponse;
 use crate::error::IntoStatusCode;
 use axum::Json;
 use axum::response::IntoResponse;
+use core_executor::error::ExecutionError;
 use core_metastore::error::MetastoreError;
 use http::StatusCode;
 use snafu::prelude::*;
@@ -20,7 +21,7 @@ pub enum DatabasesAPIError {
     #[snafu(display("Update database error: {source}"))]
     Update { source: MetastoreError },
     #[snafu(display("Get databases error: {source}"))]
-    List { source: MetastoreError },
+    List { source: ExecutionError },
 }
 
 // Select which status code to return.

crates/api-ui/src/databases/error.rs

@@ -1,20 +1,24 @@
-use crate::databases::models::DatabasesParameters;
+use crate::OrderDirection;
 use crate::state::AppState;
 use crate::{
+    SearchParameters,
     databases::error::{DatabasesAPIError, DatabasesResult},
     databases::models::{
         Database, DatabaseCreatePayload, DatabaseCreateResponse, DatabaseResponse,
         DatabaseUpdatePayload, DatabaseUpdateResponse, DatabasesResponse,
     },
+    downcast_string_column,
     error::ErrorResponse,
 };
+use api_sessions::DFSessionId;
 use axum::{
     Json,
     extract::{Path, Query, State},
 };
+use core_executor::models::QueryResultData;
+use core_executor::query::QueryContext;
 use core_metastore::Database as MetastoreDatabase;
 use core_metastore::error::MetastoreError;
-use core_utils::scan_iterator::ScanIterator;
 use utoipa::OpenApi;
 use validator::Validate;
 
@@ -81,11 +85,7 @@ pub async fn create_database(
         .create_database(&database.ident.clone(), database)
         .await
         .map_err(|e| DatabasesAPIError::Create { source: e })
-        .map(|o| {
-            Json(DatabaseCreateResponse {
-                data: o.data.into(),
-            })
-        })
+        .map(|o| Json(DatabaseCreateResponse { data: o.into() }))
 }
 
 #[utoipa::path(
@@ -113,9 +113,7 @@ pub async fn get_database(
     Path(database_name): Path<String>,
 ) -> DatabasesResult<Json<DatabaseResponse>> {
     match state.metastore.get_database(&database_name).await {
-        Ok(Some(db)) => Ok(Json(DatabaseResponse {
-            data: db.data.into(),
-        })),
+        Ok(Some(db)) => Ok(Json(DatabaseResponse { data: db.into() })),
         Ok(None) => Err(DatabasesAPIError::Get {
             source: MetastoreError::DatabaseNotFound {
                 db: database_name.clone(),
@@ -194,20 +192,18 @@ pub async fn update_database(
         .update_database(&database_name, database)
         .await
         .map_err(|e| DatabasesAPIError::Update { source: e })
-        .map(|o| {
-            Json(DatabaseUpdateResponse {
-                data: o.data.into(),
-            })
-        })
+        .map(|o| Json(DatabaseUpdateResponse { data: o.into() }))
 }
 
 #[utoipa::path(
     get,
     operation_id = "getDatabases",
     params(
-        ("cursor" = Option<String>, Query, description = "Databases cursor"),
+        ("offset" = Option<usize>, Query, description = "Databases offset"),
         ("limit" = Option<usize>, Query, description = "Databases limit"),
-        ("search" = Option<String>, Query, description = "Databases search (start with)"),
+        ("search" = Option<String>, Query, description = "Databases search"),
+        ("order_by" = Option<String>, Query, description = "Order by: database_name (default), volume_name, created_at, updated_at"),
+        ("order_direction" = Option<OrderDirection>, Query, description = "Order direction: ASC, DESC (default)"),
     ),
     tags = ["databases"],
     path = "/ui/databases",
@@ -223,30 +219,56 @@ pub async fn update_database(
     )
 )]
 #[tracing::instrument(level = "debug", skip(state), err, ret(level = tracing::Level::TRACE))]
+#[allow(clippy::unwrap_used)]
 pub async fn list_databases(
-    Query(parameters): Query<DatabasesParameters>,
+    DFSessionId(session_id): DFSessionId,
+    Query(parameters): Query<SearchParameters>,
     State(state): State<AppState>,
 ) -> DatabasesResult<Json<DatabasesResponse>> {
-    state
-        .metastore
-        .iter_databases()
-        .cursor(parameters.cursor.clone())
-        .limit(parameters.limit)
-        .token(parameters.search)
-        .collect()
+    let context = QueryContext::default();
+    let sql_string = "SELECT * FROM slatedb.public.databases".to_string();
+    let sql_string = parameters.search.map_or_else(|| sql_string.clone(), |search|
+        format!("{sql_string} WHERE (database_name ILIKE '%{search}%' OR volume_name ILIKE '%{search}%')")
+    );
+    let sql_string = parameters.order_by.map_or_else(
+        || format!("{sql_string} ORDER BY database_name"),
+        |order_by| format!("{sql_string} ORDER BY {order_by}"),
+    );
+    let sql_string = parameters.order_direction.map_or_else(
+        || format!("{sql_string} DESC"),
+        |order_direction| format!("{sql_string} {order_direction}"),
+    );
+    let sql_string = parameters.offset.map_or_else(
+        || sql_string.clone(),
+        |offset| format!("{sql_string} OFFSET {offset}"),
+    );
+    let sql_string = parameters.limit.map_or_else(
+        || sql_string.clone(),
+        |limit| format!("{sql_string} LIMIT {limit}"),
+    );
+    let QueryResultData { records, .. } = state
+        .execution_svc
+        .query(&session_id, sql_string.as_str(), context)
         .await
-        .map_err(|e| DatabasesAPIError::List {
-            source: MetastoreError::UtilSlateDB { source: e },
-        })
-        .map(|o| {
-            let next_cursor = o
-                .iter()
-                .last()
-                .map_or(String::new(), |rw_object| rw_object.ident.clone());
-            Json(DatabasesResponse {
-                items: o.into_iter().map(|x| x.data.into()).collect(),
-                current_cursor: parameters.cursor,
-                next_cursor,
-            })
-        })
+        .map_err(|e| DatabasesAPIError::List { source: e })?;
+    let mut items = Vec::new();
+    for record in records {
+        let database_names = downcast_string_column(&record, "database_name")
+            .map_err(|e| DatabasesAPIError::List { source: e })?;
+        let volume_names = downcast_string_column(&record, "volume_name")
+            .map_err(|e| DatabasesAPIError::List { source: e })?;
+        let created_at_timestamps = downcast_string_column(&record, "created_at")
+            .map_err(|e| DatabasesAPIError::List { source: e })?;
+        let updated_at_timestamps = downcast_string_column(&record, "updated_at")
+            .map_err(|e| DatabasesAPIError::List { source: e })?;
+        for i in 0..record.num_rows() {
+            items.push(Database {
+                name: database_names.value(i).to_string(),
+                volume: volume_names.value(i).to_string(),
+                created_at: created_at_timestamps.value(i).to_string(),
+                updated_at: updated_at_timestamps.value(i).to_string(),
+            });
+        }
+    }
+    Ok(Json(DatabasesResponse { items }))
 }

crates/api-ui/src/databases/handlers.rs

@@ -1,19 +1,35 @@
-use crate::default_limit;
+use core_metastore::RwObject;
 use core_metastore::models::Database as MetastoreDatabase;
 use serde::{Deserialize, Serialize};
-use utoipa::{IntoParams, ToSchema};
+use utoipa::ToSchema;
 
 #[derive(Debug, Clone, Serialize, Deserialize, ToSchema, Eq, PartialEq)]
 pub struct Database {
     pub name: String,
     pub volume: String,
+    pub created_at: String,
+    pub updated_at: String,
 }
 
 impl From<MetastoreDatabase> for Database {
     fn from(db: MetastoreDatabase) -> Self {
         Self {
             name: db.ident,
             volume: db.volume,
+            //TODO: fix this, we must use a different payload or change the test suite for dbs
+            created_at: "ERROR".to_string(),
+            updated_at: "ERROR".to_string(),
+        }
+    }
+}
+
+impl From<RwObject<MetastoreDatabase>> for Database {
+    fn from(db: RwObject<MetastoreDatabase>) -> Self {
+        Self {
+            name: db.data.ident,
+            volume: db.data.volume,
+            created_at: db.created_at.to_string(),
+            updated_at: db.updated_at.to_string(),
         }
     }
 }
@@ -70,14 +86,4 @@ pub struct DatabaseResponse {
 #[serde(rename_all = "camelCase")]
 pub struct DatabasesResponse {
     pub items: Vec<Database>,
-    pub current_cursor: Option<String>,
-    pub next_cursor: String,
-}
-
-#[derive(Debug, Deserialize, ToSchema, IntoParams)]
-pub struct DatabasesParameters {
-    pub cursor: Option<String>,
-    #[serde(default = "default_limit")]
-    pub limit: Option<u16>,
-    pub search: Option<String>,
 }

crates/api-ui/src/databases/models.rs

@@ -1,3 +1,10 @@
+use core_executor::error::ExecutionError;
+use datafusion::arrow::array::{RecordBatch, StringArray};
+use datafusion::common::DataFusionError;
+use serde::Deserialize;
+use std::fmt::Display;
+use utoipa::{IntoParams, ToSchema};
+
 pub mod auth;
 pub mod config;
 pub mod dashboard;
@@ -21,3 +28,48 @@ pub mod worksheets;
 const fn default_limit() -> Option<u16> {
     Some(250)
 }
+
+#[derive(Debug, Deserialize, ToSchema, IntoParams)]
+#[serde(rename_all = "camelCase")]
+pub struct SearchParameters {
+    pub offset: Option<usize>,
+    #[serde(default = "default_limit")]
+    pub limit: Option<u16>,
+    pub search: Option<String>,
+    pub order_by: Option<String>,
+    pub order_direction: Option<OrderDirection>,
+}
+
+#[derive(Debug, Deserialize, ToSchema)]
+#[serde(rename_all = "UPPERCASE")]
+pub enum OrderDirection {
+    ASC,
+    DESC,
+}
+
+impl Default for OrderDirection {
+    fn default() -> Self {
+        Self::DESC
+    }
+}
+
+impl Display for OrderDirection {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        match self {
+            Self::ASC => write!(f, "ASC"),
+            Self::DESC => write!(f, "DESC"),
+        }
+    }
+}
+
+fn downcast_string_column<'a>(
+    batch: &'a RecordBatch,
+    name: &str,
+) -> Result<&'a StringArray, ExecutionError> {
+    batch
+        .column_by_name(name)
+        .and_then(|col| col.as_any().downcast_ref::<StringArray>())
+        .ok_or_else(|| ExecutionError::DataFusion {
+            source: DataFusionError::Internal(format!("Missing or invalid column: '{name}'")),
+        })
+}

crates/api-ui/src/lib.rs

@@ -1,3 +1,4 @@
+use crate::downcast_string_column;
 use crate::error::ErrorResponse;
 use crate::navigation_trees::error::{NavigationTreesAPIError, NavigationTreesResult};
 use crate::navigation_trees::models::{
@@ -8,9 +9,7 @@ use crate::state::AppState;
 use api_sessions::DFSessionId;
 use axum::extract::Query;
 use axum::{Json, extract::State};
-use core_executor::{error::ExecutionError, models::QueryResultData, query::QueryContext};
-use datafusion::arrow::array::{RecordBatch, StringArray};
-use datafusion::common::DataFusionError;
+use core_executor::{models::QueryResultData, query::QueryContext};
 use std::collections::BTreeMap;
 use utoipa::OpenApi;
 
@@ -76,11 +75,14 @@ pub async fn get_navigation_trees(
         BTreeMap::new();
 
     for batch in tree_batches {
-        let databases = downcast_string_column(&batch, "database")?;
-        let schemas = downcast_string_column(&batch, "schema")?;
-        let tables = downcast_string_column(&batch, "table")?;
-        let table_types = downcast_string_column(&batch, "table_type")?;
-
+        let databases = downcast_string_column(&batch, "database")
+            .map_err(|e| NavigationTreesAPIError::Execution { source: e })?;
+        let schemas = downcast_string_column(&batch, "schema")
+            .map_err(|e| NavigationTreesAPIError::Execution { source: e })?;
+        let tables = downcast_string_column(&batch, "table")
+            .map_err(|e| NavigationTreesAPIError::Execution { source: e })?;
+        let table_types = downcast_string_column(&batch, "table_type")
+            .map_err(|e| NavigationTreesAPIError::Execution { source: e })?;
         for j in 0..batch.num_rows() {
             let database = databases.value(j).to_string();
             let schema = schemas.value(j).to_string();
@@ -133,17 +135,3 @@ pub async fn get_navigation_trees(
 
     Ok(Json(NavigationTreesResponse { items }))
 }
-
-fn downcast_string_column<'a>(
-    batch: &'a RecordBatch,
-    name: &str,
-) -> Result<&'a StringArray, NavigationTreesAPIError> {
-    batch
-        .column_by_name(name)
-        .and_then(|col| col.as_any().downcast_ref::<StringArray>())
-        .ok_or_else(|| NavigationTreesAPIError::Execution {
-            source: ExecutionError::DataFusion {
-                source: DataFusionError::Internal(format!("Missing or invalid column: '{name}'")),
-            },
-        })
-}

crates/api-ui/src/navigation_trees/handlers.rs

@@ -21,7 +21,7 @@ pub enum SchemasAPIError {
     #[snafu(display("Update schema error: {source}"))]
     Update { source: MetastoreError },
     #[snafu(display("Get schemas error: {source}"))]
-    List { source: MetastoreError },
+    List { source: ExecutionError },
 }
 
 // Select which status code to return.