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

{acryl_datahub-0.14.1.13rc4.dist-info → acryl_datahub-0.14.1.13rc5.dist-info}/RECORD

@@ -1,4 +1,4 @@
-datahub/__init__.py,sha256=WMIoTkprsOt2SjEA2gbvdoJ3VkTIeWCVPilQAqNdCQw,577
+datahub/__init__.py,sha256=sacWO6qm2tPLBhc25GFoFnt_AeiT0Qk9ZiibpHJbPhQ,577
 datahub/__main__.py,sha256=pegIvQ9hzK7IhqVeUi1MeADSZ2QlP-D3K0OQdEg55RU,106
 datahub/entrypoints.py,sha256=3-qSfXAx3Z0FEkBV5tlO8fQr4xk4ySeDRMVTpS5Xd6A,7793
 datahub/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
@@ -449,7 +449,7 @@ datahub/ingestion/source/sql/mysql.py,sha256=nDWK4YbqomcJgnit9b8geUGrp_3eix4bt0_
 datahub/ingestion/source/sql/oracle.py,sha256=exJJvMSC42Oj0IiFtpAPFatWkRu9mVu8sd7ofOdanqU,22460
 datahub/ingestion/source/sql/postgres.py,sha256=uC1kYEI8VdxiZ1Y9IxMWzwmg11wtMqYN0e2fkok1rxo,11972
 datahub/ingestion/source/sql/presto.py,sha256=PB-CS5MX2dSRFRHjlxfkLHGXLZXFNCsVAAyRBtY6HMg,3611
-datahub/ingestion/source/sql/sql_common.py,sha256=SZkPK7aUiTdMHf5alG5724i1pyAn1jFUjcuDkKUoeL0,51813
+datahub/ingestion/source/sql/sql_common.py,sha256=Qe481Pct1vfl-NVpgOufSEaNSMgJGD3bq1DPvjIkPTg,52164
 datahub/ingestion/source/sql/sql_config.py,sha256=M-l_uXau0ODolLZHBzAXhy-Rq5yYxvJ6cLbCIea7Mww,9449
 datahub/ingestion/source/sql/sql_generic.py,sha256=9AERvkK8kdJUeDOzCYJDb93xdv6Z4DGho0NfeHj5Uyg,2740
 datahub/ingestion/source/sql/sql_generic_profiler.py,sha256=vWna43rv9ClfK9MBNTx8tdoUr35nD5d8ppgeamEEhSQ,12528
@@ -462,8 +462,9 @@ datahub/ingestion/source/sql/trino.py,sha256=FEn_BQ3pm23hKx94ek5kk5IXGNYcBqZEhll
 datahub/ingestion/source/sql/two_tier_sql_source.py,sha256=YDrGBb5WKVls6qv17QU5foKrf71SydzEltc3WsVAhQc,5732
 datahub/ingestion/source/sql/vertica.py,sha256=pkx-1JDBmow7WhoEIEh0SLj7kff9L0zUOHDytoC43gk,33339
 datahub/ingestion/source/sql/mssql/__init__.py,sha256=1agpl8S_uDW40olkhCX_W19dbr5GO9qgjS3R7pLRZSk,87
-datahub/ingestion/source/sql/mssql/job_models.py,sha256=-wDCltoBAFHaydAV4Gs0vm-j7qKNCH5iEJOVWanhlZw,5981
-datahub/ingestion/source/sql/mssql/source.py,sha256=_qenSWn5mPrr5EPDB6Qbt9EOURqesK68Fm_pnh9sO58,26594
+datahub/ingestion/source/sql/mssql/job_models.py,sha256=eMyR0Efl5kvi7QNgNXzd5_6PdDKYly_552Y8OGSj9PY,6012
+datahub/ingestion/source/sql/mssql/source.py,sha256=fzpWjwexGvJgpd6Z4DCsK6Ld2vQCfPkD2M1xE4pU9Ec,29542
+datahub/ingestion/source/sql/mssql/stored_procedure_lineage.py,sha256=RpnvKPalAAaOD_eUg8bZ4VkGTSeLFWuy0mefwc4s3x8,2837
 datahub/ingestion/source/state/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 datahub/ingestion/source/state/checkpoint.py,sha256=x9Xww-MIFXSKjeg1tOZXE72LehCm5OfKy3HfucgIRWM,8833
 datahub/ingestion/source/state/entity_removal_state.py,sha256=zvIsmYg7oiIu2FhecU0VfLBNToUqvKoKyDeiFfkOcyc,6611
