Merge remote-tracking branch 'oss-datahub/master' into vertex_exp_4

Author: ryota-cloud

.github/scripts/docker_helpers.sh

@@ -24,7 +24,7 @@ function get_tag_full {
 }
 
 function get_python_docker_release_v {
-    echo $(echo ${GITHUB_REF} | sed -e "s,refs/heads/${MAIN_BRANCH},1!0.0.0+docker.${SHORT_SHA},g" -e 's,refs/tags/v\(.*\),1!\1+docker,g' -e 's,refs/pull/\([0-9]*\).*,1!0.0.0+docker.pr\1,g')
+    echo $(echo ${GITHUB_REF} | sed -e "s,refs/heads/${MAIN_BRANCH},1!0.0.0+docker.${SHORT_SHA},g" -e 's,refs/tags/v\([0-9a-zA-Z.]*\).*,1\!\1+docker,g' -e 's,refs/pull/\([0-9]*\).*,1!0.0.0+docker.pr\1,g')
 }
 
 function get_unique_tag {

docs-website/docusaurus.config.js

@@ -35,6 +35,10 @@ module.exports = {
     },
     {
       src: "https://app.revenuehero.io/scheduler.min.js"
+    },
+    {
+      src: "https://tag.clearbitscripts.com/v1/pk_2e321cabe30432a5c44c0424781aa35f/tags.js",
+      referrerPolicy: "strict-origin-when-cross-origin"
     }
   ],
   noIndex: isSaas,

metadata-ingestion/examples/ai/dh_ai_client.py

@@ -14,6 +14,7 @@
     ChangeTypeClass,
     DataProcessInstanceRunResultClass,
     DataProcessRunStatusClass,
+    EdgeClass,
 )
 from datahub.metadata.urns import (
     ContainerUrn,
@@ -255,7 +256,7 @@ def create_model(
         version_props = {
             "version": version_tag,
             "versionSet": str(version_set_urn),
-            "sortId": "AAAAAAAA",
+            "sortId": str(version_tag).zfill(10),
         }
 
         # Add alias if provided
@@ -266,22 +267,10 @@ def create_model(
             models.VersionPropertiesClass, version_props
         )
 
-        # Create version set properties
-        version_set_properties = models.VersionSetPropertiesClass(
-            latest=str(model_urn),
-            versioningScheme="ALPHANUMERIC_GENERATED_BY_DATAHUB",
-        )
-
         mcps = [
             self._create_mcp(
                 str(model_urn), properties, "mlModel", "mlModelProperties"
             ),
-            self._create_mcp(
-                str(version_set_urn),
-                version_set_properties,
-                "versionSet",
-                "versionSetProperties",
-            ),
             self._create_mcp(
                 str(model_urn), version_properties, "mlModel", "versionProperties"
             ),
@@ -429,7 +418,13 @@ def add_input_datasets_to_run(self, run_urn: str, dataset_urns: List[str]) -> No
             entity_urn=run_urn,
             entity_type="dataProcessInstance",
             aspect_name="dataProcessInstanceInput",
-            aspect=DataProcessInstanceInput(inputs=dataset_urns),
+            aspect=DataProcessInstanceInput(
+                inputs=[],
+                inputEdges=[
+                    EdgeClass(destinationUrn=str(dataset_urn))
+                    for dataset_urn in dataset_urns
+                ],
+            ),
         )
         self._emit_mcps(mcp)
         logger.info(f"Added input datasets to run {run_urn}")
@@ -440,7 +435,13 @@ def add_output_datasets_to_run(self, run_urn: str, dataset_urns: List[str]) -> N
             entity_urn=run_urn,
             entity_type="dataProcessInstance",
             aspect_name="dataProcessInstanceOutput",
-            aspect=DataProcessInstanceOutput(outputs=dataset_urns),
+            aspect=DataProcessInstanceOutput(
+                outputEdges=[
+                    EdgeClass(destinationUrn=str(dataset_urn))
+                    for dataset_urn in dataset_urns
+                ],
+                outputs=[],
+            ),
         )
         self._emit_mcps(mcp)
         logger.info(f"Added output datasets to run {run_urn}")

metadata-ingestion/examples/ai/dh_ai_client_sample.py

@@ -96,15 +96,19 @@
         end_timestamp=1628580001000,
     )
     # Create datasets
