fix: cr

Author: fengys1996

src/cmd/src/bin/greptime.rs

@@ -17,10 +17,14 @@
 use clap::{Parser, Subcommand};
 use cmd::datanode::builder::InstanceBuilder;
 use cmd::error::{InitTlsProviderSnafu, Result};
+use cmd::extension::flownode::DefaultExtensionFactory as FlowExtensionFactory;
+use cmd::extension::frontend::DefaultExtensionFactory as FrontendExtensionFactory;
+use cmd::extension::standalone::DefaultExtensionFactory as StandaloneExtensionFactory;
 use cmd::options::{EmptyOptions, GlobalOptions};
 use cmd::{App, cli, datanode, flownode, frontend, metasrv, standalone};
 use common_base::Plugins;
 use common_version::{verbose_version, version};
+use meta_srv::bootstrap::extension::DefaultExtensionFactory as MetaExtensionFactory;
 use servers::install_ring_crypto_provider;
 
 #[derive(Parser)]
@@ -113,40 +117,23 @@ async fn start(cli: Command) -> Result<()> {
             datanode::SubCommand::Objbench(ref bench) => bench.run().await,
         },
         SubCommand::Flownode(cmd) => {
-            cmd.build(
-                cmd.load_options::<EmptyOptions>(&cli.global_options)?,
-                Default::default(),
-            )
-            .await?
-            .run()
-            .await
+            let opts = cmd.load_options::<EmptyOptions>(&cli.global_options)?;
+            cmd.build(opts, FlowExtensionFactory).await?.run().await
         }
         SubCommand::Frontend(cmd) => {
-            cmd.build(
-                cmd.load_options::<EmptyOptions>(&cli.global_options)?,
-                Default::default(),
-            )
-            .await?
-            .run()
-            .await
+            let opts = cmd.load_options::<EmptyOptions>(&cli.global_options)?;
+            cmd.build(opts, FrontendExtensionFactory).await?.run().await
         }
         SubCommand::Metasrv(cmd) => {
-            cmd.build(
-                cmd.load_options::<EmptyOptions>(&cli.global_options)?,
-                Default::default(),
-            )
-            .await?
-            .run()
-            .await
+            let opts = cmd.load_options::<EmptyOptions>(&cli.global_options)?;
+            cmd.build(opts, MetaExtensionFactory).await?.run().await
         }
         SubCommand::Standalone(cmd) => {
-            cmd.build(
-                cmd.load_options::<EmptyOptions>(&cli.global_options)?,
-                Default::default(),
-            )
-            .await?
-            .run()
-            .await
+            let opts = cmd.load_options::<EmptyOptions>(&cli.global_options)?;
+            cmd.build(opts, StandaloneExtensionFactory)
+                .await?
+                .run()
+                .await
         }
         SubCommand::Cli(cmd) => {
             cmd.build(cmd.load_options(&cli.global_options)?)

src/cmd/src/cli.rs

@@ -108,6 +108,7 @@ mod tests {
     use common_telemetry::logging::LoggingOptions;
 
     use crate::error::Result as CmdResult;
+    use crate::extension::standalone::DefaultExtensionFactory as StandaloneExtensionFactory;
     use crate::options::{EmptyOptions, GlobalOptions};
     use crate::{App, cli, standalone};
 
@@ -126,7 +127,7 @@ mod tests {
             .load_options::<EmptyOptions>(&GlobalOptions::default())
             .unwrap();
         let mut instance = standalone
-            .build(standalone_opts, Default::default())
+            .build(standalone_opts, StandaloneExtensionFactory)
             .await?;
         instance.start().await?;
 

src/cmd/src/extension/common.rs

@@ -42,7 +42,7 @@ pub type InformationSchemaTableFactoriesRef = Arc<dyn InformationSchemaTableFact
 
 /// Context for information schema table factory providers.
 pub struct TableFactoryContext {
-    pub fe_client: Option<Arc<FrontendClient>>,
+    pub fe_client: Arc<FrontendClient>,
 }
 
 /// Allows extending the gRPC server with additional services (e.g., enterprise

src/cmd/src/extension/flownode.rs

@@ -12,9 +12,29 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
+use common_error::ext::BoxedError;
+
 use crate::extension::common::GrpcExtensionRef;
 
 #[derive(Default)]
 pub struct Extension {
     pub grpc: Option<GrpcExtensionRef>,
 }
+
+/// Factory trait to create Extension instances.
+pub trait ExtensionFactory: Send + Sync {
+    fn create(
+        &self,
+        ctx: ExtensionContext,
+    ) -> impl Future<Output = Result<Extension, BoxedError>> + Send;
+}
+
+pub struct ExtensionContext {}
+
+pub struct DefaultExtensionFactory;
+
+impl ExtensionFactory for DefaultExtensionFactory {
+    async fn create(&self, _ctx: ExtensionContext) -> Result<Extension, BoxedError> {
+        Ok(Extension::default())
+    }
+}

src/cmd/src/extension/frontend.rs

@@ -12,14 +12,39 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-#[cfg(feature = "enterprise")]
-use operator::statement::TriggerQuerierFactoryRef;
+use std::collections::HashMap;
 
-use crate::extension::common::InformationSchemaTableFactoriesRef;
+use catalog::information_schema::InformationSchemaTableFactoryRef;
+use common_error::ext::BoxedError;
+use common_meta::kv_backend::KvBackendRef;
+use meta_client::MetaClientRef;
+#[cfg(feature = "enterprise")]
+use operator::statement::TriggerQuerierRef;
 
 #[derive(Default)]
 pub struct Extension {
-    pub info_schema_factories: Option<InformationSchemaTableFactoriesRef>,
+    pub info_schema_factories: Option<HashMap<String, InformationSchemaTableFactoryRef>>,
     #[cfg(feature = "enterprise")]
-    pub trigger_querier_factory: Option<TriggerQuerierFactoryRef>,
+    pub trigger_querier: Option<TriggerQuerierRef>,
+}
+
+/// Factory trait to create Extension instances.
+pub trait ExtensionFactory: Send + Sync {
+    fn create(
+        &self,
+        ctx: ExtensionContext,
+    ) -> impl Future<Output = Result<Extension, BoxedError>> + Send;
+}
+
+pub struct ExtensionContext {
+    pub kv_backend: KvBackendRef,
+    pub meta_client: MetaClientRef,
+}
+
+pub struct DefaultExtensionFactory;
+
+impl ExtensionFactory for DefaultExtensionFactory {
+    async fn create(&self, _ctx: ExtensionContext) -> Result<Extension, BoxedError> {
+        Ok(Extension::default())
+    }
 }

src/cmd/src/extension/standalone.rs

@@ -12,45 +12,57 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
+use std::collections::HashMap;
+use std::sync::Arc;
+
+use catalog::CatalogManagerRef;
+use catalog::information_schema::InformationSchemaTableFactoryRef;
+use common_error::ext::BoxedError;
 #[cfg(feature = "enterprise")]
-pub use ee::*;
+use common_meta::ddl_manager::TriggerDdlManagerRef;
+use common_meta::kv_backend::KvBackendRef;
+use flow::FrontendClient;
 #[cfg(feature = "enterprise")]
-use operator::statement::TriggerQuerierFactoryRef;
+use operator::statement::TriggerQuerierRef;
 
-use crate::extension::common::InformationSchemaTableFactoriesRef;
+use crate::extension::common::{InformationSchemaTableFactories, TableFactoryContext};
 
-#[cfg(feature = "enterprise")]
-mod ee {
-    use std::sync::Arc;
-
-    use catalog::CatalogManagerRef;
-    use common_error::ext::BoxedError;
-    use common_meta::ddl_manager::TriggerDdlManagerRef;
-    use common_meta::kv_backend::KvBackendRef;
-    use flow::FrontendClient;
-
-    #[async_trait::async_trait]
-    pub trait TriggerDdlManagerFactory: Send + Sync {
-        async fn create(
-            &self,
-            ctx: TriggerDdlManagerRequest,
-        ) -> Result<TriggerDdlManagerRef, BoxedError>;
-    }
+#[derive(Default)]
+pub struct Extension {
+    #[cfg(feature = "enterprise")]
+    pub trigger_ddl_manager: Option<TriggerDdlManagerRef>,
+    #[cfg(feature = "enterprise")]
+    pub trigger_querier: Option<TriggerQuerierRef>,
+}
 
-    pub type TriggerDdlManagerFactoryRef = Arc<dyn TriggerDdlManagerFactory>;
+/// Factory trait to create Extension instances.
+pub trait ExtensionFactory: InformationSchemaTableFactories + Send + Sync {
+    fn create(
+        &self,
+        ctx: ExtensionContext,
+    ) -> impl Future<Output = Result<Extension, BoxedError>> + Send;
+}
 
-    pub struct TriggerDdlManagerRequest {
-        pub kv_backend: KvBackendRef,
-        pub catalog_manager: CatalogManagerRef,
-        pub fe_client: Arc<FrontendClient>,
+pub struct ExtensionContext {
+    pub kv_backend: KvBackendRef,
+    pub catalog_manager: CatalogManagerRef,
+    pub frontend_client: Arc<FrontendClient>,
+}
+
+pub struct DefaultExtensionFactory;
+
+#[async_trait::async_trait]
+impl InformationSchemaTableFactories for DefaultExtensionFactory {
+    async fn create_factories(
+        &self,
+        _ctx: TableFactoryContext,
+    ) -> Result<HashMap<String, InformationSchemaTableFactoryRef>, BoxedError> {
+        Ok(HashMap::new())
     }
 }
 
-#[derive(Default)]
-pub struct Extension {
-    pub info_schema_factories: Option<InformationSchemaTableFactoriesRef>,
-    #[cfg(feature = "enterprise")]
-    pub trigger_ddl_manager_factory: Option<TriggerDdlManagerFactoryRef>,
-    #[cfg(feature = "enterprise")]
-    pub trigger_querier_factory: Option<TriggerQuerierFactoryRef>,
+impl ExtensionFactory for DefaultExtensionFactory {
+    async fn create(&self, _ctx: ExtensionContext) -> Result<Extension, BoxedError> {
+        Ok(Extension::default())
+    }
 }

src/cmd/src/flownode.rs

@@ -48,7 +48,7 @@ use crate::error::{
     MissingConfigSnafu, OtherSnafu, Result, ShutdownFlownodeSnafu, StartFlownodeSnafu,
 };
 use crate::extension::common::GrpcExtensionContext;
-use crate::extension::flownode::Extension;
+use crate::extension::flownode::{ExtensionContext, ExtensionFactory};
 use crate::options::{GlobalOptions, GreptimeOptions};
 use crate::{App, create_resource_limit_metrics, log_versions, maybe_activate_heap_profile};
 
@@ -109,12 +109,12 @@ pub struct Command {
 }
 
 impl Command {
-    pub async fn build<E: Debug>(
+    pub async fn build<E: Debug, F: ExtensionFactory>(
         &self,
         opts: FlownodeOptions<E>,
-        extension: Extension,
+        extension_factory: F,
     ) -> Result<Instance> {
-        self.subcmd.build(opts, extension).await
+        self.subcmd.build(opts, extension_factory).await
     }
 
     pub fn load_options<E: Configurable>(
@@ -133,13 +133,13 @@ enum SubCommand {
 }
 
 impl SubCommand {
-    async fn build<E: Debug>(
+    async fn build<E: Debug, F: ExtensionFactory>(
         &self,
         opts: FlownodeOptions<E>,
-        extension: Extension,
+        extension_factory: F,
     ) -> Result<Instance> {
         match self {
-            SubCommand::Start(cmd) => cmd.build(opts, extension).await,
+            SubCommand::Start(cmd) => cmd.build(opts, extension_factory).await,
         }
     }
 }
@@ -262,10 +262,10 @@ impl StartCommand {
         Ok(())
     }
 
-    async fn build<E: Debug>(
+    async fn build<E: Debug, F: ExtensionFactory>(
         &self,
         opts: FlownodeOptions<E>,
-        extension: Extension,
+        extension_factory: F,
     ) -> Result<Instance> {
         common_runtime::init_global_runtimes(&opts.runtime);
 
@@ -405,14 +405,18 @@ impl StartCommand {
 
         let mut builder =
             FlownodeServiceBuilder::grpc_server_builder(&opts, flownode.flownode_server());
-        if let Some(extension) = extension.grpc.as_ref() {
+        let extension = extension_factory
+            .create(ExtensionContext {})
+            .await
+            .context(OtherSnafu)?;
+        if let Some(grpc_extension) = extension.grpc {
             let ctx = GrpcExtensionContext {
                 kv_backend: cached_meta_backend.clone(),
                 fe_client: frontend_client.clone(),
                 flownode_id: member_id,
                 catalog_manager: catalog_manager.clone(),
             };
-            extension
+            grpc_extension
                 .extend_grpc_services(&mut builder, ctx)
                 .await
                 .context(OtherSnafu)?