acryl-datahub 0.15.0.6rc2__py3-none-any.whl → 1.0.0__py3-none-any.whl

datahub/sql_parsing/split_statements.py

@@ -1,26 +1,48 @@
 import re
 from enum import Enum
-from typing import Generator, List, Tuple
+from typing import Iterator, List, Tuple
+
+SELECT_KEYWORD = "SELECT"
+CASE_KEYWORD = "CASE"
+END_KEYWORD = "END"
 
 CONTROL_FLOW_KEYWORDS = [
     "GO",
-    r"BEGIN\w+TRY",
-    r"BEGIN\w+CATCH",
+    r"BEGIN\s+TRY",
+    r"BEGIN\s+CATCH",
     "BEGIN",
-    r"END\w+TRY",
-    r"END\w+CATCH",
-    "END",
+    r"END\s+TRY",
+    r"END\s+CATCH",
+    # This isn't strictly correct, but we assume that IF | (condition) | (block) should all be split up
+    # This mainly ensures that IF statements don't get tacked onto the previous statement incorrectly
+    "IF",
+    # For things like CASE, END does not mean the end of a statement.
+    # We have special handling for this.
+    END_KEYWORD,
+    # "ELSE",  # else is also valid in CASE, so we we can't use it here.
 ]
 
 # There's an exception to this rule, which is when the statement
