acryl-datahub 0.14.1.13rc9__py3-none-any.whl → 0.15.0__py3-none-any.whl

datahub/ingestion/source/kafka/kafka_connect.py

@@ -282,10 +282,6 @@ class ConfluentJDBCSourceConnector:
         query: str
         transforms: list
 
-    def report_warning(self, key: str, reason: str) -> None:
-        logger.warning(f"{key}: {reason}")
-        self.report.report_warning(key, reason)
-
     def get_parser(
         self,
         connector_manifest: ConnectorManifest,
@@ -355,9 +351,9 @@ class ConfluentJDBCSourceConnector:
                     source_table = f"{table_name_tuple[-2]}.{source_table}"
                 else:
                     include_source_dataset = False
-                    self.report_warning(
-                        self.connector_manifest.name,
-                        f"could not find schema for table {source_table}",
+                    self.report.warning(
+                        "Could not find schema for table"
+                        f"{self.connector_manifest.name} : {source_table}",
                     )
             dataset_name: str = get_dataset_name(database_name, source_table)
             lineage = KafkaConnectLineage(
@@ -457,9 +453,9 @@ class ConfluentJDBCSourceConnector:
                     target_platform=KAFKA,
                 )
                 lineages.append(lineage)
-                self.report_warning(
+                self.report.warning(
+                    "Could not find input dataset, the connector has query configuration set",
                     self.connector_manifest.name,
-                    "could not find input dataset, the connector has query configuration set",
                 )
                 self.connector_manifest.lineages = lineages
                 return
@@ -535,24 +531,24 @@ class ConfluentJDBCSourceConnector:
                         include_source_dataset=False,
                     )
                 )
-                self.report_warning(
-                    self.connector_manifest.name,
-                    f"could not find input dataset, for connector topics {topic_names}",
+                self.report.warning(
+                    "Could not find input dataset for connector topics",
+                    f"{self.connector_manifest.name} : {topic_names}",
                 )
             self.connector_manifest.lineages = lineages
             return
         else:
             include_source_dataset = True
             if SINGLE_TRANSFORM and UNKNOWN_TRANSFORM:
-                self.report_warning(
-                    self.connector_manifest.name,
-                    f"could not find input dataset, connector has unknown transform - {transforms[0]['type']}",
+                self.report.warning(
+                    "Could not find input dataset, connector has unknown transform",
+                    f"{self.connector_manifest.name} : {transforms[0]['type']}",
                 )
                 include_source_dataset = False
             if not SINGLE_TRANSFORM and UNKNOWN_TRANSFORM:
-                self.report_warning(
+                self.report.warning(
+                    "Could not find input dataset, connector has one or more unknown transforms",
                     self.connector_manifest.name,
-                    "could not find input dataset, connector has one or more unknown transforms",
                 )
                 include_source_dataset = False
             lineages = self.default_get_lineages(
@@ -753,8 +749,10 @@ class DebeziumSourceConnector:
                     lineages.append(lineage)
             self.connector_manifest.lineages = lineages
         except Exception as e:
-            self.report.report_warning(
-                self.connector_manifest.name, f"Error resolving lineage: {e}"
+            self.report.warning(
+                "Error resolving lineage for connector",
+                self.connector_manifest.name,
+                exc=e,
             )
 
         return
@@ -783,10 +781,6 @@ class BigQuerySinkConnector:
         defaultDataset: Optional[str] = None
         version: str = "v1"
 
-    def report_warning(self, key: str, reason: str) -> None:
-        logger.warning(f"{key}: {reason}")
-        self.report.report_warning(key, reason)
-
     def get_parser(
         self,
         connector_manifest: ConnectorManifest,
@@ -917,9 +911,9 @@ class BigQuerySinkConnector:
             transformed_topic = self.apply_transformations(topic, transforms)
             dataset_table = self.get_dataset_table_for_topic(transformed_topic, parser)
             if dataset_table is None:
-                self.report_warning(
-                    self.connector_manifest.name,
-                    f"could not find target dataset for topic {transformed_topic}, please check your connector configuration",
+                self.report.warning(
+                    "Could not find target dataset for topic, please check your connector configuration"
+                    f"{self.connector_manifest.name} : {transformed_topic} ",
                 )
                 continue
             target_dataset = f"{project}.{dataset_table}"
@@ -954,10 +948,6 @@ class SnowflakeSinkConnector:
         schema_name: str
         topics_to_tables: Dict[str, str]
 
-    def report_warning(self, key: str, reason: str) -> None:
-        logger.warning(f"{key}: {reason}")
-        self.report.report_warning(key, reason)
-
     def get_table_name_from_topic_name(self, topic_name: str) -> str:
         """
         This function converts the topic name to a valid Snowflake table name using some rules.
@@ -1105,8 +1095,10 @@ class ConfluentS3SinkConnector:
                 )
             self.connector_manifest.lineages = lineages
         except Exception as e:
-            self.report.report_warning(
-                self.connector_manifest.name, f"Error resolving lineage: {e}"
+            self.report.warning(
+                "Error resolving lineage for connector",
+                self.connector_manifest.name,
+                exc=e,
             )
 
         return
@@ -1155,7 +1147,7 @@ class KafkaConnectSource(StatefulIngestionSourceBase):
             )
             self.session.auth = (self.config.username, self.config.password)
 
-        test_response = self.session.get(f"{self.config.connect_uri}")
+        test_response = self.session.get(f"{self.config.connect_uri}/connectors")
         test_response.raise_for_status()
         logger.info(f"Connection to {self.config.connect_uri} is ok")
         if not jpype.isJVMStarted():
@@ -1178,13 +1170,16 @@ class KafkaConnectSource(StatefulIngestionSourceBase):
 
         payload = connector_response.json()
 
-        for c in payload:
-            connector_url = f"{self.config.connect_uri}/connectors/{c}"
-            connector_response = self.session.get(connector_url)
-            manifest = connector_response.json()
-            connector_manifest = ConnectorManifest(**manifest)
-            if not self.config.connector_patterns.allowed(connector_manifest.name):
-                self.report.report_dropped(connector_manifest.name)
+        for connector_name in payload:
+            connector_url = f"{self.config.connect_uri}/connectors/{connector_name}"
+            connector_manifest = self._get_connector_manifest(
+                connector_name, connector_url
+            )
+            if (
+                connector_manifest is None
+                or not self.config.connector_patterns.allowed(connector_manifest.name)
+            ):
+                self.report.report_dropped(connector_name)
                 continue
 
             if self.config.provided_configs:
@@ -1195,19 +1190,11 @@ class KafkaConnectSource(StatefulIngestionSourceBase):
             connector_manifest.lineages = list()
             connector_manifest.url = connector_url
 
-            topics = self.session.get(
-                f"{self.config.connect_uri}/connectors/{c}/topics",
-            ).json()
-
-            connector_manifest.topic_names = topics[c]["topics"]
+            connector_manifest.topic_names = self._get_connector_topics(connector_name)
 
             # Populate Source Connector metadata
             if connector_manifest.type == SOURCE:
-                tasks = self.session.get(
-                    f"{self.config.connect_uri}/connectors/{c}/tasks",
-                ).json()
-
-                connector_manifest.tasks = tasks
+                connector_manifest.tasks = self._get_connector_tasks(connector_name)
 
                 # JDBC source connector lineages
                 if connector_manifest.config.get(CONNECTOR_CLASS).__eq__(
@@ -1246,7 +1233,7 @@ class KafkaConnectSource(StatefulIngestionSourceBase):
                         )
                         continue
 
-                    for topic in topics:
+                    for topic in connector_manifest.topic_names:
                         lineage = KafkaConnectLineage(
                             source_dataset=target_connector.source_dataset,
                             source_platform=target_connector.source_platform,
@@ -1286,6 +1273,49 @@ class KafkaConnectSource(StatefulIngestionSourceBase):
 
         return connectors_manifest
 
+    def _get_connector_manifest(
+        self, connector_name: str, connector_url: str
+    ) -> Optional[ConnectorManifest]:
+        try:
+            connector_response = self.session.get(connector_url)
+            connector_response.raise_for_status()
+        except Exception as e:
+            self.report.warning(
+                "Failed to get connector details", connector_name, exc=e
+            )
+            return None
+        manifest = connector_response.json()
+        connector_manifest = ConnectorManifest(**manifest)
+        return connector_manifest
+
+    def _get_connector_tasks(self, connector_name: str) -> dict:
+        try:
+            response = self.session.get(
+                f"{self.config.connect_uri}/connectors/{connector_name}/tasks",
+            )
+            response.raise_for_status()
+        except Exception as e:
+            self.report.warning(
+                "Error getting connector tasks", context=connector_name, exc=e
+            )
+            return {}
+
+        return response.json()
+
+    def _get_connector_topics(self, connector_name: str) -> List[str]:
+        try:
+            response = self.session.get(
+                f"{self.config.connect_uri}/connectors/{connector_name}/topics",
+            )
+            response.raise_for_status()
+        except Exception as e:
+            self.report.warning(
+                "Error getting connector topics", context=connector_name, exc=e
+            )
+            return []
+
+        return response.json()[connector_name]["topics"]
+
     def construct_flow_workunit(self, connector: ConnectorManifest) -> MetadataWorkUnit:
         connector_name = connector.name
         connector_type = connector.type

datahub/ingestion/source/looker/looker_liquid_tag.py

@@ -4,6 +4,7 @@ from typing import ClassVar, Optional, TextIO
 from liquid import Environment
 from liquid.ast import Node
 from liquid.context import Context
+from liquid.filter import string_filter
 from liquid.parse import expect, get_parser
 from liquid.stream import TokenStream
 from liquid.tag import Tag
@@ -81,12 +82,18 @@ class ConditionTag(Tag):
 custom_tags = [ConditionTag]
 
 
+@string_filter
+def sql_quote_filter(variable: str) -> str:
+    return f"'{variable}'"
+
+
 @lru_cache(maxsize=1)
 def _create_env() -> Environment:
-    env: Environment = Environment()
+    env: Environment = Environment(strict_filters=False)
     # register tag. One time activity
     for custom_tag in custom_tags:
         env.add_tag(custom_tag)
+    env.add_filter("sql_quote", sql_quote_filter)
     return env
 
 

datahub/ingestion/source/looker/lookml_concept_context.py

@@ -88,8 +88,7 @@ class LookerFieldContext:
         for upstream_field_match in re.finditer(r"\${TABLE}\.[\"]*([\.\w]+)", sql):
             matched_field = upstream_field_match.group(1)
             # Remove quotes from field names
-            matched_field = matched_field.replace('"', "").replace("`", "").lower()
-            column_names.append(matched_field)
+            column_names.append(matched_field.replace('"', "").replace("`", "").lower())
 
         return column_names
 

datahub/ingestion/source/looker/view_upstream.py

@@ -25,11 +25,13 @@ from datahub.ingestion.source.looker.lookml_config import (
     LookMLSourceReport,
 )
 from datahub.ingestion.source.looker.urn_functions import get_qualified_table_name
+from datahub.sql_parsing.schema_resolver import match_columns_to_schema
 from datahub.sql_parsing.sqlglot_lineage import (
     ColumnLineageInfo,
     ColumnRef,
     SqlParsingResult,
     Urn,
+    create_and_cache_schema_resolver,
     create_lineage_sql_parsed_result,
 )
 
@@ -200,7 +202,7 @@ def _generate_fully_qualified_name(
 class AbstractViewUpstream(ABC):
     """
     Implementation of this interface extracts the view upstream as per the way the view is bound to datasets.
-    For detail explanation please refer lookml_concept_context.LookerViewContext documentation.
+    For detail explanation, please refer lookml_concept_context.LookerViewContext documentation.
     """
 
     view_context: LookerViewContext
@@ -236,6 +238,47 @@ class AbstractViewUpstream(ABC):
     def create_fields(self) -> List[ViewField]:
         return []  # it is for the special case
 
+    def create_upstream_column_refs(
+        self, upstream_urn: str, downstream_looker_columns: List[str]
+    ) -> List[ColumnRef]:
+        """
+        - **`upstream_urn`**: The URN of the upstream dataset.
+
+        - **`expected_columns`**: These are the columns identified by the Looker connector as belonging to the `upstream_urn` dataset. However, there is potential for human error in specifying the columns of the upstream dataset. For example, a user might declare a column in lowercase, while on the actual platform, it may exist in uppercase, or vice versa.
+
+        - This function ensures consistency in column-level lineage by consulting GMS before creating the final `ColumnRef` instance, avoiding discrepancies.
+        """
+        schema_resolver = create_and_cache_schema_resolver(
+            platform=self.view_context.view_connection.platform,
+            platform_instance=self.view_context.view_connection.platform_instance,
+            env=self.view_context.view_connection.platform_env or self.config.env,
+            graph=self.ctx.graph,
+        )
+
+        urn, schema_info = schema_resolver.resolve_urn(urn=upstream_urn)
+
+        if schema_info:
+            actual_columns = match_columns_to_schema(
+                schema_info, downstream_looker_columns
+            )
+        else:
+            logger.info(
+                f"schema_info not found for dataset {urn} in GMS. Using expected_columns to form ColumnRef"
+            )
+            actual_columns = [column.lower() for column in downstream_looker_columns]
+
+        upstream_column_refs: List[ColumnRef] = []
+
+        for column in actual_columns:
+            upstream_column_refs.append(
+                ColumnRef(
+                    column=column,
+                    table=upstream_urn,
+                )
+            )
+
+        return upstream_column_refs
+
 
 class SqlBasedDerivedViewUpstream(AbstractViewUpstream, ABC):
     """
@@ -372,15 +415,12 @@ class SqlBasedDerivedViewUpstream(AbstractViewUpstream, ABC):
         # in-case of "select * from look_ml_view.SQL_TABLE_NAME" or extra field are defined in the looker view which is
         # referring to upstream table
         if self._get_upstream_dataset_urn() and not upstreams_column_refs:
-            upstreams_column_refs = [
-                ColumnRef(
-                    table=self._get_upstream_dataset_urn()[
-                        0
-                    ],  # 0th index has table of from clause
-                    column=column,
-                )
-                for column in field_context.column_name_in_sql_attribute()
-            ]
+            upstreams_column_refs = self.create_upstream_column_refs(
+                upstream_urn=self._get_upstream_dataset_urn()[
+                    0
+                ],  # 0th index has table of from clause,
+                downstream_looker_columns=field_context.column_name_in_sql_attribute(),
+            )
 
         # fix any derived view reference present in urn
         upstreams_column_refs = resolve_derived_view_urn_of_col_ref(
@@ -487,18 +527,18 @@ class NativeDerivedViewUpstream(AbstractViewUpstream):
             return upstream_column_refs
 
         explore_urn: str = self._get_upstream_dataset_urn()[0]
+        expected_columns: List[str] = []
 
         for column in field_context.column_name_in_sql_attribute():
             if column in self._get_explore_column_mapping():
                 explore_column: Dict = self._get_explore_column_mapping()[column]
-                upstream_column_refs.append(
-                    ColumnRef(
-                        column=explore_column.get("field", explore_column[NAME]),
-                        table=explore_urn,
-                    )
+                expected_columns.append(
+                    explore_column.get("field", explore_column[NAME])
                 )
 
-        return upstream_column_refs
+        return self.create_upstream_column_refs(
+            upstream_urn=explore_urn, downstream_looker_columns=expected_columns
+        )
 
     def get_upstream_dataset_urn(self) -> List[Urn]:
         return self._get_upstream_dataset_urn()
@@ -548,14 +588,10 @@ class RegularViewUpstream(AbstractViewUpstream):
     def get_upstream_column_ref(
         self, field_context: LookerFieldContext
     ) -> List[ColumnRef]:
-        upstream_column_ref: List[ColumnRef] = []
-
-        for column_name in field_context.column_name_in_sql_attribute():
-            upstream_column_ref.append(
-                ColumnRef(table=self._get_upstream_dataset_urn(), column=column_name)
-            )
-
-        return upstream_column_ref
+        return self.create_upstream_column_refs(
+            upstream_urn=self._get_upstream_dataset_urn(),
+            downstream_looker_columns=field_context.column_name_in_sql_attribute(),
+        )
 
     def get_upstream_dataset_urn(self) -> List[Urn]:
         return [self._get_upstream_dataset_urn()]
@@ -609,15 +645,14 @@ class DotSqlTableNameViewUpstream(AbstractViewUpstream):
         self, field_context: LookerFieldContext
     ) -> List[ColumnRef]:
         upstream_column_ref: List[ColumnRef] = []
+
         if not self._get_upstream_dataset_urn():
             return upstream_column_ref
 
-        for column_name in field_context.column_name_in_sql_attribute():
-            upstream_column_ref.append(
-                ColumnRef(table=self._get_upstream_dataset_urn()[0], column=column_name)
-            )
-
-        return upstream_column_ref
+        return self.create_upstream_column_refs(
+            upstream_urn=self._get_upstream_dataset_urn()[0],
+            downstream_looker_columns=field_context.column_name_in_sql_attribute(),
+        )
 
     def get_upstream_dataset_urn(self) -> List[Urn]:
         return self._get_upstream_dataset_urn()

datahub/ingestion/source/metadata/business_glossary.py

@@ -45,6 +45,9 @@ class Owners(ConfigModel):
     groups: Optional[List[str]] = None
 
 
+OwnersMultipleTypes = Union[List[Owners], Owners]
+
+
 class KnowledgeCard(ConfigModel):
     url: Optional[str] = None
     label: Optional[str] = None
@@ -57,7 +60,7 @@ class GlossaryTermConfig(ConfigModel):
     term_source: Optional[str] = None
     source_ref: Optional[str] = None
     source_url: Optional[str] = None
-    owners: Optional[Owners] = None
+    owners: Optional[OwnersMultipleTypes] = None
     inherits: Optional[List[str]] = None
     contains: Optional[List[str]] = None
     values: Optional[List[str]] = None
@@ -74,7 +77,7 @@ class GlossaryNodeConfig(ConfigModel):
     id: Optional[str] = None
     name: str
     description: str
-    owners: Optional[Owners] = None
+    owners: Optional[OwnersMultipleTypes] = None
     terms: Optional[List["GlossaryTermConfig"]] = None
     nodes: Optional[List["GlossaryNodeConfig"]] = None
     knowledge_links: Optional[List[KnowledgeCard]] = None
@@ -88,7 +91,7 @@ class DefaultConfig(ConfigModel):
     """Holds defaults for populating fields in glossary terms"""
 
     source: Optional[str] = None
-    owners: Owners
+    owners: OwnersMultipleTypes
     url: Optional[str] = None
     source_type: str = "INTERNAL"
 
@@ -153,30 +156,44 @@ def make_glossary_term_urn(
     return "urn:li:glossaryTerm:" + create_id(path, default_id, enable_auto_id)
 
 
-def get_owners(owners: Owners) -> models.OwnershipClass:
-    ownership_type, ownership_type_urn = validate_ownership_type(owners.type)
+def get_owners_multiple_types(owners: OwnersMultipleTypes) -> models.OwnershipClass:
+    """Allows owner types to be a list and maintains backward compatibility"""
+    if isinstance(owners, Owners):
+        return models.OwnershipClass(owners=list(get_owners(owners)))
+
+    owners_meta: List[models.OwnerClass] = []
+    for owner in owners:
+        owners_meta.extend(get_owners(owner))
+
+    return models.OwnershipClass(owners=owners_meta)
+
+
+def get_owners(owners: Owners) -> Iterable[models.OwnerClass]:
+    actual_type = owners.type or models.OwnershipTypeClass.DEVELOPER
+
+    if actual_type.startswith("urn:li:ownershipType:"):
+        ownership_type: str = "CUSTOM"
+        ownership_type_urn: Optional[str] = actual_type
+    else:
+        ownership_type, ownership_type_urn = validate_ownership_type(actual_type)
+
     if owners.typeUrn is not None:
         ownership_type_urn = owners.typeUrn
-    owners_meta: List[models.OwnerClass] = []
+
     if owners.users is not None:
-        owners_meta = owners_meta + [
-            models.OwnerClass(
+        for o in owners.users:
+            yield models.OwnerClass(
                 owner=make_user_urn(o),
                 type=ownership_type,
                 typeUrn=ownership_type_urn,
             )
-            for o in owners.users
-        ]
     if owners.groups is not None:
-        owners_meta = owners_meta + [
-            models.OwnerClass(
+        for o in owners.groups:
+            yield models.OwnerClass(
                 owner=make_group_urn(o),
                 type=ownership_type,
                 typeUrn=ownership_type_urn,
             )
-            for o in owners.groups
-        ]
-    return models.OwnershipClass(owners=owners_meta)
 
 
 def get_mces(
@@ -185,7 +202,7 @@ def get_mces(
     ingestion_config: BusinessGlossarySourceConfig,
     ctx: PipelineContext,
 ) -> Iterable[Union[MetadataChangeProposalWrapper, models.MetadataChangeEventClass]]:
-    root_owners = get_owners(glossary.owners)
+    root_owners = get_owners_multiple_types(glossary.owners)
 
     if glossary.nodes:
         for node in glossary.nodes:
@@ -270,7 +287,7 @@ def get_mces_from_node(
     node_owners = parentOwners
     if glossaryNode.owners is not None:
         assert glossaryNode.owners is not None
-        node_owners = get_owners(glossaryNode.owners)
+        node_owners = get_owners_multiple_types(glossaryNode.owners)
 
     node_snapshot = models.GlossaryNodeSnapshotClass(
         urn=node_urn,
@@ -426,7 +443,7 @@ def get_mces_from_term(
     ownership: models.OwnershipClass = parentOwnership
     if glossaryTerm.owners is not None:
         assert glossaryTerm.owners is not None
-        ownership = get_owners(glossaryTerm.owners)
+        ownership = get_owners_multiple_types(glossaryTerm.owners)
     aspects.append(ownership)
 
     if glossaryTerm.domain is not None:

datahub/ingestion/source/mode.py

@@ -18,7 +18,6 @@ from pydantic import Field, validator
 from requests.adapters import HTTPAdapter, Retry
 from requests.exceptions import ConnectionError
 from requests.models import HTTPBasicAuth, HTTPError
-from sqllineage.runner import LineageRunner
 from tenacity import retry_if_exception_type, stop_after_attempt, wait_exponential
 
 import datahub.emitter.mce_builder as builder
@@ -820,28 +819,6 @@ class ModeSource(StatefulIngestionSourceBase):
             )
         return None
 
-    @lru_cache(maxsize=None)
-    def _get_source_from_query(self, raw_query: str) -> set:
-        query = self._replace_definitions(raw_query)
-        parser = LineageRunner(query)
-        source_paths = set()
-        try:
-            for table in parser.source_tables:
-                sources = str(table).split(".")
-                source_schema, source_table = sources[-2], sources[-1]
-                if source_schema == "<default>":
-                    source_schema = str(self.config.default_schema)
-
-                source_paths.add(f"{source_schema}.{source_table}")
-        except Exception as e:
-            self.report.report_failure(
-                title="Failed to Extract Lineage From Query",
-                message="Unable to retrieve lineage from Mode query.",
-                context=f"Query: {raw_query}, Error: {str(e)}",
-            )
-
-        return source_paths
-
     def _get_datasource_urn(
         self,
         platform: str,