refactor: add share endpoint manager to query share endpoint (#10843)

* fix: add share endpoing from share config; init endpoint only needed

* fix: fix test fail

* fix: fix test fail, remove stateful share test now

* refactor: add share endpoint manager to query share endpoint

* fix: make clippy happy

* fix: make clippy happy

* refactor: refactor by review comment

Author: lichuang

Cargo.lock

@@ -52,7 +52,6 @@ common-meta-stoerr = { path = "../meta/stoerr" }
 common-meta-store = { path = "../meta/store" }
 common-meta-types = { path = "../meta/types" }
 common-metrics = { path = "../common/metrics" }
-common-sharing = { path = "../query/sharing" }
 common-tracing = { path = "../common/tracing" }
 databend-meta = { path = "../meta/service" }
 databend-query = { path = "../query/service" }

src/binaries/Cargo.toml

@@ -30,7 +30,6 @@ use common_exception::ErrorCode;
 use common_exception::Result;
 use common_meta_client::MIN_METASRV_SEMVER;
 use common_metrics::init_default_metrics_recorder;
-use common_sharing::init_share_endpoint_config;
 use common_tracing::set_panic_hook;
 use databend_query::api::HttpService;
 use databend_query::api::RpcService;
@@ -227,8 +226,6 @@ async fn main_entrypoint() -> Result<()> {
         );
     }
 
-    init_share_endpoint_config(&conf).await?;
-
     // Print information to users.
     println!("Databend Query");
     println!();

src/binaries/query/main.rs

@@ -168,15 +168,6 @@ impl<KV: kvapi::KVApi<Error = MetaError>> SchemaApi for KV {
             )));
         }
 