-    input_dataset_urn = client.create_dataset(
-        platform="snowflake",
-        name="iris_input",
-    )
+    input_dataset_urns = [
+        client.create_dataset(
+            platform="snowflake",
+            name="iris_input",
+        )
+    ]
 
-    output_dataset_urn = client.create_dataset(
-        platform="snowflake",
-        name="iris_ouptut",
-    )
+    output_dataset_urns = [
+        client.create_dataset(
+            platform="snowflake",
+            name="iris_ouptut",
+        )
+    ]
 
     # Add run to experiment
     client.add_run_to_experiment(run_urn=run_urn, experiment_urn=experiment_urn)
@@ -125,10 +129,6 @@
     )
 
     # Add input and output datasets to run
-    client.add_input_datasets_to_run(
-        run_urn=run_urn, dataset_urns=[str(input_dataset_urn)]
-    )
+    client.add_input_datasets_to_run(run_urn=run_urn, dataset_urns=input_dataset_urns)
 
-    client.add_output_datasets_to_run(
-        run_urn=run_urn, dataset_urns=[str(output_dataset_urn)]
-    )
+    client.add_output_datasets_to_run(run_urn=run_urn, dataset_urns=output_dataset_urns)

metadata-ingestion/src/datahub/ingestion/source/snowflake/snowflake_queries.py

@@ -403,6 +403,7 @@ def _parse_audit_log_row(
                     res["session_id"],
                     res["query_start_time"],
                     object_modified_by_ddl,
+                    res["query_type"],
                 )
             if known_ddl_entry:
                 return known_ddl_entry
@@ -537,40 +538,42 @@ def parse_ddl_query(
         session_id: str,
         timestamp: datetime,
         object_modified_by_ddl: dict,
+        query_type: str,
     ) -> Optional[Union[TableRename, TableSwap]]:
         timestamp = timestamp.astimezone(timezone.utc)
-        if object_modified_by_ddl[
-            "operationType"
-        ] == "ALTER" and object_modified_by_ddl["properties"].get("swapTargetName"):
-            urn1 = self.identifiers.gen_dataset_urn(
+        if (
+            object_modified_by_ddl["operationType"] == "ALTER"
+            and query_type == "RENAME_TABLE"
+            and object_modified_by_ddl["properties"].get("objectName")
+        ):
+            original_un = self.identifiers.gen_dataset_urn(
                 self.identifiers.get_dataset_identifier_from_qualified_name(
                     object_modified_by_ddl["objectName"]
                 )
             )
 
-            urn2 = self.identifiers.gen_dataset_urn(
+            new_urn = self.identifiers.gen_dataset_urn(
                 self.identifiers.get_dataset_identifier_from_qualified_name(
-                    object_modified_by_ddl["properties"]["swapTargetName"]["value"]
+                    object_modified_by_ddl["properties"]["objectName"]["value"]
                 )
             )
-
-            return TableSwap(urn1, urn2, query, session_id, timestamp)
+            return TableRename(original_un, new_urn, query, session_id, timestamp)
         elif object_modified_by_ddl[
             "operationType"
-        ] == "RENAME_TABLE" and object_modified_by_ddl["properties"].get("objectName"):
-            original_un = self.identifiers.gen_dataset_urn(
+        ] == "ALTER" and object_modified_by_ddl["properties"].get("swapTargetName"):
+            urn1 = self.identifiers.gen_dataset_urn(
                 self.identifiers.get_dataset_identifier_from_qualified_name(
                     object_modified_by_ddl["objectName"]
                 )
             )
 
-            new_urn = self.identifiers.gen_dataset_urn(
+            urn2 = self.identifiers.gen_dataset_urn(
                 self.identifiers.get_dataset_identifier_from_qualified_name(
-                    object_modified_by_ddl["properties"]["objectName"]["value"]
+                    object_modified_by_ddl["properties"]["swapTargetName"]["value"]
                 )
             )
 
-            return TableRename(original_un, new_urn, query, session_id, timestamp)
+            return TableSwap(urn1, urn2, query, session_id, timestamp)
         else:
             self.report.num_ddl_queries_dropped += 1
             return None

metadata-ingestion/src/datahub/ingestion/source/superset.py

@@ -23,6 +23,7 @@
     make_dataset_urn,
     make_dataset_urn_with_platform_instance,
     make_domain_urn,
+    make_schema_field_urn,
     make_user_urn,
 )
 from datahub.emitter.mcp_builder import add_domain_to_entity_wu
@@ -72,6 +73,9 @@
     DashboardInfoClass,
     DatasetLineageTypeClass,
     DatasetPropertiesClass,
+    FineGrainedLineageClass,
+    FineGrainedLineageDownstreamTypeClass,
+    FineGrainedLineageUpstreamTypeClass,
     GlobalTagsClass,
     OwnerClass,
     OwnershipClass,
@@ -80,6 +84,10 @@
     UpstreamClass,
     UpstreamLineageClass,
 )
+from datahub.sql_parsing.sqlglot_lineage import (
+    SqlParsingResult,
+    create_lineage_sql_parsed_result,
+)
 from datahub.utilities import config_clean
 from datahub.utilities.lossy_collections import LossyList
 from datahub.utilities.registries.domain_registry import DomainRegistry
@@ -342,7 +350,7 @@ def get_dataset_info(self, dataset_id: int) -> dict:
         )
         if dataset_response.status_code != 200:
             logger.warning(f"Failed to get dataset info: {dataset_response.text}")