@@ -858,9 +859,11 @@ datahub/specific/structured_property.py,sha256=IYeFyafPidNrDbn1sU65rEPwIZDS-wLY1
 datahub/sql_parsing/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 datahub/sql_parsing/_models.py,sha256=il-xm1RcLdi1phJUV3xrTecdOGH31akqheuSC2N4YhQ,3141
 datahub/sql_parsing/_sqlglot_patch.py,sha256=iYJ8zOThHqqbamD5jdNr9iHTWD7ewNeHzPiTb6-rO3Y,7043
-datahub/sql_parsing/query_types.py,sha256=AxoWZiLuXM6iZ2AXdSlD82gqPNr1ZFh78wn2MSTu2wY,2479
+datahub/sql_parsing/datajob.py,sha256=1X8KpEk-y3_8xJuA_Po27EHZgOcxK9QADI6Om9gSGn0,1751
+datahub/sql_parsing/query_types.py,sha256=FKjDzszZzsrCfYfm7dgD6T_8865qxWl767fdGyHWBh4,2720
 datahub/sql_parsing/schema_resolver.py,sha256=9wbJT80K4nsIHHOuQLso9QoFQAYwfSJZnqHwsaU3UTY,10197
-datahub/sql_parsing/sql_parsing_aggregator.py,sha256=YzYmWNQjXuQYcfh91ZnqjuzCDWdmZpE-ZX14Kz6XFMs,69938
+datahub/sql_parsing/split_statements.py,sha256=uZhAXLaRxDfmK0lPBW2oM_YVdJfSMhdgndnfd9iIXuA,5001
+datahub/sql_parsing/sql_parsing_aggregator.py,sha256=gLelf5l73EufB8qijb9ZDLANkt4o05schGg4DY-bOJs,69937
 datahub/sql_parsing/sql_parsing_common.py,sha256=h_V_m54hJ9EUh5kczq7cYOIeNeo4bgf0Px0H-Nq-UIg,2602
 datahub/sql_parsing/sql_parsing_result_utils.py,sha256=prwWTj1EB2fRPv1eMB4EkpFNafIYAt-X8TIK0NWqank,796
 datahub/sql_parsing/sqlglot_lineage.py,sha256=K532_zvj7vEt2Ci2d3kxYZLZkq4S0ECBRZ2nHuY11a4,45927
@@ -967,8 +970,8 @@ datahub_provider/operators/datahub_assertion_operator.py,sha256=uvTQ-jk2F0sbqqxp
 datahub_provider/operators/datahub_assertion_sensor.py,sha256=lCBj_3x1cf5GMNpHdfkpHuyHfVxsm6ff5x2Z5iizcAo,140
 datahub_provider/operators/datahub_operation_operator.py,sha256=aevDp2FzX7FxGlXrR0khoHNbxbhKR2qPEX5e8O2Jyzw,174
 datahub_provider/operators/datahub_operation_sensor.py,sha256=8fcdVBCEPgqy1etTXgLoiHoJrRt_nzFZQMdSzHqSG7M,168
