openmetadata-ingestion 1.6.3.2.dev1__py3-none-any.whl → 1.6.4.0.dev0__py3-none-any.whl

metadata/ingestion/lineage/sql_lineage.py

@@ -16,8 +16,10 @@ import traceback
 from collections import defaultdict
 from typing import Any, Iterable, List, Optional, Tuple, Union
 
+import networkx as nx
 from collate_sqllineage.core.models import Column, DataFunction
 from collate_sqllineage.core.models import Table as LineageTable
+from networkx import DiGraph
 
 from metadata.generated.schema.api.lineage.addLineage import AddLineageRequest
 from metadata.generated.schema.entity.data.storedProcedure import (
@@ -509,6 +511,7 @@ def _create_lineage_by_table_name(
     column_lineage_map: dict,
     lineage_source: LineageSource = LineageSource.QueryLineage,
     procedure: Optional[EntityReference] = None,
+    graph: DiGraph = None,
 ) -> Iterable[Either[AddLineageRequest]]:
     """
     This method is to create a lineage between two tables
@@ -539,6 +542,11 @@ def _create_lineage_by_table_name(
                 logger.debug(
                     f"WARNING: Table entity [{table_name}] not found in OpenMetadata"
                 )
+        if graph is not None and (not from_table_entities or not to_table_entities):
+            graph.add_node(from_table, entity=from_table_entities)
+            graph.add_node(to_table, entity=to_table_entities)
+            graph.add_edge(from_table, to_table, query=masked_query)
+            return
 
         for from_entity, to_entity in itertools.product(
             from_table_entities or [], to_table_entities or []
@@ -606,6 +614,7 @@ def get_lineage_by_query(
     dialect: Dialect,
     timeout_seconds: int = LINEAGE_PARSING_TIMEOUT,
     lineage_source: LineageSource = LineageSource.QueryLineage,
+    graph: DiGraph = None,
 ) -> Iterable[Either[AddLineageRequest]]:
     """
     This method parses the query to get source, target and intermediate table names to create lineage,
@@ -616,8 +625,8 @@ def get_lineage_by_query(
 
     try:
         lineage_parser = LineageParser(query, dialect, timeout_seconds=timeout_seconds)
-        masked_query = lineage_parser.masked_query or query
-        logger.debug(f"Running lineage with query: {masked_query}")
+        masked_query = lineage_parser.masked_query
+        logger.debug(f"Running lineage with query: {masked_query or query}")
 
         raw_column_lineage = lineage_parser.column_lineage
         column_lineage.update(populate_column_lineage_map(raw_column_lineage))
@@ -645,6 +654,7 @@ def get_lineage_by_query(
                         column_lineage_map=column_lineage,
                         lineage_source=lineage_source,
                         procedure=procedure,
+                        graph=graph,
                     )
             for target_table in lineage_parser.target_tables:
                 yield from _create_lineage_by_table_name(
@@ -682,11 +692,12 @@ def get_lineage_by_query(
                             column_lineage_map=column_lineage,
                             lineage_source=lineage_source,
                             procedure=procedure,
+                            graph=graph,
                         )
         if not lineage_parser.query_parsing_success:
             query_parsing_failures.add(
                 QueryParsingError(
-                    query=masked_query,
+                    query=masked_query or query,
                     error=lineage_parser.query_parsing_failure_reason,
                 )
             )
@@ -710,6 +721,7 @@ def get_lineage_via_table_entity(
     dialect: Dialect,
     timeout_seconds: int = LINEAGE_PARSING_TIMEOUT,
     lineage_source: LineageSource = LineageSource.QueryLineage,
+    graph: DiGraph = None,
 ) -> Iterable[Either[AddLineageRequest]]:
     """Get lineage from table entity"""
     column_lineage = {}
@@ -717,8 +729,10 @@ def get_lineage_via_table_entity(
 
     try:
         lineage_parser = LineageParser(query, dialect, timeout_seconds=timeout_seconds)
-        masked_query = lineage_parser.masked_query or query
-        logger.debug(f"Getting lineage via table entity using query: {masked_query}")
+        masked_query = lineage_parser.masked_query
+        logger.debug(
+            f"Getting lineage via table entity using query: {masked_query or query}"
+        )
         to_table_name = table_entity.name.root
 
         for from_table_name in lineage_parser.source_tables:
@@ -743,6 +757,7 @@ def get_lineage_via_table_entity(
                     column_lineage_map=column_lineage,
                     lineage_source=lineage_source,
                     procedure=procedure,
+                    graph=graph,
                 ) or []
         if not lineage_parser.query_parsing_success:
             query_parsing_failures.add(
@@ -759,3 +774,94 @@ def get_lineage_via_table_entity(
                 stackTrace=traceback.format_exc(),
             )
         )
+
+
+def _process_sequence(
+    sequence: List[Any], graph: DiGraph
+) -> Iterable[Either[AddLineageRequest]]:
+    """
+    Process a sequence of nodes to generate lineage information.
+    """
+    from_node = None
+    queries = set()
+    clean_queries = False
+    previous_node = None
+    for node in sequence:
+        try:
+            if clean_queries:
+                queries.clear()
+                clean_queries = False
+            current_node = graph.nodes[node]
+            current_entity = current_node.get("entity")
+
+            if (
+                previous_node is not None
+                and graph.edges[(previous_node, node)].get("query") is not None
+            ):
+                queries.add(graph.edges[(previous_node, node)].get("query"))
+
+            if current_entity and from_node is not None:
+                for from_entity, to_entity in itertools.product(
+                    from_node.get("entity") or [], current_entity or []
+                ):
+                    if to_entity and from_entity:
+                        yield _build_table_lineage(
+                            to_entity=to_entity,
+                            from_entity=from_entity,
+                            to_table_raw_name=str(node),
+                            from_table_raw_name=str(from_node),
+                            masked_query="\n--------\n".join(queries),
+                            column_lineage_map={},
+                            lineage_source=LineageSource.QueryLineage,
+                            procedure=None,
+                        )
+                        clean_queries = True
+
+            if current_entity:
+                from_node = graph.nodes[node]
+            previous_node = node
+        except Exception as exc:
+            logger.debug(traceback.format_exc())
+            logger.error(f"Error creating lineage for node [{node}]: {exc}")
+
+
+def get_lineage_by_graph(
+    graph: DiGraph,
+) -> Iterable[Either[AddLineageRequest]]:
+    """
+    Generate lineage information from a directed graph.
+    This method processes a directed graph to extract lineage information by identifying
+    weakly connected components and traversing each component to generate sequences of nodes.
+    It then yields lineage information for each sequence.
+    Args:
+        graph (DiGraph): A directed graph representing the lineage.
+    Raises:
+        Exception: If an error occurs during the lineage creation process, it logs the error.
+    """
+    if graph is None:
+        return
+
+    # Get all weakly connected components
+    components = list(nx.weakly_connected_components(graph))
+
+    # Extract each component as an independent subgraph
+    independent_subtrees = [
+        graph.subgraph(component).copy() for component in components
+    ]
+
+    # Print results in the desired format
+    for subtree in independent_subtrees:
+        # Find a root node (node with no incoming edges)
+        root = [node for node in subtree if subtree.in_degree(node) == 0][0]
+
+        # Traverse from the root to get the sequence of nodes
+        current = root
+        sequence = [current]
+        while True:
+            successors = list(subtree.successors(current))
+            if not successors:
+                break
+            current = successors[0]
+            sequence.append(current)
+
+        yield from _process_sequence(sequence, subtree)

metadata/ingestion/ometa/mixins/query_mixin.py

@@ -42,6 +42,8 @@ class OMetaQueryMixin:
         return str(result.hexdigest())
 
     def _get_or_create_query(self, query: CreateQueryRequest) -> Optional[Query]:
+        if query.query.root is None:
+            return None
         query_hash = self._get_query_hash(query=query.query.root)
         query_entity = self.get_by_name(entity=Query, fqn=query_hash)
         if query_entity is None:

metadata/ingestion/ometa/routes.py

@@ -92,6 +92,7 @@ from metadata.generated.schema.api.services.createStorageService import (
 from metadata.generated.schema.api.services.ingestionPipelines.createIngestionPipeline import (
     CreateIngestionPipelineRequest,
 )
+from metadata.generated.schema.api.teams.createPersona import CreatePersonaRequest
 from metadata.generated.schema.api.teams.createRole import CreateRoleRequest
 from metadata.generated.schema.api.teams.createTeam import CreateTeamRequest
 from metadata.generated.schema.api.teams.createUser import CreateUserRequest
@@ -157,6 +158,7 @@ from metadata.generated.schema.entity.services.mlmodelService import MlModelServ
 from metadata.generated.schema.entity.services.pipelineService import PipelineService
 from metadata.generated.schema.entity.services.searchService import SearchService
 from metadata.generated.schema.entity.services.storageService import StorageService
+from metadata.generated.schema.entity.teams.persona import Persona
 from metadata.generated.schema.entity.teams.role import Role
 from metadata.generated.schema.entity.teams.team import Team
 from metadata.generated.schema.entity.teams.user import AuthenticationMechanism, User
@@ -217,6 +219,8 @@ ROUTES = {
     CreateTeamRequest.__name__: "/teams",
     User.__name__: "/users",
     CreateUserRequest.__name__: "/users",
+    Persona.__name__: "/personas",
+    CreatePersonaRequest.__name__: "/personas",
     AuthenticationMechanism.__name__: "/users/auth-mechanism",
     Bot.__name__: "/bots",
     CreateBot.__name__: "/bots",

metadata/ingestion/source/api/rest/metadata.py

@@ -171,7 +171,7 @@ class RestSource(ApiServiceSource):
                     break
             return filtered_paths
         except Exception as err:
-            logger.info(
+            logger.warning(
                 f"Error while filtering endpoints for collection {collection.name.root}"
             )
             return None
@@ -184,7 +184,7 @@ class RestSource(ApiServiceSource):
                 endpoint.name = f"{path} - {method_type}"
             return endpoint
         except Exception as err:
-            logger.info(f"Error while parsing endpoint data: {err}")
+            logger.warning(f"Error while parsing endpoint data: {err}")
         return None
 
     def _generate_collection_url(self, collection_name: str) -> Optional[AnyUrl]:
@@ -195,7 +195,7 @@ class RestSource(ApiServiceSource):
                     f"{self.config.serviceConnection.root.config.openAPISchemaURL}#tag/{collection_name}"
                 )
         except Exception as err:
-            logger.info(f"Error while generating collection url: {err}")
+            logger.warning(f"Error while generating collection url: {err}")
         return None
 
     def _generate_endpoint_url(self, endpoint_name: str) -> AnyUrl:
@@ -203,15 +203,14 @@ class RestSource(ApiServiceSource):
         base_url = self.config.serviceConnection.root.config.openAPISchemaURL
         if endpoint_name:
             return AnyUrl(f"{base_url}#operation/{endpoint_name}")
-        else:
-            return AnyUrl(base_url)
+        return AnyUrl(base_url)
 
     def _get_api_request_method(self, method_type: str) -> Optional[str]:
         """fetch endpoint request method"""
         try:
             return ApiRequestMethod[method_type.upper()]
         except KeyError as err:
-            logger.info(f"Keyerror while fetching request method: {err}")
+            logger.warning(f"Keyerror while fetching request method: {err}")
         return None
 
     def _get_request_schema(self, info: dict) -> Optional[APISchema]:
@@ -227,9 +226,9 @@ class RestSource(ApiServiceSource):
             if not schema_ref:
                 logger.debug("No request schema found for the endpoint")
                 return None
-            return self._process_schema(schema_ref)
+            return APISchema(schemaFields=self.process_schema_fields(schema_ref))
         except Exception as err:
-            logger.info(f"Error while parsing request schema: {err}")
+            logger.warning(f"Error while parsing request schema: {err}")
         return None
 
     def _get_response_schema(self, info: dict) -> Optional[APISchema]:
@@ -246,33 +245,43 @@ class RestSource(ApiServiceSource):
             if not schema_ref:
                 logger.debug("No response schema found for the endpoint")
                 return None
-            return self._process_schema(schema_ref)
+            return APISchema(schemaFields=self.process_schema_fields(schema_ref))
         except Exception as err:
-            logger.info(f"Error while parsing response schema: {err}")
+            logger.warning(f"Error while parsing response schema: {err}")
         return None
 
-    def _process_schema(self, schema_ref: str) -> Optional[List[APISchema]]:
-        """process schema"""
+    def process_schema_fields(self, schema_ref: str) -> Optional[List[FieldModel]]:
         try:
-            schema_ref = schema_ref.split("/")[-1]
+            schema_name = schema_ref.split("/")[-1]
             schema_fields = (
-                self.json_response.get("components").get("schemas").get(schema_ref)
+                self.json_response.get("components").get("schemas").get(schema_name)
             )
 
             fetched_fields = []
             for key, val in schema_fields.get("properties", {}).items():
                 dtype = val.get("type")
-                if not dtype:
-                    continue
-                fetched_fields.append(
-                    FieldModel(
-                        name=key,
-                        dataType=DataTypeTopic[dtype.upper()]
+                if dtype:
+                    parsed_dtype = (
+                        DataTypeTopic[dtype.upper()]
                         if dtype.upper() in DataTypeTopic.__members__
-                        else DataTypeTopic.UNKNOWN,
+                        else DataTypeTopic.UNKNOWN
                     )
-                )
-            return APISchema(schemaFields=fetched_fields)
+                    fetched_fields.append(FieldModel(name=key, dataType=parsed_dtype))
+                else:
+                    # If type of field is not defined then check for sub-schema
+                    # Check if it's `object` type field
+                    # check infinite recrusrion by comparing with parent(schema_ref)
+                    object_children = None
+                    if val.get("$ref") and val.get("$ref") != schema_ref:
+                        object_children = self.process_schema_fields(val.get("$ref"))
+                    fetched_fields.append(
+                        FieldModel(
+                            name=key,
+                            dataType=DataTypeTopic.UNKNOWN,
+                            children=object_children,
+                        )
+                    )
+            return fetched_fields
         except Exception as err:
-            logger.info(f"Error while processing request schema: {err}")
+            logger.warning(f"Error while processing schema fields: {err}")
         return None

metadata/ingestion/source/dashboard/looker/metadata.py

@@ -21,6 +21,7 @@ Notes:
 """
 import copy
 import os
+import re
 import traceback
 from datetime import datetime
 from pathlib import Path
@@ -39,6 +40,7 @@ from typing import (
 
 import giturlparse
 import lkml
+import networkx as nx
 from liquid import Template
 from looker_sdk.sdk.api40.methods import Looker40SDK
 from looker_sdk.sdk.api40.models import Dashboard as LookerDashboard
@@ -134,6 +136,10 @@ logger = ingestion_logger()
 LIST_DASHBOARD_FIELDS = ["id", "title"]
 IMPORTED_PROJECTS_DIR = "imported_projects"
 
+# we need to find the derived references in the SQL query using regex
+# https://cloud.google.com/looker/docs/derived-tables#referencing_derived_tables_in_other_derived_tables
+DERIVED_REFERENCES = r"\${([\w\s\d_.]+)\.SQL_TABLE_NAME}"
+
 # Here we can update the fields to get further information, such as:
 # created_at, updated_at, last_updater_id, deleted_at, deleter_id, favorite_count, last_viewed_at
 GET_DASHBOARD_FIELDS = [
@@ -165,6 +171,13 @@ def build_datamodel_name(model_name: str, explore_name: str) -> str:
     return clean_dashboard_name(model_name + "_" + explore_name)
 
 
+def find_derived_references(sql_query: str) -> List[str]:
+    if sql_query is None:
+        return []
+    matches = re.findall(DERIVED_REFERENCES, sql_query)
+    return matches
+
+
 class LookerSource(DashboardServiceSource):
     """
     Looker Source Class.
@@ -172,6 +185,8 @@ class LookerSource(DashboardServiceSource):
     Its client uses Looker 40 from the SDK: client = looker_sdk.init40()
     """
 
+    # pylint: disable=too-many-instance-attributes
+
     config: WorkflowSource
     metadata: OpenMetadata
     client: Looker40SDK
@@ -192,6 +207,10 @@ class LookerSource(DashboardServiceSource):
         self._main__lookml_manifest: Optional[LookMLManifest] = None
         self._view_data_model: Optional[DashboardDataModel] = None
 
+        self._parsed_views: Optional[Dict[str, str]] = {}
+        self._unparsed_views: Optional[Dict[str, str]] = {}
+        self._derived_dependencies = nx.DiGraph()
+
         self._added_lineage: Optional[Dict] = {}
 
     @classmethod
@@ -557,6 +576,68 @@ class LookerSource(DashboardServiceSource):
                     )
                 )
 
+    def replace_derived_references(self, sql_query):
+        """
+        Replace all derived references with the parsed views sql query
+        will replace the derived references in the SQL query using regex
+        for e.g. It will replace ${view_name.SQL_TABLE_NAME} with the parsed view query for view_name
+        https://cloud.google.com/looker/docs/derived-tables#referencing_derived_tables_in_other_derived_tables
+        """
+        try:
+            sql_query = re.sub(
+                DERIVED_REFERENCES,
+                # from `${view_name.SQL_TABLE_NAME}` we want the `view_name`.
+                # match.group(1) will give us the `view_name`
+                lambda match: f"({self._parsed_views.get(match.group(1), match.group(0))})",
+                sql_query,
+            )
+        except Exception as e:
+            logger.warning(
+                f"Something went wrong while replacing derived view references: {e}"
+            )
+        return sql_query
+
+    def build_lineage_for_unparsed_views(self) -> Iterable[Either[AddLineageRequest]]:
+        """
+        build lineage by parsing the unparsed views containing derived references
+        """
+        try:
+            # Doing a reversed topological sort to process the views in the right order
+            for view_name in reversed(
+                list(nx.topological_sort(self._derived_dependencies))
+            ):
+                if view_name in self._parsed_views:
+                    # Skip if already processed
+                    continue
+                sql_query = self.replace_derived_references(
+                    self._unparsed_views[view_name]
+                )
+                if view_references := find_derived_references(sql_query):
+                    # There are still derived references in the view query
+                    logger.debug(
+                        f"Views {view_references} not found for {view_name}. Skipping."
+                    )
+                    continue
+                self._parsed_views[view_name] = sql_query
+                del self._unparsed_views[view_name]
+                yield from self._build_lineage_for_view(view_name, sql_query)
+
+        except Exception as err:
+            yield Either(
+                left=StackTraceError(
+                    name="parse_unparsed_views",
+                    error=f"Error parsing unparsed views: {err}",
+                    stackTrace=traceback.format_exc(),
+                )
+            )
+
+    def _add_dependency_edge(self, view_name: str, view_references: List[str]):
+        """
+        Add a dependency edge between the view and the derived reference
+        """
+        for dependent_view_name in view_references:
+            self._derived_dependencies.add_edge(view_name, dependent_view_name)
+
     def add_view_lineage(
         self, view: LookMlView, explore: LookmlModelExplore
     ) -> Iterable[Either[AddLineageRequest]]:
@@ -589,6 +670,7 @@ class LookerSource(DashboardServiceSource):
                 for db_service_name in db_service_names or []:
                     dialect = self._get_db_dialect(db_service_name)
                     source_table_name = self._clean_table_name(sql_table_name, dialect)
+                    self._parsed_views[view.name] = source_table_name
 
                     # View to the source is only there if we are informing the dbServiceNames
                     yield self.build_lineage_request(
@@ -601,20 +683,19 @@ class LookerSource(DashboardServiceSource):
                 sql_query = view.derived_table.sql
                 if not sql_query:
                     return
+                if find_derived_references(sql_query):
+                    sql_query = self.replace_derived_references(sql_query)
+                    # If we still have derived references, we cannot process the view
+                    if view_references := find_derived_references(sql_query):
+                        self._add_dependency_edge(view.name, view_references)
+                        logger.warning(
+                            f"Not all references are replaced for view [{view.name}]. Parsing it later."
+                        )
+                        return
                 logger.debug(f"Processing view [{view.name}] with SQL: \n[{sql_query}]")
-                for db_service_name in db_service_names or []:
-                    lineage_parser = LineageParser(
-                        sql_query,
-                        self._get_db_dialect(db_service_name),
-                        timeout_seconds=30,
-                    )
-                    if lineage_parser.source_tables:
-                        for from_table_name in lineage_parser.source_tables:
-                            yield self.build_lineage_request(
-                                source=str(from_table_name),
-                                db_service_name=db_service_name,
-                                to_entity=self._view_data_model,
-                            )
+                yield from self._build_lineage_for_view(view.name, sql_query)
+                if self._unparsed_views:
+                    self.build_lineage_for_unparsed_views()
 
         except Exception as err:
             yield Either(
@@ -625,6 +706,27 @@ class LookerSource(DashboardServiceSource):
                 )
             )
 
+    def _build_lineage_for_view(
+        self, view_name: str, sql_query: str
+    ) -> Iterable[Either[AddLineageRequest]]:
+        """
+        Parse the SQL query and build lineage for the view.
+        """
+        for db_service_name in self.get_db_service_names() or []:
+            lineage_parser = LineageParser(
+                sql_query,
+                self._get_db_dialect(db_service_name),
+                timeout_seconds=30,
+            )
+            if lineage_parser.source_tables:
+                self._parsed_views[view_name] = sql_query
+                for from_table_name in lineage_parser.source_tables:
+                    yield self.build_lineage_request(
+                        source=str(from_table_name),
+                        db_service_name=db_service_name,
+                        to_entity=self._view_data_model,
+                    )
+
     def _get_db_dialect(self, db_service_name) -> Dialect:
         db_service = self.metadata.get_by_name(DatabaseService, db_service_name)
         return ConnectionTypeDialectMapper.dialect_of(