refactor: move Database to client crate behind testing feature (#4059)

* refactor: move Database to client crate behind testing feature Signed-off-by: tison <wander4096@gmail.com> * partial move Signed-off-by: tison <wander4096@gmail.com> * catch up more Signed-off-by: tison <wander4096@gmail.com> * fix imports Signed-off-by: tison <wander4096@gmail.com> * finish Signed-off-by: tison <wander4096@gmail.com> * tidy Signed-off-by: tison <wander4096@gmail.com> --------- Signed-off-by: tison <wander4096@gmail.com>
2026-05-23 16:30:39 +00:00 · 2024-05-28 11:21:43 +08:00
parent 097f62f459
commit 9dd6e033a7
17 changed files with 167 additions and 114 deletions
--- a/src/client/src/database.rs
+++ b/src/client/src/database.rs
@@ -0,0 +1,314 @@
+// Copyright 2023 Greptime Team
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+use api::v1::auth_header::AuthScheme;
+use api::v1::ddl_request::Expr as DdlExpr;
+use api::v1::greptime_database_client::GreptimeDatabaseClient;
+use api::v1::greptime_request::Request;
+use api::v1::query_request::Query;
+use api::v1::{
+    AlterExpr, AuthHeader, CreateTableExpr, DdlRequest, GreptimeRequest, InsertRequests,
+    QueryRequest, RequestHeader,
+};
+use arrow_flight::Ticket;
+use async_stream::stream;
+use common_error::ext::{BoxedError, ErrorExt};
+use common_grpc::flight::{FlightDecoder, FlightMessage};
+use common_query::Output;
+use common_recordbatch::error::ExternalSnafu;
+use common_recordbatch::RecordBatchStreamWrapper;
+use common_telemetry::error;
+use common_telemetry::tracing_context::W3cTrace;
+use futures_util::StreamExt;
+use prost::Message;
+use snafu::{ensure, ResultExt};
+use tonic::transport::Channel;
+
+use crate::error::{ConvertFlightDataSnafu, Error, IllegalFlightMessagesSnafu, ServerSnafu};
+use crate::{from_grpc_response, Client, Result};
+
+#[derive(Clone, Debug, Default)]
+pub struct Database {
+    // The "catalog" and "schema" to be used in processing the requests at the server side.
+    // They are the "hint" or "context", just like how the "database" in "USE" statement is treated in MySQL.
+    // They will be carried in the request header.
+    catalog: String,
+    schema: String,
+    // The dbname follows naming rule as out mysql, postgres and http
+    // protocol. The server treat dbname in priority of catalog/schema.
+    dbname: String,
+    // The time zone indicates the time zone where the user is located.
+    // Some queries need to be aware of the user's time zone to perform some specific actions.
+    timezone: String,
+
+    client: Client,
+    ctx: FlightContext,
+}
+
+pub struct DatabaseClient {
+    pub inner: GreptimeDatabaseClient<Channel>,
+}
+
+fn make_database_client(client: &Client) -> Result<DatabaseClient> {
+    let (_, channel) = client.find_channel()?;
+    Ok(DatabaseClient {
+        inner: GreptimeDatabaseClient::new(channel)
+            .max_decoding_message_size(client.max_grpc_recv_message_size())
+            .max_encoding_message_size(client.max_grpc_send_message_size()),
+    })
+}
+
+impl Database {
+    /// Create database service client using catalog and schema
+    pub fn new(catalog: impl Into<String>, schema: impl Into<String>, client: Client) -> Self {
+        Self {
+            catalog: catalog.into(),
+            schema: schema.into(),
+            dbname: String::default(),
+            timezone: String::default(),
+            client,
+            ctx: FlightContext::default(),
+        }
+    }
+
+    /// Create database service client using dbname.
+    ///
+    /// This API is designed for external usage. `dbname` is:
+    ///
+    /// - the name of database when using GreptimeDB standalone or cluster
+    /// - the name provided by GreptimeCloud or other multi-tenant GreptimeDB
+    /// environment
+    pub fn new_with_dbname(dbname: impl Into<String>, client: Client) -> Self {
+        Self {
+            catalog: String::default(),
+            schema: String::default(),
+            timezone: String::default(),
+            dbname: dbname.into(),
+            client,
+            ctx: FlightContext::default(),
+        }
+    }
+
+    pub fn set_catalog(&mut self, catalog: impl Into<String>) {
+        self.catalog = catalog.into();
+    }
+
+    pub fn catalog(&self) -> &String {
+        &self.catalog
+    }
+
+    pub fn set_schema(&mut self, schema: impl Into<String>) {
+        self.schema = schema.into();
+    }
+
+    pub fn schema(&self) -> &String {
+        &self.schema
+    }
+
+    pub fn set_timezone(&mut self, timezone: impl Into<String>) {
+        self.timezone = timezone.into();
+    }
+
+    pub fn set_auth(&mut self, auth: AuthScheme) {
+        self.ctx.auth_header = Some(AuthHeader {
+            auth_scheme: Some(auth),
+        });
+    }
+
+    pub async fn insert(&self, requests: InsertRequests) -> Result<u32> {
+        self.handle(Request::Inserts(requests)).await
+    }
+
+    async fn handle(&self, request: Request) -> Result<u32> {
+        let mut client = make_database_client(&self.client)?.inner;
+        let request = self.to_rpc_request(request);
+        let response = client.handle(request).await?.into_inner();
+        from_grpc_response(response)
+    }
+
+    #[inline]
+    fn to_rpc_request(&self, request: Request) -> GreptimeRequest {
+        GreptimeRequest {
+            header: Some(RequestHeader {
+                catalog: self.catalog.clone(),
+                schema: self.schema.clone(),
+                authorization: self.ctx.auth_header.clone(),
+                dbname: self.dbname.clone(),
+                timezone: self.timezone.clone(),
+                // TODO(Taylor-lagrange): add client grpc tracing
+                tracing_context: W3cTrace::new(),
+            }),
+            request: Some(request),
+        }
+    }
+
+    pub async fn sql<S>(&self, sql: S) -> Result<Output>
+    where
+        S: AsRef<str>,
+    {
+        self.do_get(Request::Query(QueryRequest {
+            query: Some(Query::Sql(sql.as_ref().to_string())),
+        }))
+        .await
+    }
+
+    pub async fn logical_plan(&self, logical_plan: Vec<u8>) -> Result<Output> {
+        self.do_get(Request::Query(QueryRequest {
+            query: Some(Query::LogicalPlan(logical_plan)),
+        }))
+        .await
+    }
+
+    pub async fn create(&self, expr: CreateTableExpr) -> Result<Output> {
+        self.do_get(Request::Ddl(DdlRequest {
+            expr: Some(DdlExpr::CreateTable(expr)),
+        }))
+        .await
+    }
+
+    pub async fn alter(&self, expr: AlterExpr) -> Result<Output> {
+        self.do_get(Request::Ddl(DdlRequest {
+            expr: Some(DdlExpr::Alter(expr)),
+        }))
+        .await
+    }
+
+    async fn do_get(&self, request: Request) -> Result<Output> {
+        let request = self.to_rpc_request(request);
+        let request = Ticket {
+            ticket: request.encode_to_vec().into(),
+        };
+
+        let mut client = self.client.make_flight_client()?;
+
+        let response = client.mut_inner().do_get(request).await.map_err(|e| {
+            let tonic_code = e.code();
+            let e: Error = e.into();
+            let code = e.status_code();
+            let msg = e.to_string();
+            let error = Error::FlightGet {
+                tonic_code,
+                addr: client.addr().to_string(),
+                source: BoxedError::new(ServerSnafu { code, msg }.build()),
+            };
+            error!(
+                "Failed to do Flight get, addr: {}, code: {}, source: {:?}",
+                client.addr(),
+                tonic_code,
+                error
+            );
+            error
+        })?;
+
+        let flight_data_stream = response.into_inner();
+        let mut decoder = FlightDecoder::default();
+
+        let mut flight_message_stream = flight_data_stream.map(move |flight_data| {
+            flight_data
+                .map_err(Error::from)
+                .and_then(|data| decoder.try_decode(data).context(ConvertFlightDataSnafu))
+        });
+
+        let Some(first_flight_message) = flight_message_stream.next().await else {
+            return IllegalFlightMessagesSnafu {
+                reason: "Expect the response not to be empty",
+            }
+            .fail();
+        };
+
+        let first_flight_message = first_flight_message?;
+
+        match first_flight_message {
+            FlightMessage::AffectedRows(rows) => {
+                ensure!(
+                    flight_message_stream.next().await.is_none(),
+                    IllegalFlightMessagesSnafu {
+                        reason: "Expect 'AffectedRows' Flight messages to be the one and the only!"
+                    }
+                );
+                Ok(Output::new_with_affected_rows(rows))
+            }
+            FlightMessage::Recordbatch(_) | FlightMessage::Metrics(_) => {
+                IllegalFlightMessagesSnafu {
+                    reason: "The first flight message cannot be a RecordBatch or Metrics message",
+                }
+                .fail()
+            }
+            FlightMessage::Schema(schema) => {
+                let stream = Box::pin(stream!({
+                    while let Some(flight_message) = flight_message_stream.next().await {
+                        let flight_message = flight_message
+                            .map_err(BoxedError::new)
+                            .context(ExternalSnafu)?;
+                        match flight_message {
+                            FlightMessage::Recordbatch(record_batch) => yield Ok(record_batch),
+                            FlightMessage::Metrics(_) => {}
+                            FlightMessage::AffectedRows(_) | FlightMessage::Schema(_) => {
+                                yield IllegalFlightMessagesSnafu {reason: format!("A Schema message must be succeeded exclusively by a set of RecordBatch messages, flight_message: {:?}", flight_message)}
+                                        .fail()
+                                        .map_err(BoxedError::new)
+                                        .context(ExternalSnafu);
+                                break;
+                            }
+                        }
+                    }
+                }));
+                let record_batch_stream = RecordBatchStreamWrapper {
+                    schema,
+                    stream,
+                    output_ordering: None,
+                    metrics: Default::default(),
+                };
+                Ok(Output::new_with_stream(Box::pin(record_batch_stream)))
+            }
+        }
+    }
+}
+
+#[derive(Default, Debug, Clone)]
+struct FlightContext {
+    auth_header: Option<AuthHeader>,
+}
+
+#[cfg(test)]
+mod tests {
+    use std::assert_matches::assert_matches;
+
+    use api::v1::auth_header::AuthScheme;
+    use api::v1::{AuthHeader, Basic};
+
+    use super::*;
+
+    #[test]
+    fn test_flight_ctx() {
+        let mut ctx = FlightContext::default();
+        assert!(ctx.auth_header.is_none());
+
+        let basic = AuthScheme::Basic(Basic {
+            username: "u".to_string(),
+            password: "p".to_string(),
+        });
+
+        ctx.auth_header = Some(AuthHeader {
+            auth_scheme: Some(basic),
+        });
+
+        assert_matches!(
+            ctx.auth_header,
+            Some(AuthHeader {
+                auth_scheme: Some(AuthScheme::Basic(_)),
+            })
+        )
+    }
+}
--- a/src/client/src/lib.rs
+++ b/src/client/src/lib.rs
@@ -12,8 +12,12 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.

+#![feature(assert_matches)]
+
 mod client;
 pub mod client_manager;
+#[cfg(feature = "testing")]
+mod database;
 pub mod error;
 pub mod load_balance;
 mod metrics;
@@ -29,6 +33,8 @@ pub use common_recordbatch::{RecordBatches, SendableRecordBatchStream};
 use snafu::OptionExt;

 pub use self::client::Client;
+#[cfg(feature = "testing")]
+pub use self::database::Database;
 pub use self::error::{Error, Result};
 use crate::error::{IllegalDatabaseResponseSnafu, ServerSnafu};

--- a/src/cmd/Cargo.toml
+++ b/src/cmd/Cargo.toml
@@ -80,6 +80,7 @@ tracing-appender = "0.2"
 tikv-jemallocator = "0.5"

 [dev-dependencies]
+client = { workspace = true, features = ["testing"] }
 common-test-util.workspace = true
 serde.workspace = true
 temp-env = "0.3"
--- a/src/cmd/src/cli.rs
+++ b/src/cmd/src/cli.rs
@@ -22,8 +22,8 @@ mod helper;

 // Wait for https://github.com/GreptimeTeam/greptimedb/issues/2373
 #[allow(unused)]
-// mod repl;
-// TODO(weny): Removes it
+mod repl;
+// TODO(tisonkun): migrate deprecated methods
 #[allow(deprecated)]
 mod upgrade;

@@ -31,8 +31,8 @@ use async_trait::async_trait;
 use bench::BenchTableMetadataCommand;
 use clap::Parser;
 use common_telemetry::logging::{LoggingOptions, TracingOptions};
+pub use repl::Repl;
 use tracing_appender::non_blocking::WorkerGuard;
-// pub use repl::Repl;
 use upgrade::UpgradeCommand;

 use self::export::ExportCommand;
--- a/src/cmd/src/cli/export.rs
+++ b/src/cmd/src/cli/export.rs
@@ -434,3 +434,80 @@ fn split_database(database: &str) -> Result<(String, Option<String>)> {
        Ok((catalog.to_string(), Some(schema.to_string())))
    }
 }