-acryl_datahub-0.14.1.13rc4.dist-info/METADATA,sha256=pKvm7N_prnqdEylMVyCtB6aReDQm2T-0LmvkhiwUdrU,171138
-acryl_datahub-0.14.1.13rc4.dist-info/WHEEL,sha256=PZUExdf71Ui_so67QXpySuHtCi3-J3wvF4ORK6k_S8U,91
-acryl_datahub-0.14.1.13rc4.dist-info/entry_points.txt,sha256=VcQx0dnqaYLyeY_L5OaX7bLmmE-Il7TAXkxCKvEn2bA,9432
-acryl_datahub-0.14.1.13rc4.dist-info/top_level.txt,sha256=iLjSrLK5ox1YVYcglRUkcvfZPvKlobBWx7CTUXx8_GI,25
-acryl_datahub-0.14.1.13rc4.dist-info/RECORD,,
+acryl_datahub-0.14.1.13rc5.dist-info/METADATA,sha256=9ZZBowQ_PYKf6d4V6EHqTTlJDY77BqxSX8uKfiIqFT4,171138
+acryl_datahub-0.14.1.13rc5.dist-info/WHEEL,sha256=PZUExdf71Ui_so67QXpySuHtCi3-J3wvF4ORK6k_S8U,91
+acryl_datahub-0.14.1.13rc5.dist-info/entry_points.txt,sha256=VcQx0dnqaYLyeY_L5OaX7bLmmE-Il7TAXkxCKvEn2bA,9432
+acryl_datahub-0.14.1.13rc5.dist-info/top_level.txt,sha256=iLjSrLK5ox1YVYcglRUkcvfZPvKlobBWx7CTUXx8_GI,25
+acryl_datahub-0.14.1.13rc5.dist-info/RECORD,,

datahub/__init__.py

@@ -3,7 +3,7 @@ import warnings
 
 # Published at https://pypi.org/project/acryl-datahub/.
 __package_name__ = "acryl-datahub"
-__version__ = "0.14.1.13rc4"
+__version__ = "0.14.1.13rc5"
 
 
 def is_dev_mode() -> bool:

datahub/ingestion/source/sql/mssql/job_models.py

@@ -101,6 +101,7 @@ class StoredProcedure:
     flow: Union[MSSQLJob, MSSQLProceduresContainer]
     type: str = "STORED_PROCEDURE"
     source: str = "mssql"
+    code: Optional[str] = None
 
     @property
     def full_type(self) -> str:

datahub/ingestion/source/sql/mssql/source.py

@@ -24,6 +24,8 @@ from datahub.ingestion.api.decorators import (
     platform_name,
     support_status,
 )
