fix(query): check enterprise license embedded (#18202)

* fix(query): check enterprise license embedded

* fix: return if env changed

* z

* z

Author: everpcpc

.github/workflows/reuse.linux.yml

@@ -124,70 +124,70 @@ jobs:
         timeout-minutes: 60
 
   test_metactl:
+    needs: [build, check]
     runs-on:
       - self-hosted
       - X64
       - Linux
       - 2c8g
       - "${{ inputs.runner_provider }}"
       - "${{ inputs.runner_capacity }}"
-    needs: [build, check]
     steps:
       - uses: actions/checkout@v4
       - uses: ./.github/actions/test_metactl
         timeout-minutes: 10
 
   test_compat_meta_query:
+    needs: [build, check]
     runs-on:
       - self-hosted
       - X64
       - Linux
       - 2c8g
       - "${{ inputs.runner_provider }}"
       - "${{ inputs.runner_capacity }}"
-    needs: [build, check]
     steps:
       - uses: actions/checkout@v4
       - uses: ./.github/actions/test_compat_meta_query
         timeout-minutes: 10
 
   test_compat_fuse:
+    needs: [build, check]
     runs-on:
       - self-hosted
       - X64
       - Linux
       - 2c8g
       - "${{ inputs.runner_provider }}"
       - "${{ inputs.runner_capacity }}"
-    needs: [build, check]
     steps:
       - uses: actions/checkout@v4
       - uses: ./.github/actions/test_compat_fuse
         timeout-minutes: 20
 
   test_compat_meta_meta:
+    needs: [build, check]
     runs-on:
       - self-hosted
       - X64
       - Linux
       - 2c8g
       - "${{ inputs.runner_provider }}"
       - "${{ inputs.runner_capacity }}"
-    needs: [build, check]
     steps:
       - uses: actions/checkout@v4
       - uses: ./.github/actions/test_compat_meta_meta
         timeout-minutes: 20
 
   test_logs:
+    needs: [build, check]
     runs-on:
       - self-hosted
       - X64
       - Linux
       - 2c8g
       - "${{ inputs.runner_provider }}"
       - "${{ inputs.runner_capacity }}"
-    needs: [build, check]
     steps:
       - uses: actions/checkout@v4
       - uses: ./.github/actions/setup_license
@@ -198,42 +198,42 @@ jobs:
         timeout-minutes: 20
 
   test_meta_cluster:
+    needs: [build, check]
     runs-on:
       - self-hosted
       - X64
       - Linux
       - 2c8g
       - "${{ inputs.runner_provider }}"
       - "${{ inputs.runner_capacity }}"
-    needs: [build, check]
     steps:
       - uses: actions/checkout@v4
       - uses: ./.github/actions/test_meta_cluster
         timeout-minutes: 10
 
   test_stateless_standalone:
+    needs: [build, check]
     runs-on:
       - self-hosted
       - X64
       - Linux
       - 2c8g
       - "${{ inputs.runner_provider }}"
       - "${{ inputs.runner_capacity }}"
-    needs: [build, check]
     steps:
       - uses: actions/checkout@v4
       - uses: ./.github/actions/test_stateless_standalone_linux
         timeout-minutes: 18
 
   test_stateless_cluster:
+    needs: [build, check]
     runs-on:
       - self-hosted
       - X64
       - Linux
       - 2c8g
       - "${{ inputs.runner_provider }}"
       - "${{ inputs.runner_capacity }}"
-    needs: [build, check]
     steps:
       - uses: actions/checkout@v4
       - uses: ./.github/actions/setup_license
@@ -244,13 +244,14 @@ jobs:
         timeout-minutes: 18
 
   test_stateful_standalone:
+    needs: [build, check]
     runs-on:
       - self-hosted
       - X64
       - Linux
       - 2c8g
       - "${{ inputs.runner_provider }}"
-    needs: [build, check]
+      - "${{ inputs.runner_capacity }}"
     steps:
       - uses: actions/checkout@v4
       - uses: ./.github/actions/test_stateful_standalone_linux
@@ -262,14 +263,14 @@ jobs:
           name: test-stateful-standalone-linux
 
   test_stateful_cluster:
+    needs: [build, check]
     runs-on:
       - self-hosted
       - X64
       - Linux
       - 2c8g
       - "${{ inputs.runner_provider }}"
       - "${{ inputs.runner_capacity }}"
-    needs: [build, check]
     steps:
       - uses: actions/checkout@v4
       - uses: ./.github/actions/setup_license
@@ -286,28 +287,28 @@ jobs:
 
   test_stateful_large_data:
     if: contains(github.event.pull_request.labels.*.name, 'ci-largedata')
+    needs: [build, check]
     runs-on:
       - self-hosted
       - X64
       - Linux
       - 2c8g
       - "${{ inputs.runner_provider }}"
       - "${{ inputs.runner_capacity }}"
-    needs: [build, check]
     steps:
       - uses: actions/checkout@v4
       - uses: ./.github/actions/test_stateful_large_data
         timeout-minutes: 60
 
   test_stateful_iceberg_catalogs:
+    needs: [build, check]
     runs-on:
       - self-hosted
       - X64
       - Linux
       - 2c8g
       - "${{ inputs.runner_provider }}"
       - "${{ inputs.runner_capacity }}"
-    needs: [build, check]
     steps:
       - uses: actions/checkout@v4
       - uses: ./.github/actions/test_stateful_iceberg_catalogs_standalone

Cargo.lock

@@ -119,10 +119,8 @@ pub async fn verify_query_license(cfg: &InnerConfig) -> Result<()> {
     let session = session_manager.register_session(session)?;
     let settings = session.get_settings();
 
-    LicenseManagerSwitch::instance().check_enterprise_enabled(
-        unsafe { settings.get_enterprise_license().unwrap_or_default() },
-        Feature::SystemManagement,
-    )?;
+    LicenseManagerSwitch::instance()
+        .check_enterprise_enabled(settings.get_enterprise_license(), Feature::SystemManagement)?;
 
     debug!("databend license check passed");
     Ok(())

src/bendsave/src/storage.rs

@@ -73,6 +73,7 @@ pub fn set_env_config() {
 }
 
 pub fn add_env_version() {
+    println!("cargo:rerun-if-env-changed=DATABEND_RELEASE_VERSION");
     let version = discover_version().expect(VERSION_ERROR_MESSAGE);
     println!("cargo:rustc-env=DATABEND_GIT_SEMVER={}", version);
 }
@@ -90,11 +91,13 @@ fn discover_version() -> Result<String> {
 }
 
 pub fn add_env_license() {
+    println!("cargo:rerun-if-env-changed=DATABEND_ENTERPRISE_LICENSE_EMBEDDED");
     let v = env::var("DATABEND_ENTERPRISE_LICENSE_EMBEDDED").unwrap_or_default();
     println!("cargo:rustc-env=DATABEND_ENTERPRISE_LICENSE_EMBEDDED={v}");
 }
 
 pub fn add_license_public_key() {
+    println!("cargo:rerun-if-env-changed=DATABEND_ENTERPRISE_LICENSE_PUBLIC_KEY");
     let v = env::var("DATABEND_ENTERPRISE_LICENSE_PUBLIC_KEY").unwrap_or_default();
     let v = base64::Engine::encode(&base64::engine::general_purpose::STANDARD, v.as_bytes());
     println!("cargo:rustc-env=DATABEND_ENTERPRISE_LICENSE_PUBLIC_KEY={v}");

src/common/building/src/lib.rs

@@ -41,7 +41,7 @@ impl StorageEncryptionHandler for RealStorageEncryptionHandler {
 
         // check for valid license
         LicenseManagerSwitch::instance().check_enterprise_enabled(
-            unsafe { settings.get_enterprise_license().unwrap_or_default() },
+            settings.get_enterprise_license(),
             Feature::StorageEncryption,
         )
     }

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

@@ -39,8 +39,7 @@ impl StorageQuotaHandler for RealStorageQuotaHandler {
 
         let settings = session.get_settings();
         // check for valid license
-        LicenseManagerSwitch::instance()
-            .get_storage_quota(unsafe { settings.get_enterprise_license().unwrap_or_default() })
+        LicenseManagerSwitch::instance().get_storage_quota(settings.get_enterprise_license())
     }
 }
 

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

@@ -121,7 +121,6 @@ use databend_common_storages_stream::stream_table::StreamTable;
 use databend_common_users::GrantObjectVisibilityChecker;
 use databend_common_users::UserApiProvider;
 use databend_common_version::DATABEND_COMMIT_VERSION;
-use databend_common_version::DATABEND_ENTERPRISE_LICENSE_EMBEDDED;
 use databend_storages_common_session::drop_table_by_id;
 use databend_storages_common_session::SessionState;
 use databend_storages_common_session::TxnManagerRef;
@@ -1177,11 +1176,7 @@ impl TableContext for QueryContext {
             && ThreadTracker::capture_log_settings().is_none()
         {
             LicenseManagerSwitch::instance().check_enterprise_enabled(
-                unsafe {
-                    self.get_settings()
-                        .get_enterprise_license()
-                        .unwrap_or_default()
-                },
+                self.get_settings().get_enterprise_license(),
                 Feature::SystemHistory,
             )?;
         }
@@ -1411,10 +1406,7 @@ impl TableContext for QueryContext {
     }
 
     fn get_license_key(&self) -> String {
-        unsafe { self.get_settings().get_enterprise_license() }.unwrap_or_else(|_| {
-            // Try load license from embedded env if failed to load from settings.
-            DATABEND_ENTERPRISE_LICENSE_EMBEDDED.to_string()
-        })
+        self.get_settings().get_enterprise_license()
     }
 
     fn get_query_profiles(&self) -> Vec<PlanProfile> {

src/query/service/src/sessions/query_ctx.rs

@@ -194,17 +194,7 @@ impl AsyncSource for LicenseInfoSource {
         let settings = self.ctx.get_settings();
         // sync global changes on distributed node cluster.
         settings.load_changes().await?;
-        let license = unsafe {
-            settings.get_enterprise_license().map_err_to_code(
-                ErrorCode::LicenseKeyInvalid,
-                || {
-                    format!(
-                        "failed to get license for {}",
-                        self.ctx.get_tenant().display()
-                    )
-                },
-            )?
-        };
+        let license = settings.get_enterprise_license();
 
         LicenseManagerSwitch::instance()
             .check_enterprise_enabled(license.clone(), Feature::LicenseInfo)?;

src/query/service/src/table_functions/others/license_info.rs

@@ -16,6 +16,7 @@ databend-common-exception = { workspace = true }
 databend-common-io = { workspace = true }
 databend-common-meta-app = { workspace = true }
 databend-common-users = { workspace = true }
+databend-common-version = { workspace = true }
 itertools = { workspace = true }
 log = { workspace = true }
 num_cpus = { workspace = true }

src/query/settings/Cargo.toml

@@ -19,6 +19,7 @@ use databend_common_exception::ErrorCode;
 use databend_common_exception::Result;
 use databend_common_io::GeometryDataType;
 use databend_common_meta_app::principal::UserSettingValue;
+use databend_common_version::DATABEND_ENTERPRISE_LICENSE_EMBEDDED;
 
 use crate::settings::Settings;
 use crate::settings_default::DefaultSettings;
@@ -559,9 +560,17 @@ impl Settings {
         self.try_get_u64("acquire_lock_timeout")
     }
 
-    /// # Safety
-    pub unsafe fn get_enterprise_license(&self) -> Result<String> {
-        self.unchecked_try_get_string("enterprise_license")
+    pub fn get_enterprise_license(&self) -> String {
+        let license = unsafe {
+            self.unchecked_try_get_string("enterprise_license")
+                .unwrap_or_default()
+        };
+        if license.is_empty() {
+            // Try load license from embedded env if failed to load from settings.
+            DATABEND_ENTERPRISE_LICENSE_EMBEDDED.to_string()
+        } else {
+            license
+        }
     }
 
     /// # Safety