fix: Set param from config (#553)

* fix: Set param from config #545
redo rebase and clean submit

* fix: misidentify column name as lateral alias (#540)

* fix: misidentify-column-name-as-alias (#539)

* add LATERAL_COLUMN_ALIAS_REFERENCE in SQLLineageConfig

* adjust import order

* add test_column_top_level_enable_lateral_ref_with_metadata_from_nested_subquery

* unknown

* refactor: rebase master and convert LATERAL_COLUMN_ALIAS_REFERENCE to bool type

* refactor: use as few condition as possible: SQLLineageConfig.LATERAL_COLUMN_ALIAS_REFERENCE

* refactor: rebase master and resolve conflict

* refactor: move logic from to_source_columns to end_of_query_cleanup

* refactor: rebase master and fix black format

* docs: LATERAL_COLUMN_ALIAS_REFERENCE how-to guide

* docs: starting version for each config

---------

Co-authored-by: reata <[email protected]>

* feat:SQLLineageConfig supports set value and thread safety

* fix: Fix mypy error

* fix: Fix pytest cov

* fix: Fix the scenario of direct assignment without using with. Add the test of multi-process scenario.

* fix: add  SQLLineageConfigLoader set function

* feat: disable setattr for SQLLineageConfig

* feat: make SQLLineageConfig context manager non-reentrant

* feat: disable set unknown config

* feat: access config in parallel

* chore: disable A005 for module name builtin conflict

* refactor: classmethod to staticmethod

---------

Co-authored-by: liuzhou <[email protected]>
Co-authored-by: maoxd <[email protected]>
Co-authored-by: reata <[email protected]>

Author: 4 people

sqllineage/config.py

@@ -1,4 +1,8 @@
 import os
+import threading
+from typing import Any, Dict, Set
+
+from sqllineage.exceptions import ConfigException
 
 
 class _SQLLineageConfigLoader:
@@ -17,10 +21,18 @@ class _SQLLineageConfigLoader:
         # lateral column alias reference supported by some dialect (redshift, spark 3.4+, etc)
         "LATERAL_COLUMN_ALIAS_REFERENCE": (bool, False),
     }
-    BOOLEAN_TRUE_STRINGS = ("true", "on", "ok", "y", "yes", "1")
 
-    def __getattr__(self, item):
-        if item in self.config:
+    def __init__(self) -> None:
+        self._thread_config: Dict[int, Dict[str, Any]] = {}
+        self._thread_in_context_manager: Set[int] = set()
+
+    def __getattr__(self, item: str):
+        if item in self.config.keys():
+            if (
+                value := self._thread_config.get(self.get_ident(), {}).get(item)
+            ) is not None:
+                return value
+
             type_, default = self.config[item]
             # require SQLLINEAGE_ prefix from environment variable
             return self.parse_value(
@@ -29,23 +41,57 @@ def __getattr__(self, item):
         else:
             return super().__getattribute__(item)
 
-    @classmethod
-    def parse_value(cls, value, cast):
-        """Parse and cast provided value
+    def __setattr__(self, key, value) -> None:
+        if key in self.config:
+            raise ConfigException(
+                "SQLLineageConfig is read-only. Use context manager to update thread level config."
+            )
+        else:
+            super().__setattr__(key, value)
+
+    def __call__(self, *args, **kwargs):
+        if self.get_ident() not in self._thread_config.keys():
+            self._thread_config[self.get_ident()] = {}
+        for key, value in kwargs.items():
+            if key in self.config.keys():
+                self._thread_config[self.get_ident()][key] = self.parse_value(
+                    value, self.config[key][0]
+                )
+            else:
+                raise ConfigException(f"Invalid config key: {key}")
+        return self
 
+    def __enter__(self):
+        if (thread_id := self.get_ident()) not in self._thread_in_context_manager:
+            self._thread_in_context_manager.add(thread_id)
+        else:
+            raise ConfigException("SQLLineageConfig context manager is not reentrant")
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        thread_id = self.get_ident()
+        if thread_id in self._thread_config:
+            self._thread_config.pop(self.get_ident())
+        if thread_id in self._thread_in_context_manager:
+            self._thread_in_context_manager.remove(thread_id)
+
+    @staticmethod
+    def get_ident() -> int:
+        return threading.get_ident()
+
+    @staticmethod
+    def parse_value(value, cast) -> Any:
+        """Parse and cast provided value
         :param value: Stringed value.
         :param cast: Type to cast return value as.
-
-        :returns: Casted value
+        :returns: cast value
         """
         if cast is bool:
             try:
                 value = int(value) != 0
             except ValueError:
-                value = value.lower().strip() in cls.BOOLEAN_TRUE_STRINGS
+                value = value.lower().strip() in ("true", "on", "ok", "y", "yes", "1")
         else:
             value = cast(value)
-
         return value
 
 

sqllineage/exceptions.py

@@ -12,3 +12,7 @@ class InvalidSyntaxException(SQLLineageException):
 
 class MetaDataProviderException(SQLLineageException):
     """Raised for MetaDataProvider errors"""
+
+
+class ConfigException(SQLLineageException):
+    """Raised for configuration errors"""

tests/core/test_config.py

@@ -1,7 +1,13 @@
+import concurrent.futures
 import os
+import random
+import time
 from unittest.mock import patch
 
+import pytest
+
 from sqllineage.config import SQLLineageConfig
+from sqllineage.exceptions import ConfigException
 
 
 @patch(
@@ -21,3 +27,53 @@ def test_config():
 
     assert type(SQLLineageConfig.TSQL_NO_SEMICOLON) is bool
     assert SQLLineageConfig.TSQL_NO_SEMICOLON is True
+
+
+def test_disable_direct_update_config():
+    with pytest.raises(ConfigException):
+        SQLLineageConfig.DEFAULT_SCHEMA = "ods"
+
+
+def test_update_config_using_context_manager():
+    with SQLLineageConfig(LATERAL_COLUMN_ALIAS_REFERENCE=True):
+        assert SQLLineageConfig.LATERAL_COLUMN_ALIAS_REFERENCE is True
+    assert SQLLineageConfig.LATERAL_COLUMN_ALIAS_REFERENCE is False
+
+    with SQLLineageConfig(DEFAULT_SCHEMA="ods"):
+        assert SQLLineageConfig.DEFAULT_SCHEMA == "ods"
+    assert SQLLineageConfig.DEFAULT_SCHEMA == ""
+
+    with SQLLineageConfig(DIRECTORY=""):
+        assert SQLLineageConfig.DIRECTORY == ""
+    assert SQLLineageConfig.DIRECTORY != ""
+
+
+def test_update_config_context_manager_non_reentrant():
+    with pytest.raises(ConfigException):
+        with SQLLineageConfig(DEFAULT_SCHEMA="ods"):
+            with SQLLineageConfig(DEFAULT_SCHEMA="dwd"):
+                pass
+
+
+def test_disable_update_unknown_config():
+    with pytest.raises(ConfigException):
+        with SQLLineageConfig(UNKNOWN_KEY="value"):
+            pass
+
+
+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):
+        # randomly sleep [0, 0.1) second to simulate real parsing scenario
+        time.sleep(random.random() * 0.1)
+        return SQLLineageConfig.DEFAULT_SCHEMA
+
+
+@pytest.mark.parametrize("pool", ["ThreadPoolExecutor", "ProcessPoolExecutor"])
+def test_config_parallel(pool: str):
+    executor_class = getattr(concurrent.futures, pool)
+    schemas = [f"db{i}" for i in range(100)]
+    with executor_class() as executor:
+        futures = [executor.submit(_check_schema, schema) for schema in schemas]
+        for i, future in enumerate(futures):
+            assert future.result() == schemas[i]

tox.ini

@@ -14,7 +14,7 @@ commands =
 [flake8]
 exclude = .tox,.git,__pycache__,build,sqllineagejs,venv,env
 max-line-length = 120
-# ignore = D100,D101
+ignore = A005,W503
 show-source = true
 enable-extensions=G
 application-import-names = sqllineage