+from datahub.ingestion.api.source import StructuredLogLevel
+from datahub.ingestion.api.source_helpers import auto_workunit
 from datahub.ingestion.api.workunit import MetadataWorkUnit
 from datahub.ingestion.source.sql.mssql.job_models import (
     JobStep,
@@ -36,6 +38,9 @@ from datahub.ingestion.source.sql.mssql.job_models import (
     ProcedureParameter,
     StoredProcedure,
 )
+from datahub.ingestion.source.sql.mssql.stored_procedure_lineage import (
+    generate_procedure_lineage,
+)
 from datahub.ingestion.source.sql.sql_common import (
     SQLAlchemySource,
     SqlWorkUnit,
@@ -51,6 +56,7 @@ from datahub.metadata.schema_classes import (
     StringTypeClass,
     UnionTypeClass,
 )
+from datahub.utilities.file_backed_collections import FileBackedList
 
 logger: logging.Logger = logging.getLogger(__name__)
 
@@ -99,6 +105,10 @@ class SQLServerConfig(BasicSQLAlchemyConfig):
         default=False,
         description="Enable to convert the SQL Server assets urns to lowercase",
     )
+    include_lineage: bool = Field(
+        default=True,
+        description="Enable lineage extraction for stored procedures",
+    )
 
     @pydantic.validator("uri_args")
     def passwords_match(cls, v, values, **kwargs):
@@ -161,6 +171,7 @@ class SQLServerSource(SQLAlchemySource):
         self.current_database = None
         self.table_descriptions: Dict[str, str] = {}
         self.column_descriptions: Dict[str, str] = {}
+        self.stored_procedures: FileBackedList[StoredProcedure] = FileBackedList()
         if self.config.include_descriptions:
             for inspector in self.get_inspectors():
                 db_name: str = self.get_db_name(inspector)
@@ -374,7 +385,7 @@ class SQLServerSource(SQLAlchemySource):
     def loop_job_steps(
         self, job: MSSQLJob, job_steps: Dict[str, Any]
     ) -> Iterable[MetadataWorkUnit]:
-        for step_id, step_data in job_steps.items():
+        for _step_id, step_data in job_steps.items():
             step = JobStep(
                 job_name=job.formatted_name,
                 step_name=step_data["step_name"],
@@ -412,37 +423,44 @@ class SQLServerSource(SQLAlchemySource):
             if procedures:
                 yield from self.construct_flow_workunits(data_flow=data_flow)
             for procedure in procedures:
-                upstream = self._get_procedure_upstream(conn, procedure)
-                downstream = self._get_procedure_downstream(conn, procedure)
-                data_job = MSSQLDataJob(
-                    entity=procedure,
-                )
-                # TODO: because of this upstream and downstream are more dependencies,
-                #  can't be used as DataJobInputOutput.
-                #  Should be reorganized into lineage.
-                data_job.add_property("procedure_depends_on", str(upstream.as_property))
-                data_job.add_property(
-                    "depending_on_procedure", str(downstream.as_property)
-                )
-                procedure_definition, procedure_code = self._get_procedure_code(
-                    conn, procedure
-                )
-                if procedure_definition:
-                    data_job.add_property("definition", procedure_definition)
-                if sql_config.include_stored_procedures_code and procedure_code:
-                    data_job.add_property("code", procedure_code)
-                procedure_inputs = self._get_procedure_inputs(conn, procedure)
-                properties = self._get_procedure_properties(conn, procedure)
-                data_job.add_property(
-                    "input parameters", str([param.name for param in procedure_inputs])
-                )
-                for param in procedure_inputs:
-                    data_job.add_property(
-                        f"parameter {param.name}", str(param.properties)
-                    )
-                for property_name, property_value in properties.items():
-                    data_job.add_property(property_name, str(property_value))
-                yield from self.construct_job_workunits(data_job)
+                yield from self._process_stored_procedure(conn, procedure)
+
+    def _process_stored_procedure(
+        self, conn: Connection, procedure: StoredProcedure
+    ) -> Iterable[MetadataWorkUnit]:
+        upstream = self._get_procedure_upstream(conn, procedure)
+        downstream = self._get_procedure_downstream(conn, procedure)
+        data_job = MSSQLDataJob(
+            entity=procedure,
+        )
+        # TODO: because of this upstream and downstream are more dependencies,
+        #  can't be used as DataJobInputOutput.
+        #  Should be reorganized into lineage.
+        data_job.add_property("procedure_depends_on", str(upstream.as_property))
+        data_job.add_property("depending_on_procedure", str(downstream.as_property))
+        procedure_definition, procedure_code = self._get_procedure_code(conn, procedure)
+        procedure.code = procedure_code
+        if procedure_definition:
+            data_job.add_property("definition", procedure_definition)
+        if procedure_code and self.config.include_stored_procedures_code:
+            data_job.add_property("code", procedure_code)
+        procedure_inputs = self._get_procedure_inputs(conn, procedure)
+        properties = self._get_procedure_properties(conn, procedure)
+        data_job.add_property(
+            "input parameters", str([param.name for param in procedure_inputs])
+        )
+        for param in procedure_inputs:
+            data_job.add_property(f"parameter {param.name}", str(param.properties))
+        for property_name, property_value in properties.items():
+            data_job.add_property(property_name, str(property_value))
+        if self.config.include_lineage:
+            # These will be used to construct lineage
+            self.stored_procedures.append(procedure)
+        yield from self.construct_job_workunits(
+            data_job,
+            # For stored procedure lineage is ingested later
+            include_lineage=False,
+        )
 
     @staticmethod
     def _get_procedure_downstream(
@@ -546,8 +564,8 @@ class SQLServerSource(SQLAlchemySource):
                 code_list.append(row["Text"])
                 if code_slice_text in re.sub(" +", " ", row["Text"].lower()).strip():
                     code_slice_index = index
-            definition = "\n".join(code_list[:code_slice_index])
-            code = "\n".join(code_list[code_slice_index:])
+            definition = "".join(code_list[:code_slice_index])
+            code = "".join(code_list[code_slice_index:])
         except ResourceClosedError:
             logger.warning(
                 "Connection was closed from procedure '%s'",
@@ -602,16 +620,18 @@ class SQLServerSource(SQLAlchemySource):
     def construct_job_workunits(
         self,
         data_job: MSSQLDataJob,
+        include_lineage: bool = True,
     ) -> Iterable[MetadataWorkUnit]:
         yield MetadataChangeProposalWrapper(
             entityUrn=data_job.urn,
             aspect=data_job.as_datajob_info_aspect,
         ).as_workunit()
 
-        yield MetadataChangeProposalWrapper(
-            entityUrn=data_job.urn,
-            aspect=data_job.as_datajob_input_output_aspect,
-        ).as_workunit()
+        if include_lineage:
+            yield MetadataChangeProposalWrapper(
+                entityUrn=data_job.urn,
+                aspect=data_job.as_datajob_input_output_aspect,
+            ).as_workunit()
         # TODO: Add SubType when it appear
 
     def construct_flow_workunits(
@@ -664,3 +684,58 @@ class SQLServerSource(SQLAlchemySource):
             if self.config.convert_urns_to_lowercase
             else qualified_table_name
         )
+
+    def get_workunits_internal(self) -> Iterable[MetadataWorkUnit]:
+        yield from super().get_workunits_internal()
+
+        # This is done at the end so that we will have access to tables
+        # from all databases in schema_resolver and discovered_tables
+        for procedure in self.stored_procedures:
+            with self.report.report_exc(
+                message="Failed to parse stored procedure lineage",
+                context=procedure.full_name,
+                level=StructuredLogLevel.WARN,
+            ):
+                yield from auto_workunit(
+                    generate_procedure_lineage(
+                        schema_resolver=self.schema_resolver,
+                        procedure=procedure,
+                        procedure_job_urn=MSSQLDataJob(entity=procedure).urn,
+                        is_temp_table=self.is_temp_table,
+                    )
+                )
+
+    def is_temp_table(self, name: str) -> bool:
+        try:
+            parts = name.split(".")
+            table_name = parts[-1]
+            schema_name = parts[-2]
+            db_name = parts[-3]
+
+            if table_name.startswith("#"):
+                return True
+
+            # This is also a temp table if
+            #   1. this name would be allowed by the dataset patterns, and
+            #   2. we have a list of discovered tables, and
+            #   3. it's not in the discovered tables list
+            if (
+                self.config.database_pattern.allowed(db_name)
+                and self.config.schema_pattern.allowed(schema_name)
+                and self.config.table_pattern.allowed(name)
+                and self.standardize_identifier_case(name)
+                not in self.discovered_datasets
+            ):
+                logger.debug(f"inferred as temp table {name}")
+                return True
+
+        except Exception:
+            logger.warning(f"Error parsing table name {name} ")
+        return False
+
+    def standardize_identifier_case(self, table_ref_str: str) -> str:
+        return (
+            table_ref_str.lower()
+            if self.config.convert_urns_to_lowercase
+            else table_ref_str
+        )

datahub/ingestion/source/sql/mssql/stored_procedure_lineage.py

@@ -0,0 +1,84 @@
+import logging
+from typing import Callable, Iterable, Optional
+
+from datahub.emitter.mcp import MetadataChangeProposalWrapper
+from datahub.ingestion.source.sql.mssql.job_models import StoredProcedure
+from datahub.metadata.schema_classes import DataJobInputOutputClass
+from datahub.sql_parsing.datajob import to_datajob_input_output
+from datahub.sql_parsing.schema_resolver import SchemaResolver
+from datahub.sql_parsing.split_statements import split_statements
+from datahub.sql_parsing.sql_parsing_aggregator import (
+    ObservedQuery,
+    SqlParsingAggregator,
+)
+
+logger = logging.getLogger(__name__)
+
+
+def parse_procedure_code(
+    *,
+    schema_resolver: SchemaResolver,
+    default_db: Optional[str],
+    default_schema: Optional[str],
+    code: str,
+    is_temp_table: Callable[[str], bool],
+    raise_: bool = False,
+) -> Optional[DataJobInputOutputClass]:
+    aggregator = SqlParsingAggregator(
+        platform=schema_resolver.platform,
+        env=schema_resolver.env,
+        schema_resolver=schema_resolver,
+        generate_lineage=True,
+        generate_queries=False,
+        generate_usage_statistics=False,
+        generate_operations=False,
+        generate_query_subject_fields=False,
+        generate_query_usage_statistics=False,
+        is_temp_table=is_temp_table,
+    )
+    for query in split_statements(code):
+        # TODO: We should take into account `USE x` statements.
+        aggregator.add_observed_query(
+            observed=ObservedQuery(
+                default_db=default_db,
+                default_schema=default_schema,
+                query=query,
+            )
+        )
+    if aggregator.report.num_observed_queries_failed and raise_:
+        logger.info(aggregator.report.as_string())
+        raise ValueError(
+            f"Failed to parse {aggregator.report.num_observed_queries_failed} queries."
+        )
+
+    mcps = list(aggregator.gen_metadata())
+    return to_datajob_input_output(
+        mcps=mcps,
+        ignore_extra_mcps=True,
+    )
+
+
+# Is procedure handling generic enough to be added to SqlParsingAggregator?
+def generate_procedure_lineage(
+    *,
+    schema_resolver: SchemaResolver,
+    procedure: StoredProcedure,
+    procedure_job_urn: str,
+    is_temp_table: Callable[[str], bool] = lambda _: False,
+    raise_: bool = False,
+) -> Iterable[MetadataChangeProposalWrapper]:
+    if procedure.code:
+        datajob_input_output = parse_procedure_code(
+            schema_resolver=schema_resolver,
+            default_db=procedure.db,
+            default_schema=procedure.schema,
+            code=procedure.code,
+            is_temp_table=is_temp_table,
+            raise_=raise_,
+        )
+
+        if datajob_input_output:
+            yield MetadataChangeProposalWrapper(
+                entityUrn=procedure_job_urn,
+                aspect=datajob_input_output,
+            )

datahub/ingestion/source/sql/sql_common.py

@@ -392,6 +392,7 @@ class SQLAlchemySource(StatefulIngestionSourceBase, TestableSource):
             platform_instance=self.config.platform_instance,
             env=self.config.env,
         )
+        self.discovered_datasets: Set[str] = set()
         self._view_definition_cache: MutableMapping[str, str]
         if self.config.use_file_backed_cache:
             self._view_definition_cache = FileBackedDict[str]()
@@ -831,8 +832,9 @@ class SQLAlchemySource(StatefulIngestionSourceBase, TestableSource):
         self._classify(dataset_name, schema, table, data_reader, schema_metadata)
 
         dataset_snapshot.aspects.append(schema_metadata)
-        if self.config.include_view_lineage:
+        if self._save_schema_to_resolver():
             self.schema_resolver.add_schema_metadata(dataset_urn, schema_metadata)
+            self.discovered_datasets.add(dataset_name)
         db_name = self.get_db_name(inspector)
 
         yield from self.add_table_to_schema_container(
@@ -1126,8 +1128,9 @@ class SQLAlchemySource(StatefulIngestionSourceBase, TestableSource):
                 columns,
                 canonical_schema=schema_fields,
             )
-            if self.config.include_view_lineage:
+            if self._save_schema_to_resolver():
                 self.schema_resolver.add_schema_metadata(dataset_urn, schema_metadata)
+                self.discovered_datasets.add(dataset_name)
         description, properties, _ = self.get_table_properties(inspector, schema, view)
         try:
             view_definition = inspector.get_view_definition(view, schema)
@@ -1190,6 +1193,11 @@ class SQLAlchemySource(StatefulIngestionSourceBase, TestableSource):
                 domain_registry=self.domain_registry,
             )
 
+    def _save_schema_to_resolver(self):
+        return self.config.include_view_lineage or (
+            hasattr(self.config, "include_lineage") and self.config.include_lineage
+        )
+
     def _run_sql_parser(
         self, view_identifier: str, query: str, schema_resolver: SchemaResolver
     ) -> Optional[SqlParsingResult]:

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.