+
+#[cfg(test)]
+mod tests {
+    use clap::Parser;
+    use client::{Client, Database};
+    use common_catalog::consts::{DEFAULT_CATALOG_NAME, DEFAULT_SCHEMA_NAME};
+    use common_telemetry::logging::LoggingOptions;
+
+    use crate::error::Result as CmdResult;
+    use crate::options::GlobalOptions;
+    use crate::{cli, standalone, App};
+
+    #[tokio::test(flavor = "multi_thread")]
+    async fn test_export_create_table_with_quoted_names() -> CmdResult<()> {
+        let output_dir = tempfile::tempdir().unwrap();
+
+        let standalone = standalone::Command::parse_from([
+            "standalone",
+            "start",
+            "--data-home",
+            &*output_dir.path().to_string_lossy(),
+        ]);
+
+        let standalone_opts = standalone.load_options(&GlobalOptions::default()).unwrap();
+        let mut instance = standalone.build(standalone_opts).await?;
+        instance.start().await?;
+
+        let client = Client::with_urls(["127.0.0.1:4001"]);
+        let database = Database::new(DEFAULT_CATALOG_NAME, DEFAULT_SCHEMA_NAME, client);
+        database
+            .sql(r#"CREATE DATABASE "cli.export.create_table";"#)
+            .await
+            .unwrap();
+        database
+            .sql(
+                r#"CREATE TABLE "cli.export.create_table"."a.b.c"(
+                        ts TIMESTAMP,
+                        TIME INDEX (ts)
+                    ) engine=mito;
+                "#,
+            )
+            .await
+            .unwrap();
+
+        let output_dir = tempfile::tempdir().unwrap();
+        let cli = cli::Command::parse_from([
+            "cli",
+            "export",
+            "--addr",
+            "127.0.0.1:4000",
+            "--output-dir",
+            &*output_dir.path().to_string_lossy(),
+            "--target",
+            "create-table",
+        ]);
+        let mut cli_app = cli.build(LoggingOptions::default()).await?;
+        cli_app.start().await?;
+
+        instance.stop().await?;
+
+        let output_file = output_dir
+            .path()
+            .join("greptime-cli.export.create_table.sql");
+        let res = std::fs::read_to_string(output_file).unwrap();
+        let expect = r#"CREATE TABLE IF NOT EXISTS "a.b.c" (
+  "ts" TIMESTAMP(3) NOT NULL,
+  TIME INDEX ("ts")
+)
+
+ENGINE=mito
+;
+"#;
+        assert_eq!(res.trim(), expect.trim());
+
+        Ok(())
+    }
+}
--- a/src/cmd/src/cli/repl.rs
+++ b/src/cmd/src/cli/repl.rs
@@ -16,14 +16,18 @@ use std::path::PathBuf;
 use std::sync::Arc;
 use std::time::Instant;

