acryl-datahub 0.14.1.13rc3__py3-none-any.whl → 0.14.1.13rc5__py3-none-any.whl

datahub/sql_parsing/datajob.py

@@ -0,0 +1,50 @@
+import logging
+from typing import Iterable, List, Optional
+
+from datahub.emitter.mcp import MetadataChangeProposalWrapper
+from datahub.metadata.schema_classes import (
+    DataJobInputOutputClass,
+    FineGrainedLineageClass,
+    UpstreamLineageClass,
+)
+
+logger = logging.getLogger(__name__)
+
+
+def to_datajob_input_output(
+    *, mcps: Iterable[MetadataChangeProposalWrapper], ignore_extra_mcps: bool = True
+) -> Optional[DataJobInputOutputClass]:
+    inputDatasets: List[str] = []
+    outputDatasets: List[str] = []
+    fineGrainedLineages: List[FineGrainedLineageClass] = []
+    for mcp in mcps:
+        # TODO: Represent simple write operations without lineage as outputDatasets.
+
+        upstream_lineage = mcp.as_workunit().get_aspect_of_type(UpstreamLineageClass)
+        if upstream_lineage is not None:
+            if mcp.entityUrn and mcp.entityUrn not in outputDatasets:
+                outputDatasets.append(mcp.entityUrn)
+
+            for upstream in upstream_lineage.upstreams:
+                if upstream.dataset not in inputDatasets:
+                    inputDatasets.append(upstream.dataset)
+
+            if upstream_lineage.fineGrainedLineages:
+                for fineGrainedLineage in upstream_lineage.fineGrainedLineages:
+                    fineGrainedLineages.append(fineGrainedLineage)
+
+        elif ignore_extra_mcps:
+            pass
+        else:
+            raise ValueError(
+                f"Expected an upstreamLineage aspect, got {mcp.aspectName} for {mcp.entityUrn}"
+            )
+
+    if not inputDatasets and not outputDatasets:
+        return None
+
+    return DataJobInputOutputClass(
+        inputDatasets=inputDatasets,
+        outputDatasets=outputDatasets,
+        fineGrainedLineages=fineGrainedLineages,
+    )

datahub/sql_parsing/query_types.py

@@ -14,7 +14,16 @@ def _is_temp_table(table: sqlglot.exp.Table, dialect: sqlglot.Dialect) -> bool:
     identifier: sqlglot.exp.Identifier = table.this
 
     return identifier.args.get("temporary") or (
-        is_dialect_instance(dialect, "redshift") and identifier.name.startswith("#")
+        # These dialects use # as a prefix for temp tables.
+        is_dialect_instance(
+            dialect,
+            [
+                "redshift",
+                "mssql",
+                # sybase is another one, but we don't support that dialect yet.
+            ],
+        )
+        and identifier.name.startswith("#")
     )
 
 

datahub/sql_parsing/split_statements.py

@@ -0,0 +1,163 @@
+import re
+from enum import Enum
+from typing import Generator, List, Tuple
+
+CONTROL_FLOW_KEYWORDS = [
+    "GO",
+    r"BEGIN\w+TRY",
+    r"BEGIN\w+CATCH",
+    "BEGIN",
+    r"END\w+TRY",
+    r"END\w+CATCH",
+    "END",
+]
+
+# There's an exception to this rule, which is when the statement
+# is preceeded by a CTE.
+FORCE_NEW_STATEMENT_KEYWORDS = [
+    # SELECT is used inside queries as well, so we can't include it here.
+    "INSERT",
+    "UPDATE",
+    "DELETE",
+    "MERGE",
+]
+
+
+class ParserState(Enum):
+    NORMAL = 1
+    STRING = 2
+    COMMENT = 3
+    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
+
+    # 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"^{re.escape(keyword)}\b"
+    match = re.match(pattern, sql[pos:], re.IGNORECASE)
+    return bool(match)
+
+
+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.
+    """
+
+    for keyword in keywords:
+        if _is_keyword_at_position(sql, pos, keyword):
+            return True, keyword, len(keyword)
+    return False, "", 0
+
+
+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
+
+    current_statement: List[str] = []
+    state = ParserState.NORMAL
+    i = 0
+
+    def yield_if_complete() -> Generator[str, None, None]:
+        statement = "".join(current_statement).strip()
+        if 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
+
+                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()
+                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()

datahub/sql_parsing/sql_parsing_aggregator.py

@@ -762,7 +762,6 @@ class SqlParsingAggregator(Closeable):
 
         This assumes that queries come in order of increasing timestamps.
         """
-
         self.report.num_observed_queries += 1
 
         # All queries with no session ID are assumed to be part of the same session.

datahub/sql_parsing/sql_parsing_common.py

