refactor(query): refactor license manager (#16489)

src/common/license/src/license_manager.rs

@@ -12,6 +12,7 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
+use std::ops::Deref;
 use std::sync::Arc;
 
 use databend_common_base::base::GlobalInstance;
@@ -58,21 +59,34 @@ pub trait LicenseManager: Sync + Send {
     fn get_storage_quota(&self, license_key: String) -> Result<StorageQuota>;
 }
 
-pub struct LicenseManagerWrapper {
-    pub manager: Box<dyn LicenseManager>,
+pub struct LicenseManagerSwitch {
+    manager: Box<dyn LicenseManager>,
+}
+
+impl LicenseManagerSwitch {
+    pub fn create(manager: Box<dyn LicenseManager>) -> LicenseManagerSwitch {
+        LicenseManagerSwitch { manager }
+    }
+
+    pub fn instance() -> Arc<LicenseManagerSwitch> {
+        GlobalInstance::get()
+    }
+}
+
+impl Deref for LicenseManagerSwitch {
+    type Target = dyn LicenseManager;
+
+    fn deref(&self) -> &Self::Target {
+        self.manager.as_ref()
+    }
 }
-unsafe impl Send for LicenseManagerWrapper {}
-unsafe impl Sync for LicenseManagerWrapper {}
 
 pub struct OssLicenseManager {}
 
 impl LicenseManager for OssLicenseManager {
     fn init(_tenant: String) -> Result<()> {
         let rm = OssLicenseManager {};
-        let wrapper = LicenseManagerWrapper {
-            manager: Box::new(rm),
-        };
-        GlobalInstance::set(Arc::new(wrapper));
+        GlobalInstance::set(Arc::new(LicenseManagerSwitch::create(Box::new(rm))));
         Ok(())
     }
 
@@ -97,7 +111,3 @@ impl LicenseManager for OssLicenseManager {
         Ok(StorageQuota::default())
     }
 }
-
-pub fn get_license_manager() -> Arc<LicenseManagerWrapper> {
-    GlobalInstance::get()
-}

src/query/ee/src/background_service/background_service_handler.rs

@@ -23,7 +23,7 @@ use databend_common_config::InnerConfig;
 use databend_common_exception::ErrorCode;
 use databend_common_exception::Result;
 use databend_common_license::license::Feature;
-use databend_common_license::license_manager::get_license_manager;
+use databend_common_license::license_manager::LicenseManagerSwitch;
 use databend_common_meta_api::BackgroundApi;
 use databend_common_meta_app::background::BackgroundJobIdent;
 use databend_common_meta_app::background::BackgroundJobInfo;
@@ -275,7 +275,7 @@ impl RealBackgroundService {
         let settings = session.get_settings();
 
         // check for valid license
-        get_license_manager().manager.check_enterprise_enabled(
+        LicenseManagerSwitch::instance().check_enterprise_enabled(
             unsafe { settings.get_enterprise_license().unwrap_or_default() },
             Feature::BackgroundService,
         )

src/query/ee/src/license/license_mgr.rs

@@ -23,7 +23,7 @@ use databend_common_license::license::Feature;
 use databend_common_license::license::LicenseInfo;
 use databend_common_license::license::StorageQuota;
 use databend_common_license::license_manager::LicenseManager;
-use databend_common_license::license_manager::LicenseManagerWrapper;
+use databend_common_license::license_manager::LicenseManagerSwitch;
 use jwt_simple::algorithms::ES256PublicKey;
 use jwt_simple::claims::JWTClaims;
 use jwt_simple::prelude::Clock;
@@ -49,10 +49,8 @@ impl LicenseManager for RealLicenseManager {
             cache: DashMap::new(),
             public_key: LICENSE_PUBLIC_KEY.to_string(),
         };
-        let wrapper = LicenseManagerWrapper {
-            manager: Box::new(rm),
-        };
-        GlobalInstance::set(Arc::new(wrapper));
+
+        GlobalInstance::set(Arc::new(LicenseManagerSwitch::create(Box::new(rm))));
         Ok(())
     }
 

src/query/ee/src/storage_encryption/handler.rs

@@ -18,7 +18,7 @@ use databend_common_base::base::GlobalInstance;
 use databend_common_config::InnerConfig;
 use databend_common_exception::Result;
 use databend_common_license::license::Feature;
-use databend_common_license::license_manager::get_license_manager;
+use databend_common_license::license_manager::LicenseManagerSwitch;
 use databend_enterprise_storage_encryption::StorageEncryptionHandler;
 use databend_enterprise_storage_encryption::StorageEncryptionHandlerWrapper;
 use databend_query::sessions::SessionManager;
@@ -40,7 +40,7 @@ impl StorageEncryptionHandler for RealStorageEncryptionHandler {
         let settings = session.get_settings();
 
         // check for valid license
-        get_license_manager().manager.check_enterprise_enabled(
+        LicenseManagerSwitch::instance().check_enterprise_enabled(
             unsafe { settings.get_enterprise_license().unwrap_or_default() },
             Feature::StorageEncryption,
         )

src/query/ee/src/storage_quota/handler.rs

@@ -18,7 +18,7 @@ use databend_common_base::base::GlobalInstance;
 use databend_common_config::InnerConfig;
 use databend_common_exception::Result;
 use databend_common_license::license::StorageQuota;
-use databend_common_license::license_manager::get_license_manager;
+use databend_common_license::license_manager::LicenseManagerSwitch;
 use databend_enterprise_storage_quota::StorageQuotaHandler;
 use databend_enterprise_storage_quota::StorageQuotaHandlerWrapper;
 use databend_query::sessions::SessionManager;
@@ -39,8 +39,7 @@ impl StorageQuotaHandler for RealStorageQuotaHandler {
 
         let settings = session.get_settings();
         // check for valid license
-        get_license_manager()
-            .manager
+        LicenseManagerSwitch::instance()
             .get_storage_quota(unsafe { settings.get_enterprise_license().unwrap_or_default() })
     }
 }

src/query/ee/src/test_kits/mock_services.rs

@@ -17,7 +17,7 @@ use std::sync::Arc;
 use databend_common_base::base::GlobalInstance;
 use databend_common_config::InnerConfig;
 use databend_common_exception::Result;
-use databend_common_license::license_manager::LicenseManagerWrapper;
+use databend_common_license::license_manager::LicenseManagerSwitch;
 
 use crate::aggregating_index::RealAggregatingIndexHandler;
 use crate::data_mask::RealDatamaskHandler;
@@ -32,10 +32,7 @@ impl MockServices {
     #[async_backtrace::framed]
     pub async fn init(cfg: &InnerConfig, public_key: String) -> Result<()> {
         let rm = RealLicenseManager::new(cfg.query.tenant_id.tenant_name().to_string(), public_key);
-        let wrapper = LicenseManagerWrapper {
-            manager: Box::new(rm),
-        };
-        GlobalInstance::set(Arc::new(wrapper));
+        GlobalInstance::set(Arc::new(LicenseManagerSwitch::create(Box::new(rm))));
         RealVacuumHandler::init()?;
         RealAggregatingIndexHandler::init()?;
         RealDatamaskHandler::init()?;

src/query/service/src/interpreters/common/stream.rs

@@ -19,7 +19,7 @@ use chrono::Utc;
 use databend_common_exception::ErrorCode;
 use databend_common_exception::Result;
 use databend_common_license::license::Feature;
-use databend_common_license::license_manager::get_license_manager;
+use databend_common_license::license_manager::LicenseManagerSwitch;
 use databend_common_meta_app::schema::TableInfo;
 use databend_common_meta_app::schema::UpdateStreamMetaReq;
 use databend_common_meta_app::schema::UpdateTableMetaReq;
@@ -49,9 +49,7 @@ pub async fn dml_build_update_stream_req(
         return Ok(vec![]);
     }
 
-    let license_manager = get_license_manager();
-    license_manager
-        .manager
+    LicenseManagerSwitch::instance()
         .check_enterprise_enabled(ctx.get_license_key(), Feature::Stream)?;
 
     let mut reqs = Vec::with_capacity(tables.len());
@@ -134,9 +132,7 @@ pub async fn query_build_update_stream_req(
         return Ok(None);
     }
 
-    let license_manager = get_license_manager();
-    license_manager
-        .manager
+    LicenseManagerSwitch::instance()
         .check_enterprise_enabled(ctx.get_license_key(), Feature::Stream)?;
 
     let cap = streams.len();

src/query/service/src/interpreters/hook/vacuum_hook.rs

@@ -19,7 +19,7 @@ use databend_common_base::runtime::GlobalIORuntime;
 use databend_common_catalog::table_context::TableContext;
 use databend_common_exception::Result;
 use databend_common_license::license::Feature::Vacuum;
-use databend_common_license::license_manager::get_license_manager;
+use databend_common_license::license_manager::LicenseManagerSwitch;
 use databend_common_pipeline_core::query_spill_prefix;
 use databend_common_storage::DataOperator;
 use databend_enterprise_vacuum_handler::get_vacuum_handler;
@@ -31,13 +31,11 @@ pub fn hook_vacuum_temp_files(query_ctx: &Arc<QueryContext>) -> Result<()> {
     let tenant = query_ctx.get_tenant();
     let settings = query_ctx.get_settings();
     let spill_prefix = query_spill_prefix(tenant.tenant_name(), &query_ctx.get_id());
-    let license_manager = get_license_manager();
     let vacuum_limit = settings.get_max_vacuum_temp_files_after_query()?;
 
     // disable all s3 operator if vacuum limit = 0
     if vacuum_limit != 0
-        && license_manager
-            .manager
+        && LicenseManagerSwitch::instance()
             .check_enterprise_enabled(query_ctx.get_license_key(), Vacuum)
             .is_ok()
     {

src/query/service/src/interpreters/interpreter_data_mask_create.rs

@@ -16,7 +16,7 @@ use std::sync::Arc;
 
 use databend_common_exception::Result;
 use databend_common_license::license::Feature;
-use databend_common_license::license_manager::get_license_manager;
+use databend_common_license::license_manager::LicenseManagerSwitch;
 use databend_common_sql::plans::CreateDatamaskPolicyPlan;
 use databend_common_users::UserApiProvider;
 use databend_enterprise_data_mask_feature::get_datamask_handler;
@@ -49,9 +49,7 @@ impl Interpreter for CreateDataMaskInterpreter {
 
     #[async_backtrace::framed]
     async fn execute2(&self) -> Result<PipelineBuildResult> {
-        let license_manager = get_license_manager();
-        license_manager
-            .manager
+        LicenseManagerSwitch::instance()
             .check_enterprise_enabled(self.ctx.get_license_key(), Feature::DataMask)?;
         let meta_api = UserApiProvider::instance().get_meta_store_client();
         let handler = get_datamask_handler();

src/query/service/src/interpreters/interpreter_data_mask_desc.rs

@@ -20,7 +20,7 @@ use databend_common_expression::types::StringType;
 use databend_common_expression::DataBlock;
 use databend_common_expression::FromData;
 use databend_common_license::license::Feature;
-use databend_common_license::license_manager::get_license_manager;
+use databend_common_license::license_manager::LicenseManagerSwitch;
 use databend_common_sql::plans::DescDatamaskPolicyPlan;
 use databend_common_users::UserApiProvider;
 use databend_enterprise_data_mask_feature::get_datamask_handler;
@@ -54,9 +54,7 @@ impl Interpreter for DescDataMaskInterpreter {
 
     #[async_backtrace::framed]
     async fn execute2(&self) -> Result<PipelineBuildResult> {
-        let license_manager = get_license_manager();
-        license_manager
-            .manager
+        LicenseManagerSwitch::instance()
             .check_enterprise_enabled(self.ctx.get_license_key(), Feature::DataMask)?;
         let meta_api = UserApiProvider::instance().get_meta_store_client();
         let handler = get_datamask_handler();

src/query/service/src/interpreters/interpreter_data_mask_drop.rs

@@ -16,7 +16,7 @@ use std::sync::Arc;
 
 use databend_common_exception::Result;
 use databend_common_license::license::Feature;
-use databend_common_license::license_manager::get_license_manager;
+use databend_common_license::license_manager::LicenseManagerSwitch;
 use databend_common_sql::plans::DropDatamaskPolicyPlan;
 use databend_common_users::UserApiProvider;
 use databend_enterprise_data_mask_feature::get_datamask_handler;
@@ -49,9 +49,7 @@ impl Interpreter for DropDataMaskInterpreter {
 
     #[async_backtrace::framed]
     async fn execute2(&self) -> Result<PipelineBuildResult> {
-        let license_manager = get_license_manager();
-        license_manager
-            .manager
+        LicenseManagerSwitch::instance()
             .check_enterprise_enabled(self.ctx.get_license_key(), Feature::DataMask)?;
         let meta_api = UserApiProvider::instance().get_meta_store_client();
         let handler = get_datamask_handler();

src/query/service/src/interpreters/interpreter_index_create.rs

@@ -18,7 +18,7 @@ use chrono::Utc;
 use databend_common_exception::ErrorCode;
 use databend_common_exception::Result;
 use databend_common_license::license::Feature;
-use databend_common_license::license_manager::get_license_manager;
+use databend_common_license::license_manager::LicenseManagerSwitch;
 use databend_common_meta_app::schema::CreateIndexReq;
 use databend_common_meta_app::schema::IndexMeta;
 use databend_common_meta_app::schema::IndexNameIdent;
@@ -56,9 +56,7 @@ impl Interpreter for CreateIndexInterpreter {
     async fn execute2(&self) -> Result<PipelineBuildResult> {
         let tenant = self.ctx.get_tenant();
 
-        let license_manager = get_license_manager();
-        license_manager
-            .manager
+        LicenseManagerSwitch::instance()
             .check_enterprise_enabled(self.ctx.get_license_key(), Feature::AggregateIndex)?;
 
         let index_name = self.plan.index_name.clone();

src/query/service/src/interpreters/interpreter_index_drop.rs

@@ -16,7 +16,7 @@ use std::sync::Arc;
 
 use databend_common_exception::Result;
 use databend_common_license::license::Feature;
-use databend_common_license::license_manager::get_license_manager;
+use databend_common_license::license_manager::LicenseManagerSwitch;
 use databend_common_meta_app::schema::DropIndexReq;
 use databend_common_meta_app::schema::IndexNameIdent;
 use databend_common_sql::plans::DropIndexPlan;
@@ -52,9 +52,7 @@ impl Interpreter for DropIndexInterpreter {
     async fn execute2(&self) -> Result<PipelineBuildResult> {
         let tenant = self.ctx.get_tenant();
 
-        let license_manager = get_license_manager();
-        license_manager
-            .manager
+        LicenseManagerSwitch::instance()
             .check_enterprise_enabled(self.ctx.get_license_key(), Feature::AggregateIndex)?;
 
         let index_name = self.plan.index.clone();

src/query/service/src/interpreters/interpreter_index_refresh.rs

@@ -27,7 +27,7 @@ use databend_common_expression::TableField;
 use databend_common_expression::TableSchema;
 use databend_common_expression::BLOCK_NAME_COL_NAME;
 use databend_common_license::license::Feature;
-use databend_common_license::license_manager::get_license_manager;
+use databend_common_license::license_manager::LicenseManagerSwitch;
 use databend_common_meta_app::schema::IndexMeta;
 use databend_common_meta_app::schema::UpdateIndexReq;
 use databend_common_pipeline_core::ExecutionInfo;
@@ -214,9 +214,7 @@ impl Interpreter for RefreshIndexInterpreter {
 
     #[async_backtrace::framed]
     async fn execute2(&self) -> Result<PipelineBuildResult> {
-        let license_manager = get_license_manager();
-        license_manager
-            .manager
+        LicenseManagerSwitch::instance()
             .check_enterprise_enabled(self.ctx.get_license_key(), Feature::AggregateIndex)?;
         let (mut query_plan, output_schema, select_columns) = match self.plan.query_plan.as_ref() {
             Plan::Query {

src/query/service/src/interpreters/interpreter_stream_create.rs

@@ -16,7 +16,7 @@ use std::sync::Arc;
 
 use databend_common_exception::Result;
 use databend_common_license::license::Feature;
-use databend_common_license::license_manager::get_license_manager;
+use databend_common_license::license_manager::LicenseManagerSwitch;
 use databend_common_sql::plans::CreateStreamPlan;
 use databend_common_storages_fuse::TableContext;
 use databend_enterprise_stream_handler::get_stream_handler;
@@ -48,9 +48,7 @@ impl Interpreter for CreateStreamInterpreter {
 
     #[async_backtrace::framed]
     async fn execute2(&self) -> Result<PipelineBuildResult> {
-        let license_manager = get_license_manager();
-        license_manager
-            .manager
+        LicenseManagerSwitch::instance()
             .check_enterprise_enabled(self.ctx.get_license_key(), Feature::Stream)?;
 
         let handler = get_stream_handler();

src/query/service/src/interpreters/interpreter_stream_drop.rs

@@ -16,7 +16,7 @@ use std::sync::Arc;
 
 use databend_common_exception::Result;
 use databend_common_license::license::Feature;
-use databend_common_license::license_manager::get_license_manager;
+use databend_common_license::license_manager::LicenseManagerSwitch;
 use databend_common_sql::plans::DropStreamPlan;
 use databend_enterprise_stream_handler::get_stream_handler;
 
@@ -48,9 +48,7 @@ impl Interpreter for DropStreamInterpreter {
 
     #[async_backtrace::framed]
     async fn execute2(&self) -> Result<PipelineBuildResult> {
-        let license_manager = get_license_manager();
-        license_manager
-            .manager
+        LicenseManagerSwitch::instance()
             .check_enterprise_enabled(self.ctx.get_license_key(), Feature::Stream)?;
 
         let handler = get_stream_handler();

src/query/service/src/interpreters/interpreter_table_add_column.rs

@@ -19,7 +19,7 @@ use databend_common_catalog::table::TableExt;
 use databend_common_exception::ErrorCode;
 use databend_common_exception::Result;
 use databend_common_license::license::Feature::ComputedColumn;
-use databend_common_license::license_manager::get_license_manager;
+use databend_common_license::license_manager::LicenseManagerSwitch;
 use databend_common_meta_app::schema::DatabaseType;
 use databend_common_meta_app::schema::TableMeta;
 use databend_common_meta_app::schema::UpdateTableMetaReq;
@@ -94,9 +94,7 @@ impl Interpreter for AddTableColumnInterpreter {
         let mut new_table_meta = table_info.meta.clone();
         let field = self.plan.field.clone();
         if field.computed_expr().is_some() {
-            let license_manager = get_license_manager();
-            license_manager
-                .manager
+            LicenseManagerSwitch::instance()
                 .check_enterprise_enabled(self.ctx.get_license_key(), ComputedColumn)?;
         }