-            dataset_response.raise_for_status()
+            return {}
         return dataset_response.json()
 
     def get_datasource_urn_from_id(
@@ -680,6 +688,88 @@ def gen_dataset_urn(self, datahub_dataset_name: str) -> str:
             env=self.config.env,
         )
 
+    def generate_virtual_dataset_lineage(
+        self,
+        parsed_query_object: SqlParsingResult,
+        datasource_urn: str,
+    ) -> UpstreamLineageClass:
+        cll = (
+            parsed_query_object.column_lineage
+            if parsed_query_object.column_lineage is not None
+            else []
+        )
+
+        fine_grained_lineages: List[FineGrainedLineageClass] = []
+
+        for cll_info in cll:
+            downstream = (
+                [make_schema_field_urn(datasource_urn, cll_info.downstream.column)]
+                if cll_info.downstream and cll_info.downstream.column
+                else []
+            )
+            upstreams = [
+                make_schema_field_urn(column_ref.table, column_ref.column)
+                for column_ref in cll_info.upstreams
+            ]
+            fine_grained_lineages.append(
+                FineGrainedLineageClass(
+                    downstreamType=FineGrainedLineageDownstreamTypeClass.FIELD,
+                    downstreams=downstream,
+                    upstreamType=FineGrainedLineageUpstreamTypeClass.FIELD_SET,
+                    upstreams=upstreams,
+                )
+            )
+
+        upstream_lineage = UpstreamLineageClass(
+            upstreams=[
+                UpstreamClass(
+                    type=DatasetLineageTypeClass.TRANSFORMED,
+                    dataset=input_table_urn,
+                )
+                for input_table_urn in parsed_query_object.in_tables
+            ],
+            fineGrainedLineages=fine_grained_lineages,
+        )
+        return upstream_lineage
+
+    def generate_physical_dataset_lineage(
+        self,
+        dataset_response: dict,
+        upstream_dataset: str,
+        datasource_urn: str,
+    ) -> UpstreamLineageClass:
+        # To generate column level lineage, we can manually decode the metadata
+        # to produce the ColumnLineageInfo
+        columns = dataset_response.get("result", {}).get("columns", [])
+        fine_grained_lineages: List[FineGrainedLineageClass] = []
+
+        for column in columns:
+            column_name = column.get("column_name", "")
+            if not column_name:
+                continue
+
+            downstream = [make_schema_field_urn(datasource_urn, column_name)]
+            upstreams = [make_schema_field_urn(upstream_dataset, column_name)]
+            fine_grained_lineages.append(
+                FineGrainedLineageClass(
+                    downstreamType=FineGrainedLineageDownstreamTypeClass.FIELD,
+                    downstreams=downstream,
+                    upstreamType=FineGrainedLineageUpstreamTypeClass.FIELD_SET,
+                    upstreams=upstreams,
+                )
+            )
+
+        upstream_lineage = UpstreamLineageClass(
+            upstreams=[
+                UpstreamClass(
+                    type=DatasetLineageTypeClass.TRANSFORMED,
+                    dataset=upstream_dataset,
+                )
+            ],
+            fineGrainedLineages=fine_grained_lineages,
+        )
+        return upstream_lineage
+
     def construct_dataset_from_dataset_data(
         self, dataset_data: dict
     ) -> DatasetSnapshot:
