From 4d14d75e6241ad2fdd6b7841fe114ce383b1d554 Mon Sep 17 00:00:00 2001
From: Xuanwo <github@xuanwo.io>
Date: Thu, 12 Sep 2024 09:37:08 +0800
Subject: [PATCH] feat: Load DATABEND_ENTERPRISE_LICENSE_EMBEDDED while
 building (#16440)

Signed-off-by: Xuanwo <github@xuanwo.io>
Co-authored-by: Bohu <overred.shuttler@gmail.com>
---
 src/common/building/src/lib.rs              | 6 ++++++
 src/query/service/src/sessions/query_ctx.rs | 9 ++++++++-
 2 files changed, 14 insertions(+), 1 deletion(-)

diff --git a/src/common/building/src/lib.rs b/src/common/building/src/lib.rs
index fff9c2978f91..21a7459e8e98 100644
--- a/src/common/building/src/lib.rs
+++ b/src/common/building/src/lib.rs
@@ -57,6 +57,7 @@ pub fn add_building_env_vars() {
     add_env_credits_info();
     add_target_features();
     add_env_version();
+    add_env_license();
 }
 
 pub fn set_env_config() {
@@ -87,6 +88,11 @@ fn discover_version() -> Result<String> {
     }
 }
 
+pub fn add_env_license() {
+    let v = env::var("DATABEND_ENTERPRISE_LICENSE_EMBEDDED").unwrap_or_default();
+    println!("cargo:rustc-env=DATABEND_ENTERPRISE_LICENSE_EMBEDDED={v}");
+}
+
 pub fn add_env_commit_authors(repo: &Repository) {
     match git::get_commit_authors(repo) {
         Ok(authors) => println!("cargo:rustc-env=DATABEND_COMMIT_AUTHORS={}", authors),
diff --git a/src/query/service/src/sessions/query_ctx.rs b/src/query/service/src/sessions/query_ctx.rs
index f839a243c321..a466fd31c91c 100644
--- a/src/query/service/src/sessions/query_ctx.rs
+++ b/src/query/service/src/sessions/query_ctx.rs
@@ -1105,11 +1105,18 @@ impl TableContext for QueryContext {
     }
 
     fn get_license_key(&self) -> String {
-        unsafe {
+        let mut license = unsafe {
             self.get_settings()
                 .get_enterprise_license()
                 .unwrap_or_default()
+        };
+
+        // Try load license from embedded env if failed to load from settings.
+        if license.is_empty() {
+            license = env!("DATABEND_ENTERPRISE_LICENSE_EMBEDDED").to_string();
         }
+
+        license
     }
 
     fn get_query_profiles(&self) -> Vec<PlanProfile> {