@@ -21,6 +21,9 @@ DIALECTS_WITH_CASE_INSENSITIVE_COLS = {
     # See more below:
     # https://documentation.sas.com/doc/en/pgmsascdc/9.4_3.5/acreldb/n0ejgx4895bofnn14rlguktfx5r3.htm
     "teradata",
+    # For SQL server, the default collation rules mean that all identifiers (schema, table, column names)
+    # are case preserving but case insensitive.
+    "mssql",
 }
 DIALECTS_WITH_DEFAULT_UPPERCASE_COLS = {
     # In some dialects, column identifiers are effectively case insensitive
@@ -28,6 +31,9 @@ DIALECTS_WITH_DEFAULT_UPPERCASE_COLS = {
     # automatically lowercase unquoted identifiers.
     "snowflake",
 }
+assert DIALECTS_WITH_DEFAULT_UPPERCASE_COLS.issubset(
+    DIALECTS_WITH_CASE_INSENSITIVE_COLS
+)
 
 
 class QueryType(enum.Enum):

datahub/sql_parsing/sqlglot_lineage.py

@@ -5,7 +5,7 @@ import functools
 import logging
 import traceback
 from collections import defaultdict
-from typing import Any, Dict, List, Optional, Set, Tuple, Union
+from typing import Any, Dict, List, Optional, Set, Tuple, TypeVar, Union
 
 import pydantic.dataclasses
 import sqlglot
@@ -873,6 +873,49 @@ def _translate_internal_column_lineage(
     )
 
 
+_StrOrNone = TypeVar("_StrOrNone", str, Optional[str])
+
+
+def _normalize_db_or_schema(
+    db_or_schema: _StrOrNone,
+    dialect: sqlglot.Dialect,
+) -> _StrOrNone:
+    if db_or_schema is None:
+        return None
+
+    # In snowflake, table identifiers must be uppercased to match sqlglot's behavior.
+    if is_dialect_instance(dialect, "snowflake"):
+        return db_or_schema.upper()
+
+    # In mssql, table identifiers must be lowercased.
+    elif is_dialect_instance(dialect, "mssql"):
+        return db_or_schema.lower()
+
+    return db_or_schema
+
+
+def _simplify_select_into(statement: sqlglot.exp.Expression) -> sqlglot.exp.Expression:
+    """
+    Check if the expression is a SELECT INTO statement. If so, converts it into a CTAS.
+    Other expressions are returned as-is.
+    """
+
+    if not (isinstance(statement, sqlglot.exp.Select) and statement.args.get("into")):
+        return statement
+
+    # Convert from SELECT <cols> INTO <out> <expr>
+    # to CREATE TABLE <out> AS SELECT <cols> <expr>
+    into_expr: sqlglot.exp.Into = statement.args["into"].pop()
+    into_table = into_expr.this
+
+    create = sqlglot.exp.Create(
+        this=into_table,
+        kind="TABLE",
+        expression=statement,
+    )
+    return create
+
+
 def _sqlglot_lineage_inner(
     sql: sqlglot.exp.ExpOrStr,
     schema_resolver: SchemaResolverInterface,
@@ -885,12 +928,9 @@ def _sqlglot_lineage_inner(
     else:
         dialect = get_dialect(default_dialect)
 
-    if is_dialect_instance(dialect, "snowflake"):
-        # in snowflake, table identifiers must be uppercased to match sqlglot's behavior.
-        if default_db:
-            default_db = default_db.upper()
-        if default_schema:
-            default_schema = default_schema.upper()
+    default_db = _normalize_db_or_schema(default_db, dialect)
+    default_schema = _normalize_db_or_schema(default_schema, dialect)
+
     if is_dialect_instance(dialect, "redshift") and not default_schema:
         # On Redshift, there's no "USE SCHEMA <schema>" command. The default schema
         # is public, and "current schema" is the one at the front of the search path.
@@ -918,6 +958,8 @@ def _sqlglot_lineage_inner(
     #     original_statement.sql(pretty=True, dialect=dialect),
     # )
 
+    statement = _simplify_select_into(statement)
+
     # Make sure the tables are resolved with the default db / schema.
     # This only works for Unionable statements. For other types of statements,
     # we have to do it manually afterwards, but that's slightly lower accuracy

datahub/sql_parsing/sqlglot_utils.py

@@ -61,7 +61,7 @@ def is_dialect_instance(
     else:
         platforms = list(platforms)
 
-    dialects = [sqlglot.Dialect.get_or_raise(platform) for platform in platforms]
+    dialects = [get_dialect(platform) for platform in platforms]
 
     if any(isinstance(dialect, dialect_class.__class__) for dialect_class in dialects):
         return True

datahub/utilities/file_backed_collections.py

@@ -228,6 +228,12 @@ class FileBackedDict(MutableMapping[str, _VT], Closeable, Generic[_VT]):
         else:
             self._conn = ConnectionWrapper()
 
+        if sqlite3.sqlite_version_info < (3, 24, 0):
+            # We use the ON CONFLICT clause to implement UPSERTs with sqlite.
+            # This was added in 3.24.0 from 2018-06-04.
+            # See https://www.sqlite.org/lang_conflict.html
+            raise RuntimeError("SQLite version 3.24.0 or later is required")
+
         # We keep a small cache in memory to avoid having to serialize/deserialize
         # data from the database too often. We use an OrderedDict to build
         # a poor-man's LRU cache.