@@ -700,6 +790,14 @@ def construct_dataset_from_dataset_data(
         upstream_warehouse_platform = (
             dataset_response.get("result", {}).get("database", {}).get("backend")
         )
+        upstream_warehouse_db_name = (
+            dataset_response.get("result", {}).get("database", {}).get("database_name")
+        )
+
+        # if we have rendered sql, we always use that and defualt back to regular sql
+        sql = dataset_response.get("result", {}).get(
+            "rendered_sql"
+        ) or dataset_response.get("result", {}).get("sql")
 
         # Preset has a way of naming their platforms differently than
         # how datahub names them, so map the platform name to the correct naming
@@ -712,22 +810,28 @@ def construct_dataset_from_dataset_data(
         if upstream_warehouse_platform in warehouse_naming:
             upstream_warehouse_platform = warehouse_naming[upstream_warehouse_platform]
 
-        # TODO: Categorize physical vs virtual upstream dataset
-        # mark all upstream dataset as physical for now, in the future we would ideally like
-        # to differentiate physical vs virtual upstream datasets
-        tag_urn = f"urn:li:tag:{self.platform}:physical"
         upstream_dataset = self.get_datasource_urn_from_id(
             dataset_response, upstream_warehouse_platform
         )
-        upstream_lineage = UpstreamLineageClass(
-            upstreams=[
-                UpstreamClass(
-                    type=DatasetLineageTypeClass.TRANSFORMED,
-                    dataset=upstream_dataset,
-                    properties={"externalUrl": dataset_url},
-                )
-            ]
-        )
+
+        # Sometimes the field will be null instead of not existing
+        if sql == "null" or not sql:
+            tag_urn = f"urn:li:tag:{self.platform}:physical"
+            upstream_lineage = self.generate_physical_dataset_lineage(
+                dataset_response, upstream_dataset, datasource_urn
+            )
+        else:
+            tag_urn = f"urn:li:tag:{self.platform}:virtual"
+            parsed_query_object = create_lineage_sql_parsed_result(
+                query=sql,
+                default_db=upstream_warehouse_db_name,
+                platform=upstream_warehouse_platform,
+                platform_instance=None,
+                env=self.config.env,
+            )
+            upstream_lineage = self.generate_virtual_dataset_lineage(
+                parsed_query_object, datasource_urn
+            )
 
         dataset_info = DatasetPropertiesClass(
             name=dataset.table_name,

metadata-ingestion/src/datahub/ingestion/source/vertexai.py

@@ -805,8 +805,7 @@ def _search_endpoint(self, model: Model) -> List[Endpoint]:
                     endpoint_dict[resource.model].append(endpoint)
             self.endpoints = endpoint_dict
 
-        endpoints = self.endpoints.get(model.resource_name, [])
-        return endpoints
+        return self.endpoints.get(model.resource_name, [])
 
     def _make_ml_model_urn(self, model_version: VersionInfo, model_name: str) -> str:
         urn = builder.make_ml_model_urn(