Repository: reata/sqllineage

Issue:
support split SQL statements without semicolon in tsql
Hello,
I have Q about column lineage -f scratch.txt -l column
Issue 1:
I am trying to parse a large Store Procedure with multiple MS SQL statements and it produces bad results. So I tried to check and init__.py gives error at    line # 27        if len(holder.write) > 1:
I believe this happens because generally in SP, the statements are not separated by semicolon ;

Is there a way to detect ; automatically and add at end of each statement?

For example 
insert table x1 select b1 from foo
insert table x2 select b2 from bar 

fails but when I add semicolon ; at end of first statement, it works. The problem is most do not write ; at end while writing Stored Procs. Is the a way out?

NOTE : For TSQL syntax is a bit diff but same issue 
insert into x1 select b1 from foo 
insert into x2 select b2 from bar


Issue 2: (This is a suggestion)
TSQL has syntax like
    if object_id('tempdb..#tbl') is not null
        drop table #tbl

Can you parse this and ignore it for lineage instead of creating error?

Thanks

Discussion:
codecov[bot]: ## [Codecov](https://app.codecov.io/gh/reata/sqllineage/pull/442?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=None) Report
> Merging [#442](https://app.codecov.io/gh/reata/sqllineage/pull/442?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=None) (52558c6) into [master](https://app.codecov.io/gh/reata/sqllineage/commit/a619ae57df53823ed8ff5d2a0bdbc5e5e80c8053?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=None) (a619ae5) will **decrease** coverage by `0.15%`.
> The diff coverage is `100.00%`.

```diff
@@            Coverage Diff             @@
##           master     #442      +/-   ##
==========================================
- Coverage   99.54%   99.39%   -0.15%     
==========================================
  Files          37       38       +1     
  Lines        1959     1983      +24     
==========================================
+ Hits         1950     1971      +21     
- Misses          9       12       +3     
```


| [Files Changed](https://app.codecov.io/gh/reata/sqllineage/pull/442?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=None) | Coverage Δ | |
|---|---|---|
| [sqllineage/\_\_init\_\_.py](https://app.codecov.io/gh/reata/sqllineage/pull/442?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=None#diff-c3FsbGluZWFnZS9fX2luaXRfXy5weQ==) | `100.00% <ø> (ø)` | |
| [sqllineage/cli.py](https://app.codecov.io/gh/reata/sqllineage/pull/442?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=None#diff-c3FsbGluZWFnZS9jbGkucHk=) | `97.50% <100.00%> (+0.06%)` | :arrow_up: |
| [sqllineage/config.py](https://app.codecov.io/gh/reata/sqllineage/pull/442?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=None#diff-c3FsbGluZWFnZS9jb25maWcucHk=) | `100.00% <100.00%> (ø)` | |
| [sqllineage/core/parser/sqlfluff/analyzer.py](https://app.codecov.io/gh/reata/sqllineage/pull/442?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=None#diff-c3FsbGluZWFnZS9jb3JlL3BhcnNlci9zcWxmbHVmZi9hbmFseXplci5weQ==) | `100.00% <100.00%> (ø)` | |
| [sqllineage/drawing.py](https://app.codecov.io/gh/reata/sqllineage/pull/442?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=None#diff-c3FsbGluZWFnZS9kcmF3aW5nLnB5) | `100.00% <100.00%> (ø)` | |
| [sqllineage/runner.py](https://app.codecov.io/gh/reata/sqllineage/pull/442?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=None#diff-c3FsbGluZWFnZS9ydW5uZXIucHk=) | `100.00% <100.00%> (ø)` | |

... and [2 files with indirect coverage changes](https://app.codecov.io/gh/reata/sqllineage/pull/442/indirect-changes?src=pr&el=tree-more&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=None)



Patch:
diff --git a/sqllineage/__init__.py b/sqllineage/__init__.py
index 1aac942..958b5db 100644
--- a/sqllineage/__init__.py
+++ b/sqllineage/__init__.py
@@ -1,6 +1,3 @@
-import os
-
-
 NAME = "sqllineage"
 VERSION = "1.4.7"
 DEFAULT_LOGGING = {
@@ -31,9 +28,6 @@ DEFAULT_LOGGING = {
 }
 
 STATIC_FOLDER = "build"
-DATA_FOLDER = os.environ.get(
-    "SQLLINEAGE_DIRECTORY", os.path.join(os.path.dirname(__file__), "data")
-)
 DEFAULT_HOST = "localhost"
 DEFAULT_PORT = 5000
 SQLPARSE_DIALECT = "non-validating"
diff --git a/sqllineage/cli.py b/sqllineage/cli.py
index a80998f..20a580e 100644
--- a/sqllineage/cli.py
+++ b/sqllineage/cli.py
@@ -1,7 +1,7 @@
 import argparse
 import logging
 import logging.config
-
+import warnings
 
 from sqllineage import DEFAULT_DIALECT, DEFAULT_HOST, DEFAULT_LOGGING, DEFAULT_PORT
 from sqllineage.drawing import draw_lineage_graph
@@ -78,9 +78,7 @@ def main(args=None) -> None:
     )
     args = parser.parse_args(args)
     if args.e and args.f:
-        logging.warning(
-            "Both -e and -f options are specified. -e option will be ignored"
-        )
+        warnings.warn("Both -e and -f options are specified. -e option will be ignored")
     if args.f or args.e:
         sql = extract_sql_from_args(args)
         runner = LineageRunner(
diff --git a/sqllineage/config.py b/sqllineage/config.py
new file mode 100644
index 0000000..68bf0af
--- /dev/null
+++ b/sqllineage/config.py
@@ -0,0 +1,26 @@
+import os
+
+
+class _SQLLineageConfigLoader:
+    """
+    Load all configurable items from environment variable, otherwise fallback to default
+    """
+
+    # inspired by https://github.com/joke2k/django-environ
+    config = {
+        # for frontend directory drawer
+        "DIRECTORY": (str, os.path.join(os.path.dirname(__file__), "data")),
+        # to enable tsql no semicolon splitter mode
+        "TSQL_NO_SEMICOLON": (bool, False),
+    }
+
+    def __getattr__(self, item):
+        if item in self.config:
+            type_, default = self.config[item]
+            # require SQLLINEAGE_ prefix from environment variable
+            return type_(os.environ.get("SQLLINEAGE_" + item, default))
+        else:
+            return super().__getattribute__(item)
+
+
+SQLLineageConfig = _SQLLineageConfigLoader()
diff --git a/sqllineage/core/parser/sqlfluff/analyzer.py b/sqllineage/core/parser/sqlfluff/analyzer.py
index 03d6ada..9a31997 100644
--- a/sqllineage/core/parser/sqlfluff/analyzer.py
+++ b/sqllineage/core/parser/sqlfluff/analyzer.py
@@ -1,6 +1,8 @@
 import warnings
+from typing import Dict, List
 
 from sqlfluff.core import Linter, SQLLexError, SQLParseError
+from sqlfluff.core.parser import BaseSegment
 
 from sqllineage.core.analyzer import LineageAnalyzer
 from sqllineage.core.holders import StatementLineageHolder
@@ -17,8 +19,46 @@ class SqlFluffLineageAnalyzer(LineageAnalyzer):
 
     def __init__(self, dialect: str):
         self._dialect = dialect
+        self.tsql_split_cache: Dict[str, BaseSegment] = {}
+
+    def split_tsql(self, sql: str) -> List[str]:
+        """
+        use sqlfluff parse to split tsql statements. This is in particular for semicolon not present cases.
+        The result is cached so that later analyze method doesn't have to parse regarding single statement sql.
+        """
+        sqls = []
+        for segment in self._list_specific_statement_segment(sql):
+            self.tsql_split_cache[segment.raw] = segment
+            sqls.append(segment.raw)
+        return sqls
 
     def analyze(self, sql: str) -> StatementLineageHolder:
+        if sql in self.tsql_split_cache:
+            statement_segments = [self.tsql_split_cache[sql]]
+        else:
+            statement_segments = self._list_specific_statement_segment(sql)
+        if len(statement_segments) == 0:
+            raise UnsupportedStatementException(
+                f"SQLLineage cannot parse SQL:" f"{sql}"
+            )  # pragma: no cover
+        else:
+            statement_segment = statement_segments[0]
+            for extractor in [
+                extractor_cls(self._dialect)
+                for extractor_cls in BaseExtractor.__subclasses__()
+            ]:
+                if extractor.can_extract(statement_segment.type):
+                    lineage_holder = extractor.extract(
+                        statement_segment, AnalyzerContext()
+                    )
+                    return StatementLineageHolder.of(lineage_holder)
+            else:
+                raise UnsupportedStatementException(
+                    f"SQLLineage doesn't support analyzing statement type [{statement_segment.type}] for SQL:"
+                    f"{sql}"
+                )
+
+    def _list_specific_statement_segment(self, sql: str):
         parsed = Linter(dialect=self._dialect).parse_string(sql)
         violations = [
             str(e)
@@ -32,12 +72,10 @@ class SqlFluffLineageAnalyzer(LineageAnalyzer):
                 f"{sql}\n"
                 f"{violation_msg}"
             )
-
-        statement_segment = None
+        segments = []
         for top_segment in getattr(parsed.tree, "segments", []):
             if top_segment.type == "statement":
-                statement_segment = top_segment.segments[0]
-                break
+                segments.append(top_segment.segments[0])
             elif top_segment.type == "batch":
                 statements = top_segment.get_children("statement")
                 if len(statements) > 1:
@@ -47,24 +85,6 @@ class SqlFluffLineageAnalyzer(LineageAnalyzer):
                         SyntaxWarning,
                         stacklevel=2,
                     )
-                statement_segment = statements[0].segments[0]
-                break
-        if statement_segment is None:
-            raise UnsupportedStatementException(
-                f"SQLLineage cannot parse SQL:" f"{sql}"
-            )  # pragma: no cover
-        else:
-            for extractor in [
-                extractor_cls(self._dialect)
-                for extractor_cls in BaseExtractor.__subclasses__()
-            ]:
-                if extractor.can_extract(statement_segment.type):
-                    lineage_holder = extractor.extract(
-                        statement_segment, AnalyzerContext()
-                    )
-                    return StatementLineageHolder.of(lineage_holder)
-            else:
-                raise UnsupportedStatementException(
-                    f"SQLLineage doesn't support analyzing statement type [{statement_segment.type}] for SQL:"
-                    f"{sql}"
-                )
+                for statement in statements:
+                    segments.append(statement.segments[0])
+        return segments
diff --git a/sqllineage/drawing.py b/sqllineage/drawing.py
index 4c6bf1a..8659ff8 100644
--- a/sqllineage/drawing.py
+++ b/sqllineage/drawing.py
@@ -17,13 +17,8 @@ from typing import Any, Callable, Dict, List
 from urllib.parse import urlencode
 from wsgiref.simple_server import make_server
 
-from sqllineage import (
-    DATA_FOLDER,
-    DEFAULT_DIALECT,
-    DEFAULT_HOST,
-    DEFAULT_PORT,
-)
-from sqllineage import STATIC_FOLDER
+from sqllineage import DEFAULT_DIALECT, DEFAULT_HOST, DEFAULT_PORT, STATIC_FOLDER
+from sqllineage.config import SQLLineageConfig
 from sqllineage.exceptions import SQLLineageException
 from sqllineage.utils.constant import LineageLevel
 from sqllineage.utils.helpers import extract_sql_from_args
@@ -34,7 +29,7 @@ logger = logging.getLogger(__name__)
 class SQLLineageApp:
     def __init__(self) -> None:
         self.routes: Dict[str, Callable[[Dict[str, Any]], Dict[str, Any]]] = {}
-        self.root_path = Path(DATA_FOLDER)
+        self.root_path = Path(SQLLineageConfig.DIRECTORY)
 
     def route(self, path: str):
         def wrapper(handler):
@@ -193,7 +188,7 @@ def directory(payload):
     elif payload.get("d"):
         root = Path(payload["d"])
     else:
-        root = Path(DATA_FOLDER)
+        root = Path(SQLLineageConfig.DIRECTORY)
     data = {
         "id": str(root),
         "name": root.name,
diff --git a/sqllineage/runner.py b/sqllineage/runner.py
index 8a21124..231fa99 100644
--- a/sqllineage/runner.py
+++ b/sqllineage/runner.py
@@ -3,6 +3,7 @@ import warnings
 from typing import Dict, List, Optional, Tuple
 
 from sqllineage import DEFAULT_DIALECT, SQLPARSE_DIALECT
+from sqllineage.config import SQLLineageConfig
 from sqllineage.core.holders import SQLLineageHolder
 from sqllineage.core.models import Column, Table
 from sqllineage.core.parser.sqlfluff.analyzer import SqlFluffLineageAnalyzer
@@ -168,12 +169,20 @@ Target Tables:
         print(str(self))
 
     def _eval(self):
-        self._stmt = split(self._sql.strip())
         analyzer = (
             SqlParseLineageAnalyzer()
             if self._dialect == SQLPARSE_DIALECT
             else SqlFluffLineageAnalyzer(self._dialect)
         )
+        if SQLLineageConfig.TSQL_NO_SEMICOLON and self._dialect == "tsql":
+            self._stmt = analyzer.split_tsql(self._sql.strip())
+        else:
+            if SQLLineageConfig.TSQL_NO_SEMICOLON and self._dialect != "tsql":
+                warnings.warn(
+                    f"Dialect={self._dialect}, TSQL_NO_SEMICOLON will be ignored unless dialect is tsql"
+                )
+            self._stmt = split(self._sql.strip())
+
         self._stmt_holders = [analyzer.analyze(stmt) for stmt in self._stmt]
         self._sql_holder = SQLLineageHolder.of(*self._stmt_holders)
         self._evaluated = True


Test patch:
diff --git a/tests/core/test_cli.py b/tests/core/test_cli.py
index e5b8852..0752dd7 100644
--- a/tests/core/test_cli.py
+++ b/tests/core/test_cli.py
@@ -4,8 +4,8 @@ from unittest.mock import patch
 
 import pytest
 
-from sqllineage import DATA_FOLDER
 from sqllineage.cli import main
+from sqllineage.config import SQLLineageConfig
 
 
 @patch("socketserver.BaseServer.serve_forever")
@@ -13,7 +13,7 @@ def test_cli_dummy(_):
     main([])
     main(["-e", "select * from dual"])
     main(["-e", "insert into foo select * from dual", "-l", "column"])
-    for dirname, _, files in os.walk(DATA_FOLDER):
+    for dirname, _, files in os.walk(SQLLineageConfig.DIRECTORY):
         if len(files) > 0:
             sql_file = str(Path(dirname).joinpath(Path(files[0])))
             main(["-f", sql_file])
diff --git a/tests/core/test_drawing.py b/tests/core/test_drawing.py
index 103a206..d26ec2b 100644
--- a/tests/core/test_drawing.py
+++ b/tests/core/test_drawing.py
@@ -4,7 +4,7 @@ from collections import namedtuple
 from http import HTTPStatus
 from io import StringIO
 
-from sqllineage import DATA_FOLDER
+from sqllineage.config import SQLLineageConfig
 from sqllineage.drawing import app
 
 
@@ -39,10 +39,16 @@ def test_handler():
     mock_request("POST", "/script", {"e": "SELECT * FROM dual", "p": 5000})
     assert container.status.startswith(str(HTTPStatus.OK.value))
     mock_request(
-        "POST", "/directory", {"f": os.path.join(DATA_FOLDER, "tpcds/query01.sql")}
+        "POST",
+        "/directory",
+        {"f": os.path.join(SQLLineageConfig.DIRECTORY, "tpcds/query01.sql")},
     )
     assert container.status.startswith(str(HTTPStatus.OK.value))
-    mock_request("POST", "/directory", {"d": os.path.join(DATA_FOLDER, "tpcds/")})
+    mock_request(
+        "POST",
+        "/directory",
+        {"d": os.path.join(SQLLineageConfig.DIRECTORY, "tpcds/")},
+    )
     assert container.status.startswith(str(HTTPStatus.OK.value))
     mock_request("POST", "/directory", {})
     assert container.status.startswith(str(HTTPStatus.OK.value))
@@ -62,7 +68,9 @@ def test_handler():
     mock_request("GET", "/static")
     assert container.status.startswith(str(HTTPStatus.NOT_FOUND.value))
     mock_request(
-        "POST", "/script", {"f": os.path.join(DATA_FOLDER, "tpcds/query100.sql")}
+        "POST",
+        "/script",
+        {"f": os.path.join(SQLLineageConfig.DIRECTORY, "tpcds/query100.sql")},
     )
     assert container.status.startswith(str(HTTPStatus.NOT_FOUND.value))
     mock_request("POST", "/non-exist-resource", {"e": "SELECT * FROM where foo='bar'"})
diff --git a/tests/core/test_exception.py b/tests/core/test_exception.py
index 49f4b85..1d287db 100644
--- a/tests/core/test_exception.py
+++ b/tests/core/test_exception.py
@@ -1,4 +1,4 @@
-import warnings
+from unittest.mock import patch
 
 import pytest
 
@@ -43,19 +43,25 @@ def test_unsupported_query_type_in_sqlfluff():
         )._eval()
 
 
-def test_deprecated_warning_in_sqlparse():
-    with warnings.catch_warnings(record=True) as w:
+def test_deprecation_warning_in_sqlparse():
+    with pytest.warns(DeprecationWarning):
         LineageRunner("SELECT * FROM DUAL", dialect="non-validating")._eval()
-        assert len(w) == 1
-        assert issubclass(w[0].category, DeprecationWarning)
 
 
 def test_syntax_warning_no_semicolon_in_tsql():
-    with warnings.catch_warnings(record=True) as w:
+    with pytest.warns(SyntaxWarning):
         LineageRunner(
             """SELECT * FROM foo
 SELECT * FROM bar""",
             dialect="tsql",
         )._eval()
-        assert len(w) == 1
-        assert issubclass(w[0].category, SyntaxWarning)
+
+
+@patch("os.environ", {"SQLLINEAGE_TSQL_NO_SEMICOLON": "TRUE"})
+def test_user_warning_enable_tsql_no_semicolon_with_other_dialect():
+    with pytest.warns(UserWarning):
+        LineageRunner(
+            """SELECT * FROM foo;
+SELECT * FROM bar""",
+            dialect="ansi",
+        )._eval()
diff --git a/tests/sql/sqlfluff_only/test_tsql_no_semicolon.py b/tests/sql/sqlfluff_only/test_tsql_no_semicolon.py
new file mode 100644
index 0000000..a75646e
--- /dev/null
+++ b/tests/sql/sqlfluff_only/test_tsql_no_semicolon.py
@@ -0,0 +1,21 @@
+from unittest.mock import patch
+
+import pytest
+from tests.helpers import assert_table_lineage_equal
+
+
+@patch("os.environ", {"SQLLINEAGE_TSQL_NO_SEMICOLON": "TRUE"})
+@pytest.mark.parametrize("dialect", ["tsql"])
+def test_tsql_multi_statement_no_semicolon(dialect: str):
+    """
+    tsql multiple statements without explicit semicolon as splitter.
+    """
+    sql = """insert into tab1 select * from foo
+insert into tab2 select * from bar"""
+    assert_table_lineage_equal(
+        sql,
+        {"foo", "bar"},
+        {"tab1", "tab2"},
+        dialect=dialect,
+        test_sqlparse=False,
+    )