-use catalog::kvbackend::{
-    CachedMetaKvBackend, CachedMetaKvBackendBuilder, KvBackendCatalogManager,
+use cache::{
+    build_fundamental_cache_registry, with_default_composite_cache_registry, TABLE_CACHE_NAME,
+    TABLE_ROUTE_CACHE_NAME,
 };
-use client::{Client, OutputData, DEFAULT_CATALOG_NAME, DEFAULT_SCHEMA_NAME};
+use catalog::kvbackend::{
+    CachedMetaKvBackend, CachedMetaKvBackendBuilder, KvBackendCatalogManager, MetaKvBackend,
+};
+use client::{Client, Database, OutputData, DEFAULT_CATALOG_NAME, DEFAULT_SCHEMA_NAME};
 use common_base::Plugins;
 use common_config::Mode;
 use common_error::ext::ErrorExt;
-use common_meta::cache_invalidator::MultiCacheInvalidator;
+use common_meta::cache::{CacheRegistryBuilder, LayeredCacheRegistryBuilder};
 use common_query::Output;
 use common_recordbatch::RecordBatches;
 use common_telemetry::debug;
@@ -38,12 +42,13 @@ use query::QueryEngine;
 use rustyline::error::ReadlineError;
 use rustyline::Editor;
 use session::context::QueryContext;
-use snafu::ResultExt;
+use snafu::{OptionExt, ResultExt};
 use substrait::{DFLogicalSubstraitConvertor, SubstraitPlan};

 use crate::cli::cmd::ReplCommand;
 use crate::cli::helper::RustylineHelper;
 use crate::cli::AttachCommand;
+use crate::error;
 use crate::error::{
    CollectRecordBatchesSnafu, ParseSqlSnafu, PlanStatementSnafu, PrettyPrintRecordBatchesSnafu,
    ReadlineSnafu, ReplCreationSnafu, RequestDatabaseSnafu, Result, StartMetaClientSnafu,
@@ -257,19 +262,42 @@ async fn create_query_engine(meta_addr: &str) -> Result<DatafusionQueryEngine> {

    let cached_meta_backend =
        Arc::new(CachedMetaKvBackendBuilder::new(meta_client.clone()).build());
-    let multi_cache_invalidator = Arc::new(MultiCacheInvalidator::with_invalidators(vec![
-        cached_meta_backend.clone(),
-    ]));
-    let catalog_list = KvBackendCatalogManager::new(
+    let layered_cache_builder = LayeredCacheRegistryBuilder::default().add_cache_registry(
+        CacheRegistryBuilder::default()
+            .add_cache(cached_meta_backend.clone())
+            .build(),
+    );
+    let fundamental_cache_registry =
+        build_fundamental_cache_registry(Arc::new(MetaKvBackend::new(meta_client.clone())));
+    let layered_cache_registry = Arc::new(
+        with_default_composite_cache_registry(
+            layered_cache_builder.add_cache_registry(fundamental_cache_registry),
+        )
+        .context(error::BuildCacheRegistrySnafu)?
+        .build(),
+    );
+
+    let table_cache = layered_cache_registry
+        .get()
+        .context(error::CacheRequiredSnafu {
+            name: TABLE_CACHE_NAME,
+        })?;
+    let table_route_cache = layered_cache_registry
+        .get()
+        .context(error::CacheRequiredSnafu {
+            name: TABLE_ROUTE_CACHE_NAME,
+        })?;
+    let catalog_manager = KvBackendCatalogManager::new(
        Mode::Distributed,
        Some(meta_client.clone()),
        cached_meta_backend.clone(),
-        multi_cache_invalidator,
+        table_cache,
+        table_route_cache,
    )
    .await;
    let plugins: Plugins = Default::default();
    let state = Arc::new(QueryEngineState::new(
-        catalog_list,
+        catalog_manager,
        None,
        None,
        None,
--- a/src/cmd/src/error.rs
+++ b/src/cmd/src/error.rs
@@ -163,6 +163,15 @@ pub enum Error {
        location: Location,
    },

+    #[snafu(display("Failed to request database, sql: {sql}"))]
+    RequestDatabase {
+        sql: String,
+        #[snafu(source)]
+        source: client::Error,
+        #[snafu(implicit)]
+        location: Location,
+    },
+
    #[snafu(display("Failed to collect RecordBatches"))]
    CollectRecordBatches {
        #[snafu(implicit)]
@@ -354,6 +363,7 @@ impl ErrorExt for Error {
            Error::ReplCreation { .. } | Error::Readline { .. } | Error::HttpQuerySql { .. } => {
                StatusCode::Internal
            }
+            Error::RequestDatabase { source, .. } => source.status_code(),
            Error::CollectRecordBatches { source, .. }
            | Error::PrettyPrintRecordBatches { source, .. } => source.status_code(),
            Error::StartMetaClient { source, .. } => source.status_code(),
--- a/src/servers/Cargo.toml
+++ b/src/servers/Cargo.toml
@@ -111,7 +111,7 @@ tikv-jemalloc-ctl = { version = "0.5", features = ["use_std"] }
 [dev-dependencies]
 auth = { workspace = true, features = ["testing"] }
 catalog = { workspace = true, features = ["testing"] }
-client.workspace = true
+client = { workspace = true, features = ["testing"] }
 common-base.workspace = true
 common-test-util.workspace = true
 criterion = "0.4"
@@ -125,8 +125,6 @@ serde_json.workspace = true
 session = { workspace = true, features = ["testing"] }
 table.workspace = true
 tempfile = "3.0.0"
-# TODO depend `Database` client
-tests-integration.workspace = true
 tokio-postgres = "0.7"
 tokio-postgres-rustls = "0.11"
 tokio-test = "0.4"
--- a/src/servers/tests/grpc/mod.rs
+++ b/src/servers/tests/grpc/mod.rs
@@ -21,7 +21,7 @@ use arrow_flight::flight_service_server::{FlightService, FlightServiceServer};
 use async_trait::async_trait;
 use auth::tests::MockUserProvider;
 use auth::UserProviderRef;
-use client::{Client, DEFAULT_CATALOG_NAME, DEFAULT_SCHEMA_NAME};
+use client::{Client, Database, DEFAULT_CATALOG_NAME, DEFAULT_SCHEMA_NAME};
 use common_runtime::{Builder as RuntimeBuilder, Runtime};
 use servers::error::{Result, StartGrpcSnafu, TcpBindSnafu};
 use servers::grpc::flight::FlightCraftWrapper;
@@ -31,7 +31,6 @@ use servers::server::Server;
 use snafu::ResultExt;
 use table::test_util::MemTable;
 use table::TableRef;
-use tests_integration::database::Database;
 use tokio::net::TcpListener;
 use tokio_stream::wrappers::TcpListenerStream;
 use tonic::codec::CompressionEncoding;