-        // if create a database from a share, check if the share exists and grant access, update share_meta.
-        if let Some(from_share) = &req.meta.from_share {
-            if from_share.tenant == req.name_ident.tenant {
-                return Err(KVAppError::AppError(AppError::WrongShare(WrongShare::new(
-                    req.name_ident.to_string_key(),
-                ))));
-            }
-        }
-
         let mut retry = 0;
         while retry < TXN_MAX_RETRY_TIMES {
             retry += 1;

src/meta/api/src/schema_api_impl.rs

@@ -984,32 +984,19 @@ pub fn statement(i: Input) -> IResult<StatementMsg> {
              ~ #ident
              ~ URL ~ "=" ~ #share_endpoint_uri_location
              ~ TENANT ~ "=" ~ #ident
-             ~ ARGS ~ "=" ~ #options
+             ~ ( ARGS ~ "=" ~ #options)?
              ~ ( COMMENT ~ "=" ~ #literal_string)?
         },
-        |(
-            _,
-            _,
-            _,
-            opt_if_not_exists,
-            endpoint,
-            _,
-            _,
-            url,
-            _,
-            _,
-            tenant,
-            _,
-            _,
-            args,
-            comment_opt,
-        )| {
+        |(_, _, _, opt_if_not_exists, endpoint, _, _, url, _, _, tenant, args_opt, comment_opt)| {
             Statement::CreateShareEndpoint(CreateShareEndpointStmt {
                 if_not_exists: opt_if_not_exists.is_some(),
                 endpoint,
                 url,
                 tenant,
-                args,
+                args: match args_opt {
+                    Some(opt) => opt.2,
+                    None => BTreeMap::new(),
+                },
                 comment: match comment_opt {
                     Some(opt) => Some(opt.2),
                     None => None,

src/query/ast/src/parser/statement.rs

@@ -44,7 +44,6 @@ enable-histogram-metrics = [
 # Workspace dependencies
 common-arrow = { path = "../../common/arrow" }
 common-ast = { path = "../ast" }
-common-auth = { path = "../../common/auth" }
 common-base = { path = "../../common/base" }
 common-cache = { path = "../../common/cache" }
 common-catalog = { path = "../catalog" }
@@ -73,6 +72,7 @@ common-pipeline-sources = { path = "../pipeline/sources" }
 common-pipeline-transforms = { path = "../pipeline/transforms" }
 common-profile = { path = "../../common/profile" }
 common-settings = { path = "../settings" }
+common-sharing = { path = "../sharing" }
 common-sql = { path = "../sql" }
 common-storage = { path = "../../common/storage" }
 common-storages-factory = { path = "../storages/factory" }
@@ -111,7 +111,6 @@ base64 = "0.21.0"
 bincode = "1.3.3"
 bumpalo = { workspace = true }
 byteorder = "1.4.3"
-bytes = "1"
 chrono = { workspace = true }
 chrono-tz = { workspace = true }
 ctor = "0.1.26"

src/query/service/Cargo.toml

@@ -73,6 +73,7 @@ use crate::storages::Table;
 #[derive(Clone)]
 pub struct MutableCatalog {
     ctx: CatalogContext,
+    tenant: String,
 }
 
 impl MutableCatalog {
@@ -123,13 +124,17 @@ impl MutableCatalog {
             storage_factory: Arc::new(storage_factory),
             database_factory: Arc::new(database_factory),
         };
-        Ok(MutableCatalog { ctx })
+        Ok(MutableCatalog {
+            ctx,
+            tenant: conf.query.tenant_id.clone(),
+        })
     }
 
     fn build_db_instance(&self, db_info: &Arc<DatabaseInfo>) -> Result<Arc<dyn Database>> {
         let ctx = DatabaseContext {
             meta: self.ctx.meta.clone(),
             storage_factory: self.ctx.storage_factory.clone(),
+            tenant: self.tenant.clone(),
         };
         self.ctx.database_factory.get_database(ctx, db_info)
     }

src/query/service/src/catalogs/default/mutable_catalog.rs

@@ -23,4 +23,5 @@ use crate::storages::StorageFactory;
 pub struct DatabaseContext {
     pub meta: MetaStore,
     pub storage_factory: Arc<StorageFactory>,
+    pub tenant: String,
 }

src/query/service/src/databases/database_context.rs

@@ -14,16 +14,11 @@
 
 use std::str;
 use std::sync::Arc;
-use std::sync::Mutex;
 
-use bytes::Bytes;
-use common_auth::RefreshableToken;
 use common_catalog::table::Table;
-use common_config::GlobalConfig;
 use common_exception::ErrorCode;
 use common_exception::Result;
 use common_meta_api::SchemaApi;
-use common_meta_api::ShareApi;
 use common_meta_app::schema::CreateTableReq;
 use common_meta_app::schema::DatabaseInfo;
 use common_meta_app::schema::DropTableByIdReq;
@@ -43,89 +38,23 @@ use common_meta_app::schema::UpsertTableCopiedFileReply;
 use common_meta_app::schema::UpsertTableCopiedFileReq;
 use common_meta_app::schema::UpsertTableOptionReply;
 use common_meta_app::schema::UpsertTableOptionReq;
-use common_meta_app::share::GetShareEndpointReq;
-use common_meta_app::share::TableInfoMap;
-use common_users::UserApiProvider;
-use http::header::AUTHORIZATION;
-use http::header::CONTENT_LENGTH;
-use http::Method;
-use http::Request;
-use opendal::raw::AsyncBody;
-use opendal::raw::HttpClient;
-use tracing::error;
-use tracing::info;
+use common_sharing::ShareEndpointManager;
 
 use crate::databases::Database;
 use crate::databases::DatabaseContext;
 
-const TENANT_HEADER: &str = "X-DATABEND-TENANT";
-
-#[derive(Clone, Debug)]
-struct EndpointConfig {
-    pub url: String,
-
-    pub token: RefreshableToken,
-}
-
 // Share Database implementation for `Database` trait.
 #[derive(Clone)]
 pub struct ShareDatabase {
     ctx: DatabaseContext,
 
     db_info: DatabaseInfo,
-
-    client: HttpClient,
-
-    endpoint_config: Arc<Mutex<Option<EndpointConfig>>>,
 }
 
 impl ShareDatabase {
     pub const NAME: &'static str = "SHARE";
     pub fn try_create(ctx: DatabaseContext, db_info: DatabaseInfo) -> Result<Box<dyn Database>> {
-        Ok(Box::new(Self {
-            ctx,
-            db_info,
-            client: HttpClient::new()?,
-            endpoint_config: Arc::new(Mutex::new(None)),
-        }))
-    }
-
-    async fn get_share_endpoint(&self) -> Result<EndpointConfig> {
-        let endpoint_config = {
-            let endpoint_config = self.endpoint_config.lock().unwrap();
-            endpoint_config.clone()
-        };
-
-        match endpoint_config {
-            Some(ref endpoint_config) => Ok(endpoint_config.clone()),
-            None => {
-                let share_name = self.db_info.meta.from_share.clone().unwrap();
-                let req = GetShareEndpointReq {
-                    tenant: self.db_info.name_ident.tenant.clone(),
-                    endpoint: None,
-                    to_tenant: Some(share_name.tenant.clone()),
-                };
-                let meta_api = UserApiProvider::instance().get_meta_store_client();
-                let resp = meta_api.get_share_endpoint(req).await?;
-                if let Some((_, endpoint_meta)) = resp.share_endpoint_meta_vec.into_iter().next() {
-                    let endpoint = format!(
-                        "{}tenant/{}/{}/meta",
-                        endpoint_meta.url, share_name.tenant, share_name.share_name,
-                    );
-                    let config = EndpointConfig {
-                        url: endpoint,
-                        token: RefreshableToken::Direct(self.db_info.name_ident.tenant.clone()),
-                    };
-                    let mut endpoint_config = self.endpoint_config.lock().unwrap();
-                    *endpoint_config = Some(config.clone());
-                    return Ok(config);
-                }
-
-                Err(ErrorCode::EmptyShareEndpointConfig(
-                    "EmptyShareEndpointConfig, cannot query share databases".to_string(),
-                ))
-            }
-        }
+        Ok(Box::new(Self { ctx, db_info }))
     }
 
     fn load_tables(&self, table_infos: Vec<Arc<TableInfo>>) -> Result<Vec<Arc<dyn Table>>> {
@@ -136,53 +65,11 @@ impl ShareDatabase {
         })
     }
 
-    // Read table info map from operator
-    async fn get_table_info_map(&self, req: Vec<String>) -> Result<TableInfoMap> {
-        // only when endpoint_config is Some can try again
-        let mut try_again = {
-            let endpoint_config = self.endpoint_config.lock().unwrap();
-            !endpoint_config.is_some()
-        };
-
-        loop {
-            let endpoint_config = self.get_share_endpoint().await?;
-            let bs = Bytes::from(serde_json::to_vec(&req)?);
-            let auth = endpoint_config.token.to_header().await?;
-            let requester = GlobalConfig::instance().as_ref().query.tenant_id.clone();
-            let req = Request::builder()
-                .method(Method::POST)
-                .uri(&endpoint_config.url)
-                .header(AUTHORIZATION, auth)
-                .header(CONTENT_LENGTH, bs.len())
-                .header(TENANT_HEADER, requester)
-                .body(AsyncBody::Bytes(bs))?;
-            let resp = self.client.send_async(req).await;
-            match resp {
-                Ok(resp) => {
-                    let bs = resp.into_body().bytes().await?;
-                    let table_info_map: TableInfoMap = serde_json::from_slice(&bs)?;
-
-                    return Ok(table_info_map);
-                }
-                Err(err) => {
-                    if try_again {
-                        error!("get_table_info_map error: {:?}", err);
-                        return Err(err.into());
-                    } else {
-                        // endpoint may be changed, so cleanup endpoint and try again
-                        try_again = true;
-                        let mut endpoint_config = self.endpoint_config.lock().unwrap();
-                        *endpoint_config = None;
-                        info!("get_table_info_map error: {:?}, try again", err);
-                    }
-                }
-            }
-        }
-    }
-
     async fn get_table_info(&self, table_name: &str) -> Result<Arc<TableInfo>> {
-        let table_info_map = self
-            .get_table_info_map(vec![table_name.to_string()])
+        let table_info_map = ShareEndpointManager::instance()
+            .get_table_info_map(&self.ctx.tenant, &self.db_info, vec![
+                table_name.to_string(),
+            ])
             .await?;
         match table_info_map.get(table_name) {
             None => Err(ErrorCode::UnknownTable(format!(
@@ -194,7 +81,9 @@ impl ShareDatabase {
     }
 
     async fn list_tables(&self) -> Result<Vec<Arc<TableInfo>>> {
-        let table_info_map = self.get_table_info_map(vec![]).await?;
+        let table_info_map = ShareEndpointManager::instance()
+            .get_table_info_map(&self.ctx.tenant, &self.db_info, vec![])
+            .await?;
         let table_infos: Vec<Arc<TableInfo>> = table_info_map
             .values()
             .map(|table_info| Arc::new(table_info.to_owned()))

src/query/service/src/databases/share/share_database.rs

@@ -19,6 +19,7 @@ use common_catalog::catalog::CatalogManager;
 use common_config::GlobalConfig;
 use common_config::InnerConfig;
 use common_exception::Result;
+use common_sharing::ShareEndpointManager;
 use common_storage::DataOperator;
 use common_storage::ShareTableConfig;
 use common_tracing::QueryLogger;
@@ -76,6 +77,7 @@ impl GlobalServices {
         )
         .await?;
         RoleCacheManager::init()?;
+        ShareEndpointManager::init()?;
 
         Ok(())
     }