From 6189d312b3f53d4302bb951a3297a06321a53a2d Mon Sep 17 00:00:00 2001 From: reata Date: Fri, 24 May 2024 22:53:07 +0800 Subject: [PATCH] feat: allow tsql identifier enclosed in brackets (#619) --- sqllineage/utils/helpers.py | 4 ++++ tests/sql/table/test_select_dialect_specific.py | 12 ++++++++++++ 2 files changed, 16 insertions(+) diff --git a/sqllineage/utils/helpers.py b/sqllineage/utils/helpers.py index e403ba79..df9c3ede 100644 --- a/sqllineage/utils/helpers.py +++ b/sqllineage/utils/helpers.py @@ -17,6 +17,10 @@ def escape_identifier_name(name: str): for quote_char in quote_chars: name = name.strip(quote_char) return name + elif name.startswith("[") and name.endswith("]"): + # tsql allows quoted identifier with square brackets, see reference + # https://learn.microsoft.com/en-us/sql/relational-databases/databases/database-identifiers?view=sql-server-ver16#classes-of-identifiers + return name.strip("[]") else: return name.lower() diff --git a/tests/sql/table/test_select_dialect_specific.py b/tests/sql/table/test_select_dialect_specific.py index 9755fa57..c53fc34c 100644 --- a/tests/sql/table/test_select_dialect_specific.py +++ b/tests/sql/table/test_select_dialect_specific.py @@ -24,6 +24,18 @@ def test_select_with_schema_in_backtick(dialect: str): ) +@pytest.mark.parametrize("dialect", ["tsql"]) +def test_select_with_table_name_in_bracket(dialect: str): + assert_table_lineage_equal("SELECT * FROM [tab1]", {"tab1"}, dialect=dialect) + + +@pytest.mark.parametrize("dialect", ["tsql"]) +def test_select_with_schema_in_bracket(dialect: str): + assert_table_lineage_equal( + "SELECT * FROM [schema1].[tab1]", {"schema1.tab1"}, dialect=dialect + ) + + @pytest.mark.parametrize("dialect", ["databricks", "hive", "sparksql"]) def test_select_left_semi_join(dialect: str): assert_table_lineage_equal(