acryl-datahub 0.14.1.13rc4__py3-none-any.whl → 0.14.1.13rc6__py3-none-any.whl

datahub/ingestion/source/sql/teradata.py

@@ -44,7 +44,7 @@ from datahub.ingestion.api.workunit import MetadataWorkUnit
 from datahub.ingestion.graph.client import DataHubGraph
 from datahub.ingestion.source.sql.sql_common import SqlWorkUnit, register_custom_type
 from datahub.ingestion.source.sql.sql_config import SQLCommonConfig
-from datahub.ingestion.source.sql.sql_generic_profiler import ProfilingSqlReport
+from datahub.ingestion.source.sql.sql_report import SQLSourceReport
 from datahub.ingestion.source.sql.two_tier_sql_source import (
     TwoTierSQLAlchemyConfig,
     TwoTierSQLAlchemySource,
@@ -330,7 +330,7 @@ def optimized_get_view_definition(
 
 
 @dataclass
-class TeradataReport(ProfilingSqlReport, IngestionStageReport, BaseTimeWindowReport):
+class TeradataReport(SQLSourceReport, IngestionStageReport, BaseTimeWindowReport):
     num_queries_parsed: int = 0
     num_view_ddl_parsed: int = 0
     num_table_parse_failures: int = 0

datahub/ingestion/source/sql/vertica.py

@@ -27,7 +27,6 @@ from datahub.ingestion.api.workunit import MetadataWorkUnit
 from datahub.ingestion.source.common.data_reader import DataReader
 from datahub.ingestion.source.sql.sql_common import (
     SQLAlchemySource,
-    SQLSourceReport,
     SqlWorkUnit,
     get_schema_metadata,
 )
@@ -35,6 +34,7 @@ from datahub.ingestion.source.sql.sql_config import (
     BasicSQLAlchemyConfig,
     SQLCommonConfig,
 )
+from datahub.ingestion.source.sql.sql_report import SQLSourceReport
 from datahub.ingestion.source.sql.sql_utils import get_domain_wu
 from datahub.metadata.com.linkedin.pegasus2avro.common import StatusClass
 from datahub.metadata.com.linkedin.pegasus2avro.dataset import UpstreamLineage
@@ -536,7 +536,7 @@ class VerticaSource(SQLAlchemySource):
             )
 
             if not self.is_dataset_eligible_for_profiling(
-                dataset_name, sql_config, inspector, profile_candidates
+                dataset_name, schema, inspector, profile_candidates
             ):
                 if self.config.profiling.report_dropped_profiles:
                     self.report.report_dropped(f"profile of {dataset_name}")

datahub/ingestion/source/unity/report.py

@@ -2,7 +2,7 @@ from dataclasses import dataclass, field
 from typing import Optional, Tuple
 
 from datahub.ingestion.api.report import EntityFilterReport, Report
-from datahub.ingestion.source.sql.sql_generic_profiler import ProfilingSqlReport
+from datahub.ingestion.source.sql.sql_report import SQLSourceReport
 from datahub.ingestion.source_report.ingestion_stage import IngestionStageReport
 from datahub.utilities.lossy_collections import LossyDict, LossyList
 from datahub.utilities.perf_timer import PerfTimer
@@ -19,7 +19,7 @@ class UnityCatalogUsagePerfReport(Report):
 
 
 @dataclass
-class UnityCatalogReport(IngestionStageReport, ProfilingSqlReport):
+class UnityCatalogReport(IngestionStageReport, SQLSourceReport):
     metastores: EntityFilterReport = EntityFilterReport.field(type="metastore")
     catalogs: EntityFilterReport = EntityFilterReport.field(type="catalog")
     schemas: EntityFilterReport = EntityFilterReport.field(type="schema")

datahub/metadata/schema.avsc

@@ -6005,7 +6005,7 @@
                           "fields": [
                             {
                               "Searchable": {
-                                "boostScore": 5.0,
+                                "boostScore": 1.0,
                                 "fieldName": "fieldPaths",
                                 "fieldType": "TEXT",
                                 "queryByDefault": "true"

datahub/metadata/schemas/AssertionInfo.avsc

@@ -1542,7 +1542,7 @@
                         "fields": [
                           {
                             "Searchable": {
-                              "boostScore": 5.0,
+                              "boostScore": 1.0,
                               "fieldName": "fieldPaths",
                               "fieldType": "TEXT",
                               "queryByDefault": "true"

datahub/metadata/schemas/InputFields.avsc

@@ -42,7 +42,7 @@
                   "fields": [
                     {
                       "Searchable": {
-                        "boostScore": 5.0,
+                        "boostScore": 1.0,
                         "fieldName": "fieldPaths",
                         "fieldType": "TEXT",
                         "queryByDefault": "true"

datahub/metadata/schemas/MetadataChangeEvent.avsc

@@ -4013,7 +4013,7 @@
                             "fields": [
                               {
                                 "Searchable": {
-                                  "boostScore": 5.0,
+                                  "boostScore": 1.0,
                                   "fieldName": "fieldPaths",
                                   "fieldType": "TEXT",
                                   "queryByDefault": "true"

datahub/metadata/schemas/SchemaMetadata.avsc

@@ -309,7 +309,7 @@
           "fields": [
             {
               "Searchable": {
-                "boostScore": 5.0,
+                "boostScore": 1.0,
                 "fieldName": "fieldPaths",
                 "fieldType": "TEXT",
                 "queryByDefault": "true"

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.