-# is preceeded by a CTE.
-FORCE_NEW_STATEMENT_KEYWORDS = [
+# is preceded by a CTE. For those, we have to check if the character
+# before this is a ")".
+NEW_STATEMENT_KEYWORDS = [
     # SELECT is used inside queries as well, so we can't include it here.
+    "CREATE",
     "INSERT",
     "UPDATE",
     "DELETE",
     "MERGE",
 ]
+STRICT_NEW_STATEMENT_KEYWORDS = [
+    # For these keywords, a SELECT following it does indicate a new statement.
+    "DROP",
+    "TRUNCATE",
+]
+
+
+class _AlreadyIncremented(Exception):
+    # Using exceptions for control flow isn't great - but the code is clearer so it's fine.
+    pass
 
 
 class ParserState(Enum):
@@ -30,134 +52,206 @@ class ParserState(Enum):
     MULTILINE_COMMENT = 4
 
 
-def _is_keyword_at_position(sql: str, pos: int, keyword: str) -> bool:
-    """
-    Check if a keyword exists at the given position using regex word boundaries.
-    """
-    if pos + len(keyword) > len(sql):
-        return False
+class _StatementSplitter:
+    def __init__(self, sql: str):
+        self.sql = sql
 
-    # If we're not at a word boundary, we can't generate a keyword.
-    if pos > 0 and not (
-        bool(re.match(r"\w\W", sql[pos - 1 : pos + 1]))
-        or bool(re.match(r"\W\w", sql[pos - 1 : pos + 1]))
-    ):
-        return False
+        # Main parser state.
+        self.i = 0
+        self.state = ParserState.NORMAL
+        self.current_statement: List[str] = []
 
-    pattern = rf"^{re.escape(keyword)}\b"
-    match = re.match(pattern, sql[pos:], re.IGNORECASE)
-    return bool(match)
+        # Additional parser state.
 
+        # If we see a SELECT, should we start a new statement?
+        # If we previously saw a drop/truncate/etc, a SELECT does mean a new statement.
+        # But if we're in a select/create/etc, a select could just be a subquery.
+        self.does_select_mean_new_statement = False
 
-def _look_ahead_for_keywords(
-    sql: str, pos: int, keywords: List[str]
-) -> Tuple[bool, str, int]:
-    """
-    Look ahead for SQL keywords at the current position.
-    """
+        # The END keyword terminates CASE and BEGIN blocks.
+        # We need to match the CASE statements with END blocks to determine
+        # what a given END is closing.
+        self.current_case_statements = 0
 
-    for keyword in keywords:
-        if _is_keyword_at_position(sql, pos, keyword):
-            return True, keyword, len(keyword)
-    return False, "", 0
+    def _is_keyword_at_position(self, pos: int, keyword: str) -> Tuple[bool, str]:
+        """
+        Check if a keyword exists at the given position using regex word boundaries.
+        """
+        sql = self.sql
 
+        keyword_length = len(keyword.replace(r"\s+", " "))
 
-def split_statements(sql: str) -> Generator[str, None, None]:
-    """
-    Split T-SQL code into individual statements, handling various SQL constructs.
-    """
-    if not sql or not sql.strip():
-        return
+        if pos + keyword_length > len(sql):
+            return False, ""
+
+        # If we're not at a word boundary, we can't generate a keyword.
+        if pos > 0 and not (
+            bool(re.match(r"\w\W", sql[pos - 1 : pos + 1]))
+            or bool(re.match(r"\W\w", sql[pos - 1 : pos + 1]))
+        ):
+            return False, ""
+
+        pattern = rf"^{keyword}\b"
+        match = re.match(pattern, sql[pos:], re.IGNORECASE)
+        is_match = bool(match)
+        actual_match = (
+            sql[pos:][match.start() : match.end()] if match is not None else ""
+        )
+        return is_match, actual_match
 
-    current_statement: List[str] = []
-    state = ParserState.NORMAL
-    i = 0
+    def _look_ahead_for_keywords(self, keywords: List[str]) -> Tuple[bool, str, int]:
+        """
+        Look ahead for SQL keywords at the current position.
+        """
 
-    def yield_if_complete() -> Generator[str, None, None]:
-        statement = "".join(current_statement).strip()
+        for keyword in keywords:
+            is_match, keyword = self._is_keyword_at_position(self.i, keyword)
+            if is_match:
+                return True, keyword, len(keyword)
+        return False, "", 0
+
+    def _yield_if_complete(self) -> Iterator[str]:
+        statement = "".join(self.current_statement).strip()
         if statement:
+            # Subtle - to avoid losing full whitespace, they get merged into the next statement.
             yield statement
-            current_statement.clear()
-
-    prev_real_char = "\0"  # the most recent non-whitespace, non-comment character
-    while i < len(sql):
-        c = sql[i]
-        next_char = sql[i + 1] if i < len(sql) - 1 else "\0"
-
-        if state == ParserState.NORMAL:
-            if c == "'":
-                state = ParserState.STRING
-                current_statement.append(c)
-                prev_real_char = c
-            elif c == "-" and next_char == "-":
-                state = ParserState.COMMENT
-                current_statement.append(c)
-                current_statement.append(next_char)
-                i += 1
-            elif c == "/" and next_char == "*":
-                state = ParserState.MULTILINE_COMMENT
-                current_statement.append(c)
-                current_statement.append(next_char)
-                i += 1
-            else:
-                most_recent_real_char = prev_real_char
-                if not c.isspace():
-                    prev_real_char = c
+            self.current_statement.clear()
 
-                is_control_keyword, keyword, keyword_len = _look_ahead_for_keywords(
-                    sql, i, keywords=CONTROL_FLOW_KEYWORDS
-                )
-                if is_control_keyword:
-                    # Yield current statement if any
-                    yield from yield_if_complete()
-                    # Yield keyword as its own statement
-                    yield keyword
-                    i += keyword_len
-                    continue
-
-                (
-                    is_force_new_statement_keyword,
-                    keyword,
-                    keyword_len,
-                ) = _look_ahead_for_keywords(
-                    sql, i, keywords=FORCE_NEW_STATEMENT_KEYWORDS
-                )
-                if (
-                    is_force_new_statement_keyword and most_recent_real_char != ")"
-                ):  # usually we'd have a close paren that closes a CTE
-                    # Force termination of current statement
-                    yield from yield_if_complete()
-
-                    current_statement.append(keyword)
-                    i += keyword_len
-                    continue
-
-                elif c == ";":
-                    yield from yield_if_complete()
+        # Reset current_statement-specific state.
+        self.does_select_mean_new_statement = False
+        if self.current_case_statements != 0:
+            breakpoint()
+        self.current_case_statements = 0
+
+    def process(self) -> Iterator[str]:
+        if not self.sql or not self.sql.strip():
+            yield from ()
+
+        prev_real_char = "\0"  # the most recent non-whitespace, non-comment character
+        while self.i < len(self.sql):
+            c = self.sql[self.i]
+            next_char = self.sql[self.i + 1] if self.i < len(self.sql) - 1 else "\0"
+
+            if self.state == ParserState.NORMAL:
+                if c == "'":
+                    self.state = ParserState.STRING
+                    self.current_statement.append(c)
+                    prev_real_char = c
+                elif c == "-" and next_char == "-":
+                    self.state = ParserState.COMMENT
+                    self.current_statement.append(c)
+                    self.current_statement.append(next_char)
+                    self.i += 1
+                elif c == "/" and next_char == "*":
+                    self.state = ParserState.MULTILINE_COMMENT
+                    self.current_statement.append(c)
+                    self.current_statement.append(next_char)
+                    self.i += 1
                 else:
-                    current_statement.append(c)
-
-        elif state == ParserState.STRING:
-            current_statement.append(c)
-            if c == "'" and next_char == "'":
-                current_statement.append(next_char)
-                i += 1
-            elif c == "'":
-                state = ParserState.NORMAL
-
-        elif state == ParserState.COMMENT:
-            current_statement.append(c)
-            if c == "\n":
-                state = ParserState.NORMAL
-
-        elif state == ParserState.MULTILINE_COMMENT:
-            current_statement.append(c)
-            if c == "*" and next_char == "/":
-                current_statement.append(next_char)
-                i += 1
-                state = ParserState.NORMAL
-
-        i += 1
-
-    # Handle the last statement
-    yield from yield_if_complete()
+                    most_recent_real_char = prev_real_char
+                    if not c.isspace():
+                        prev_real_char = c
+
+                    try:
+                        yield from self._process_normal(
+                            most_recent_real_char=most_recent_real_char
+                        )
+                    except _AlreadyIncremented:
+                        # Skip the normal i += 1 step.
+                        continue
+
+            elif self.state == ParserState.STRING:
+                self.current_statement.append(c)
+                if c == "'" and next_char == "'":
+                    self.current_statement.append(next_char)
+                    self.i += 1
+                elif c == "'":
+                    self.state = ParserState.NORMAL
+
+            elif self.state == ParserState.COMMENT:
+                self.current_statement.append(c)
+                if c == "\n":
+                    self.state = ParserState.NORMAL
+
+            elif self.state == ParserState.MULTILINE_COMMENT:
+                self.current_statement.append(c)
+                if c == "*" and next_char == "/":
+                    self.current_statement.append(next_char)
+                    self.i += 1
+                    self.state = ParserState.NORMAL
+
+            self.i += 1
+
+        # Handle the last statement
+        yield from self._yield_if_complete()
+
+    def _process_normal(self, most_recent_real_char: str) -> Iterator[str]:
+        c = self.sql[self.i]
+
+        if self._is_keyword_at_position(self.i, CASE_KEYWORD)[0]:
+            self.current_case_statements += 1
+
+        is_control_keyword, keyword, keyword_len = self._look_ahead_for_keywords(
+            keywords=CONTROL_FLOW_KEYWORDS
+        )
+        if (
+            is_control_keyword
+            and keyword == END_KEYWORD
+            and self.current_case_statements > 0
+        ):
+            # If we're closing a CASE statement with END, we can just decrement the counter and continue.
+            self.current_case_statements -= 1
+        elif is_control_keyword:
+            # Yield current statement if any
+            yield from self._yield_if_complete()
+            # Yield keyword as its own statement
+            yield keyword
+            self.i += keyword_len
+            self.does_select_mean_new_statement = True
+            raise _AlreadyIncremented()
+
+        (
+            is_strict_new_statement_keyword,
+            keyword,
+            keyword_len,
+        ) = self._look_ahead_for_keywords(keywords=STRICT_NEW_STATEMENT_KEYWORDS)
+        if is_strict_new_statement_keyword:
+            yield from self._yield_if_complete()
+            self.current_statement.append(keyword)
+            self.i += keyword_len
+            self.does_select_mean_new_statement = True
+            raise _AlreadyIncremented()
+
+        (
+            is_force_new_statement_keyword,
+            keyword,
+            keyword_len,
+        ) = self._look_ahead_for_keywords(
+            keywords=(
+                NEW_STATEMENT_KEYWORDS
+                + ([SELECT_KEYWORD] if self.does_select_mean_new_statement else [])
+            ),
+        )
+        if (
+            is_force_new_statement_keyword and most_recent_real_char != ")"
+        ):  # usually we'd have a close paren that closes a CTE
+            # Force termination of current statement
+            yield from self._yield_if_complete()
+
+            self.current_statement.append(keyword)
+            self.i += keyword_len
+            raise _AlreadyIncremented()
+
+        if c == ";":
+            yield from self._yield_if_complete()
+        else:
+            self.current_statement.append(c)
+
+
+def split_statements(sql: str) -> Iterator[str]:
+    """
+    Split T-SQL code into individual statements, handling various SQL constructs.
+    """
+
+    splitter = _StatementSplitter(sql)
+    yield from splitter.process()

datahub/sql_parsing/sql_parsing_common.py

@@ -24,12 +24,19 @@ DIALECTS_WITH_CASE_INSENSITIVE_COLS = {
     # For SQL server, the default collation rules mean that all identifiers (schema, table, column names)
     # are case preserving but case insensitive.
     "mssql",
+    # Oracle automatically converts unquoted identifiers to uppercase.
+    # https://docs.oracle.com/en/database/oracle/oracle-database/19/sqlrf/Database-Object-Names-and-Qualifiers.html#GUID-3C59E44A-5140-4BCA-B9E1-3039C8050C49
+    # In our Oracle connector, we then normalize column names to lowercase. This behavior
+    # actually comes from the underlying Oracle sqlalchemy dialect.
+    # https://github.com/sqlalchemy/sqlalchemy/blob/d9b4d8ff3aae504402d324f3ebf0b8faff78f5dc/lib/sqlalchemy/dialects/oracle/base.py#L2579
+    "oracle",
 }
 DIALECTS_WITH_DEFAULT_UPPERCASE_COLS = {
     # In some dialects, column identifiers are effectively case insensitive
     # because they are automatically converted to uppercase. Most other systems
     # automatically lowercase unquoted identifiers.
     "snowflake",
+    "oracle",
 }
 assert DIALECTS_WITH_DEFAULT_UPPERCASE_COLS.issubset(
     DIALECTS_WITH_CASE_INSENSITIVE_COLS

datahub/sql_parsing/sqlglot_lineage.py

@@ -473,7 +473,7 @@ def _create_table_ddl_cll(
     return column_lineage
 
 
-def _select_statement_cll(  # noqa: C901
+def _select_statement_cll(
     statement: _SupportedColumnLineageTypes,
     dialect: sqlglot.Dialect,
     root_scope: sqlglot.optimizer.Scope,

datahub/sql_parsing/sqlglot_utils.py

@@ -56,10 +56,7 @@ def get_dialect(platform: DialectOrStr) -> sqlglot.Dialect:
 def is_dialect_instance(
     dialect: sqlglot.Dialect, platforms: Union[str, Iterable[str]]
 ) -> bool:
-    if isinstance(platforms, str):
-        platforms = [platforms]
-    else:
-        platforms = list(platforms)
+    platforms = [platforms] if isinstance(platforms, str) else list(platforms)
 
     dialects = [get_dialect(platform) for platform in platforms]
 

datahub/testing/check_sql_parser_result.py

@@ -1,5 +1,4 @@
 import logging
-import os
 import pathlib
 from typing import Any, Dict, Optional
 
@@ -8,11 +7,10 @@ import deepdiff
 from datahub.ingestion.source.bigquery_v2.bigquery_audit import BigqueryTableIdentifier
 from datahub.sql_parsing.schema_resolver import SchemaInfo, SchemaResolver
 from datahub.sql_parsing.sqlglot_lineage import SqlParsingResult, sqlglot_lineage
+from datahub.testing.pytest_hooks import get_golden_settings
 
 logger = logging.getLogger(__name__)
 
-UPDATE_FILES = os.environ.get("UPDATE_SQLPARSER_FILES", "false").lower() == "true"
-
 
 def assert_sql_result_with_resolver(
     sql: str,
@@ -22,6 +20,8 @@ def assert_sql_result_with_resolver(
     allow_table_error: bool = False,
     **kwargs: Any,
 ) -> None:
+    settings = get_golden_settings()
+
     # HACK: Our BigQuery source overwrites this value and doesn't undo it.
     # As such, we need to handle that here.
     BigqueryTableIdentifier._BQ_SHARDED_TABLE_SUFFIX = "_yyyymmdd"
@@ -47,15 +47,14 @@ def assert_sql_result_with_resolver(
         )
 
     txt = res.json(indent=4)
-    if UPDATE_FILES:
+    if settings.update_golden:
         expected_file.write_text(txt)
         return
 
     if not expected_file.exists():
         expected_file.write_text(txt)
         raise AssertionError(
-            f"Expected file {expected_file} does not exist. "
-            "Created it with the expected output. Please verify it."
+            f"Missing expected golden file; run with --update-golden-files to create it: {expected_file}"
         )
 
     expected = SqlParsingResult.parse_raw(expected_file.read_text())

datahub/testing/compare_metadata_json.py

@@ -16,6 +16,7 @@ from deepdiff import DeepDiff
 from datahub.ingestion.sink.file import write_metadata_file
 from datahub.ingestion.source.file import read_metadata_file
 from datahub.testing.mcp_diff import CannotCompareMCPs, MCPDiff, get_aspects_by_urn
+from datahub.testing.pytest_hooks import get_golden_settings
 
 logger = logging.getLogger(__name__)
 
@@ -40,26 +41,26 @@ def load_json_file(filename: Union[str, os.PathLike]) -> MetadataJson:
 def assert_metadata_files_equal(
     output_path: Union[str, os.PathLike],
     golden_path: Union[str, os.PathLike],
-    update_golden: bool,
-    copy_output: bool,
     ignore_paths: Sequence[str] = (),
     ignore_paths_v2: Sequence[str] = (),
     ignore_order: bool = True,
 ) -> None:
+    settings = get_golden_settings()
+
     golden_exists = os.path.isfile(golden_path)
 
-    if copy_output:
+    if settings.copy_output:
         shutil.copyfile(str(output_path), str(golden_path) + ".output")
         logger.info(f"Copied output file to {golden_path}.output")
 
-    if not update_golden and not golden_exists:
+    if not settings.update_golden and not golden_exists:
         raise FileNotFoundError(
             "Golden file does not exist. Please run with the --update-golden-files option to create."
         )
 
     output = load_json_file(output_path)
 
-    if update_golden and not golden_exists:
+    if settings.update_golden and not golden_exists:
         shutil.copyfile(str(output_path), str(golden_path))
         return
     else:
@@ -87,7 +88,7 @@ def assert_metadata_files_equal(
     ignore_paths = (*ignore_paths, *default_exclude_paths)
 
     diff = diff_metadata_json(output, golden, ignore_paths, ignore_order=ignore_order)
-    if diff and update_golden:
+    if diff and settings.update_golden:
         if isinstance(diff, MCPDiff) and diff.is_delta_valid:
             logger.info(f"Applying delta to golden file {golden_path}")
             diff.apply_delta(golden)

datahub/testing/pytest_hooks.py

@@ -0,0 +1,56 @@
+import dataclasses
+from typing import Optional
+
+import pytest
+
+__all__ = [
+    "load_golden_flags",
+    "get_golden_settings",
+    "pytest_addoption",
+    "GoldenFileSettings",
+]
+
+
+@dataclasses.dataclass
+class GoldenFileSettings:
+    update_golden: bool
+    copy_output: bool
+
+
+_registered: bool = False
+_settings: Optional[GoldenFileSettings] = None
+
+
+def pytest_addoption(parser: pytest.Parser) -> None:
+    parser.addoption(
+        "--update-golden-files",
+        action="store_true",
+        default=False,
+    )
+
+    # TODO: Deprecate and remove this flag.
+    parser.addoption("--copy-output-files", action="store_true", default=False)
+
+    global _registered
+    _registered = True
+
+
+@pytest.fixture(scope="session", autouse=True)
+def load_golden_flags(pytestconfig: pytest.Config) -> None:
+    global _settings
+    _settings = GoldenFileSettings(
+        update_golden=pytestconfig.getoption("--update-golden-files"),
+        copy_output=pytestconfig.getoption("--copy-output-files"),
+    )
+
+
+def get_golden_settings() -> GoldenFileSettings:
+    if not _registered:
+        raise ValueError(
+            "Golden files aren't set up properly. Call register_golden_flags from a conftest pytest_addoptions method."
+        )
+    if not _settings:
+        raise ValueError(
+            "Golden files aren't set up properly. Ensure load_golden_flags is imported in your conftest."
+        )
+    return _settings

datahub/upgrade/upgrade.py

@@ -293,9 +293,9 @@ def is_client_server_compatible(client: VersionStats, server: VersionStats) -> i
         return server.version.micro - client.version.micro
 
 
-def _maybe_print_upgrade_message(  # noqa: C901
+def _maybe_print_upgrade_message(
     version_stats: Optional[DataHubVersionStats],
-) -> None:  # noqa: C901
+) -> None:
     days_before_cli_stale = 7
     days_before_quickstart_stale = 7
 

datahub/utilities/file_backed_collections.py

@@ -10,13 +10,11 @@ import tempfile
 import threading
 from dataclasses import dataclass, field
 from datetime import datetime
-from enum import Enum
 from types import TracebackType
 from typing import (
     Any,
     Callable,
     Dict,
-    Final,
     Generic,
     Iterator,
     List,
@@ -31,6 +29,7 @@ from typing import (
 )
 
 from datahub.ingestion.api.closeable import Closeable
+from datahub.utilities.sentinels import Unset, unset
 
 logger: logging.Logger = logging.getLogger(__name__)
 
@@ -59,16 +58,6 @@ SqliteValue = Union[int, float, str, bytes, datetime, None]
 _VT = TypeVar("_VT")
 
 
-class Unset(Enum):
-    token = 0
-
-
-# It's pretty annoying to create a true sentinel that works with typing.
-# https://peps.python.org/pep-0484/#support-for-singleton-types-in-unions
-# Can't wait for https://peps.python.org/pep-0661/
-_unset: Final = Unset.token
-
-
 class ConnectionWrapper:
     """
     Wraps a SQlite connection, allowing connection reuse across multiple FileBacked* objects.
@@ -372,7 +361,7 @@ class FileBackedDict(MutableMapping[str, _VT], Closeable, Generic[_VT]):
         self,
         /,
         key: str,
-        default: Union[_VT, Unset] = _unset,
+        default: Union[_VT, Unset] = unset,
     ) -> _VT:
         # If key is in the dictionary, this is similar to __getitem__ + mark_dirty.
         # If key is not in the dictionary, this is similar to __setitem__.
@@ -383,7 +372,7 @@ class FileBackedDict(MutableMapping[str, _VT], Closeable, Generic[_VT]):
             self.mark_dirty(key)
             return value
         except KeyError:
-            if default is _unset:
+            if default is unset:
                 raise
 
             self[key] = default