Repository: reata/sqllineage

Issue:
Support Change Configure Default Schema At Runtime
Continue on the undone feature on #523.

#523 add an environment variable to configure default schema. That solves part of the problem but it would still be preferable if we can support configure through LineageRunner.

Create this issue to continue tracking on that.



Discussion:
最新 master的实现 无法正常完成功能。
因为 Schema init 函数 传入时已经完成变量绑定
我马上提交测试用例和修复PR

The latest master implementation cannot complete the function properly.
Because the variable binding has been completed when Schema init function is passed in
I will submit test cases and fix PR immediately





Patch:
diff --git a/docs/gear_up/configuration.rst b/docs/gear_up/configuration.rst
index 8653358f..a0dbd027 100644
--- a/docs/gear_up/configuration.rst
+++ b/docs/gear_up/configuration.rst
@@ -7,6 +7,24 @@ The SQLLineage configuration allows user to customize the behaviour of sqllineag
 We adopt environment variable approach for global key-value mapping. The keys listed in this section should start with
 `"SQLLINEAGE_"` to be a valid config. For example, to use DEFAULT_SCHEMA, use ``SQLLINEAGE_DEFAULT_SCHEMA=default``.
 
+.. note::
+     Starting v1.5.2, we also support changing config at runtime. A local config is kept for each thread that will mask
+     global config. Local configuration must be set using context manager:
+
+     .. code-block:: python
+
+        >>> from sqllineage.config import SQLLineageConfig
+        >>> from sqllineage.runner import LineageRunner
+
+        >>> with SQLLineageConfig(DEFAULT_SCHEMA="ods"):
+        >>>     print(LineageRunner("select * from test").source_tables)
+        [Table: ods.test]
+        >>> with SQLLineageConfig(DEFAULT_SCHEMA="dwd"):
+        >>>     print(LineageRunner("select * from test").source_tables)
+        [Table: dwd.test]
+
+     Note when setting local config, the key does not start with `"SQLLINEAGE_"`.
+
 DEFAULT_SCHEMA
 ==============
 Default schema, or interchangeably called database. Tables without schema qualifier parsed from SQL is set with a schema
diff --git a/sqllineage/core/models.py b/sqllineage/core/models.py
index 01e6985c..49ce17a5 100644
--- a/sqllineage/core/models.py
+++ b/sqllineage/core/models.py
@@ -12,13 +12,17 @@ class Schema:
     """
 
     unknown = "<default>"
-    default = SQLLineageConfig.DEFAULT_SCHEMA or unknown
 
-    def __init__(self, name: str = default):
+    def __init__(self, name: Optional[str] = None):
         """
         :param name: schema name
         """
-        self.raw_name = escape_identifier_name(name)
+        if name:
+            self.raw_name = escape_identifier_name(name)
+        elif SQLLineageConfig.DEFAULT_SCHEMA:
+            self.raw_name = escape_identifier_name(SQLLineageConfig.DEFAULT_SCHEMA)
+        else:
+            self.raw_name = escape_identifier_name(Schema.unknown)
 
     def __str__(self):
         return self.raw_name
@@ -33,7 +37,7 @@ def __hash__(self):
         return hash(str(self))
 
     def __bool__(self):
-        return str(self) != self.default
+        return str(self) != self.unknown
 
 
 class Table:


Test patch:
diff --git a/tests/core/test_config.py b/tests/core/test_config.py
index 9d7ed46d..e3b96e63 100644
--- a/tests/core/test_config.py
+++ b/tests/core/test_config.py
@@ -8,6 +8,7 @@
 
 from sqllineage.config import SQLLineageConfig
 from sqllineage.exceptions import ConfigException
+from sqllineage.runner import LineageRunner
 
 
 @patch(
@@ -64,9 +65,10 @@ def test_disable_update_unknown_config():
 def _check_schema(schema: str):
     # used by test_config_parallel, must be a global function so that it can be pickled between processes
     with SQLLineageConfig(DEFAULT_SCHEMA=schema):
+        table = LineageRunner("select * from test").source_tables.pop()
         # randomly sleep [0, 0.1) second to simulate real parsing scenario
         time.sleep(random.random() * 0.1)
-        return SQLLineageConfig.DEFAULT_SCHEMA
+        return table.schema.raw_name
 
 
 @pytest.mark.parametrize("pool", ["ThreadPoolExecutor", "ProcessPoolExecutor"])