openmetadata-ingestion 1.6.0.0rc2__py3-none-any.whl → 1.6.0.0rc3__py3-none-any.whl

metadata/generated/schema/type/entityReferenceList.py

@@ -1,6 +1,6 @@
 # generated by datamodel-codegen:
 #   filename:  type/entityReferenceList.json
-#   timestamp: 2024-11-29T19:36:27+00:00
+#   timestamp: 2024-12-05T13:16:41+00:00
 
 from __future__ import annotations
 

metadata/generated/schema/type/entityRelationship.py

@@ -1,6 +1,6 @@
 # generated by datamodel-codegen:
 #   filename:  type/entityRelationship.json
-#   timestamp: 2024-11-29T19:36:27+00:00
+#   timestamp: 2024-12-05T13:16:41+00:00
 
 from __future__ import annotations
 
@@ -38,6 +38,7 @@ class RelationshipType(Enum):
     expert = 'expert'
     editedBy = 'editedBy'
     defaultsTo = 'defaultsTo'
+    relatesTo = 'relatesTo'
 
 
 class EntityRelationship(BaseModel):

metadata/generated/schema/type/entityUsage.py

@@ -1,6 +1,6 @@
 # generated by datamodel-codegen:
 #   filename:  type/entityUsage.json
-#   timestamp: 2024-11-29T19:36:27+00:00
+#   timestamp: 2024-12-05T13:16:41+00:00
 
 from __future__ import annotations
 

metadata/generated/schema/type/filterPattern.py

@@ -1,6 +1,6 @@
 # generated by datamodel-codegen:
 #   filename:  type/filterPattern.json
-#   timestamp: 2024-11-29T19:36:27+00:00
+#   timestamp: 2024-12-05T13:16:41+00:00
 
 from __future__ import annotations
 

metadata/generated/schema/type/function.py

@@ -1,6 +1,6 @@
 # generated by datamodel-codegen:
 #   filename:  type/function.json
-#   timestamp: 2024-11-29T19:36:27+00:00
+#   timestamp: 2024-12-05T13:16:41+00:00
 
 from __future__ import annotations
 

metadata/generated/schema/type/include.py

@@ -1,6 +1,6 @@
 # generated by datamodel-codegen:
 #   filename:  type/include.json
-#   timestamp: 2024-11-29T19:36:27+00:00
+#   timestamp: 2024-12-05T13:16:41+00:00
 
 from __future__ import annotations
 

metadata/generated/schema/type/jdbcConnection.py

@@ -1,6 +1,6 @@
 # generated by datamodel-codegen:
 #   filename:  type/jdbcConnection.json
-#   timestamp: 2024-11-29T19:36:27+00:00
+#   timestamp: 2024-12-05T13:16:41+00:00
 
 from __future__ import annotations
 

metadata/generated/schema/type/lifeCycle.py

@@ -1,6 +1,6 @@
 # generated by datamodel-codegen:
 #   filename:  type/lifeCycle.json
-#   timestamp: 2024-11-29T19:36:27+00:00
+#   timestamp: 2024-12-05T13:16:41+00:00
 
 from __future__ import annotations
 

metadata/generated/schema/type/paging.py

@@ -1,6 +1,6 @@
 # generated by datamodel-codegen:
 #   filename:  type/paging.json
-#   timestamp: 2024-11-29T19:36:27+00:00
+#   timestamp: 2024-12-05T13:16:41+00:00
 
 from __future__ import annotations
 

metadata/generated/schema/type/profile.py

@@ -1,6 +1,6 @@
 # generated by datamodel-codegen:
 #   filename:  type/profile.json
-#   timestamp: 2024-11-29T19:36:27+00:00
+#   timestamp: 2024-12-05T13:16:41+00:00
 
 from __future__ import annotations
 

metadata/generated/schema/type/queryParserData.py

@@ -1,6 +1,6 @@
 # generated by datamodel-codegen:
 #   filename:  type/queryParserData.json
-#   timestamp: 2024-11-29T19:36:27+00:00
+#   timestamp: 2024-12-05T13:16:41+00:00
 
 from __future__ import annotations
 

metadata/generated/schema/type/reaction.py

@@ -1,6 +1,6 @@
 # generated by datamodel-codegen:
 #   filename:  type/reaction.json
-#   timestamp: 2024-11-29T19:36:27+00:00
+#   timestamp: 2024-12-05T13:16:41+00:00
 
 from __future__ import annotations
 

metadata/generated/schema/type/schedule.py

@@ -1,6 +1,6 @@
 # generated by datamodel-codegen:
 #   filename:  type/schedule.json
-#   timestamp: 2024-11-29T19:36:27+00:00
+#   timestamp: 2024-12-05T13:16:41+00:00
 
 from __future__ import annotations
 

metadata/generated/schema/type/schema.py

@@ -1,6 +1,6 @@
 # generated by datamodel-codegen:
 #   filename:  type/schema.json
-#   timestamp: 2024-11-29T19:36:27+00:00
+#   timestamp: 2024-12-05T13:16:41+00:00
 
 from __future__ import annotations
 

metadata/generated/schema/type/tableQuery.py

@@ -1,6 +1,6 @@
 # generated by datamodel-codegen:
 #   filename:  type/tableQuery.json
-#   timestamp: 2024-11-29T19:36:27+00:00
+#   timestamp: 2024-12-05T13:16:41+00:00
 
 from __future__ import annotations
 

metadata/generated/schema/type/tableUsageCount.py

@@ -1,6 +1,6 @@
 # generated by datamodel-codegen:
 #   filename:  type/tableUsageCount.json
-#   timestamp: 2024-11-29T19:36:27+00:00
+#   timestamp: 2024-12-05T13:16:41+00:00
 
 from __future__ import annotations
 

metadata/generated/schema/type/tagLabel.py

@@ -1,6 +1,6 @@
 # generated by datamodel-codegen:
 #   filename:  type/tagLabel.json
-#   timestamp: 2024-11-29T19:36:27+00:00
+#   timestamp: 2024-12-05T13:16:41+00:00
 
 from __future__ import annotations
 

metadata/generated/schema/type/usageDetails.py

@@ -1,6 +1,6 @@
 # generated by datamodel-codegen:
 #   filename:  type/usageDetails.json
-#   timestamp: 2024-11-29T19:36:27+00:00
+#   timestamp: 2024-12-05T13:16:41+00:00
 
 from __future__ import annotations
 

metadata/generated/schema/type/usageRequest.py

@@ -1,6 +1,6 @@
 # generated by datamodel-codegen:
 #   filename:  type/usageRequest.json
-#   timestamp: 2024-11-29T19:36:27+00:00
+#   timestamp: 2024-12-05T13:16:41+00:00
 
 from __future__ import annotations
 

metadata/generated/schema/type/votes.py

@@ -1,6 +1,6 @@
 # generated by datamodel-codegen:
 #   filename:  type/votes.json
-#   timestamp: 2024-11-29T19:36:27+00:00
+#   timestamp: 2024-12-05T13:16:41+00:00
 
 from __future__ import annotations
 

metadata/ingestion/lineage/parser.py

@@ -15,12 +15,12 @@ import traceback
 from collections import defaultdict
 from copy import deepcopy
 from logging.config import DictConfigurator
-from typing import Any, Dict, List, Optional, Tuple
+from typing import Any, Dict, List, Optional, Tuple, Union
 
 import sqlparse
 from cached_property import cached_property
 from collate_sqllineage import SQLPARSE_DIALECT
-from collate_sqllineage.core.models import Column, Table
+from collate_sqllineage.core.models import Column, DataFunction, Table
 from collate_sqllineage.exceptions import SQLLineageException
 from collate_sqllineage.runner import LineageRunner
 from sqlparse.sql import Comparison, Identifier, Parenthesis, Statement
@@ -110,7 +110,7 @@ class LineageParser:
         return []
 
     @cached_property
-    def source_tables(self) -> List[Table]:
+    def source_tables(self) -> List[Union[Table, DataFunction]]:
         """
         Get a list of source tables
         """
@@ -374,7 +374,9 @@ class LineageParser:
         if not self._clean_query:
             return []
         return [
-            self.clean_table_name(table) for table in tables if isinstance(table, Table)
+            self.clean_table_name(table)
+            for table in tables
+            if isinstance(table, (Table, DataFunction))
         ]
 
     @classmethod

metadata/ingestion/lineage/sql_lineage.py

@@ -13,9 +13,18 @@ Helper functions to handle SQL lineage operations
 """
 import itertools
 import traceback
-from typing import Any, Iterable, List, Optional, Tuple
+from collections import defaultdict
+from typing import Any, Iterable, List, Optional, Tuple, Union
+
+from collate_sqllineage.core.models import Column, DataFunction
+from collate_sqllineage.core.models import Table as LineageTable
 
 from metadata.generated.schema.api.lineage.addLineage import AddLineageRequest
+from metadata.generated.schema.entity.data.storedProcedure import (
+    Language,
+    StoredProcedure,
+    StoredProcedureType,
+)
 from metadata.generated.schema.entity.data.table import Table
 from metadata.generated.schema.entity.services.ingestionPipelines.status import (
     StackTraceError,
@@ -163,6 +172,187 @@ def get_table_fqn_from_query_name(
     return database_query, schema_query, table
 
 
+def __process_intermediate_column_lineage(
+    intermediate_column_lineage: dict,
+    result: dict,
+    source_table: str,
+    intermediate_table: str,
+    intermediate_column: str,
+    source_column: str,
+):
+    # Check intermediate dictionary for mappings
+    for (
+        target_table,
+        target_mappings,
+    ) in intermediate_column_lineage[intermediate_table].items():
+        for inter_col, target_col in target_mappings:
+            if intermediate_column == inter_col:
+                # Append to the result dictionary
+                if target_table not in result[source_table]:
+                    result[source_table][target_table] = []
+                result[source_table][target_table].append((source_column, target_col))
+
+
+def __process_column_mappings(
+    mappings: dict, result: dict, source_table: str, intermediate_column_lineage: dict
+):
+    for intermediate_table, column_pairs in mappings.items():
+        # Iterate through each column mapping in the original dictionary
+        for source_column, intermediate_column in column_pairs:
+            if intermediate_table in intermediate_column_lineage:
+                __process_intermediate_column_lineage(
+                    intermediate_column_lineage,
+                    result,
+                    source_table,
+                    intermediate_table,
+                    intermediate_column,
+                    source_column,
+                )
+
+
+def handle_udf_column_lineage(
+    column_lineage_original: dict,
+    column_lineage_generated: List[Tuple[Column, Column]],
+):
+    """
+    Handle UDF column lineage
+    """
+    try:
+        result = defaultdict(dict)
+        intermediate_column_lineage = populate_column_lineage_map(
+            column_lineage_generated
+        )
+        # Iterate through the original dictionary
+        for source_table, mappings in column_lineage_original.items():
+            __process_column_mappings(
+                mappings, result, source_table, intermediate_column_lineage
+            )
+        column_lineage_original.update(result)
+    except Exception as exc:
+        logger.debug(traceback.format_exc())
+        logger.error(f"Error handling UDF column lineage: {exc}")
+
+
+# pylint: disable=too-many-arguments
+def __process_udf_es_results(
+    metadata: OpenMetadata,
+    dialect: Dialect,
+    source_table: Union[DataFunction, LineageTable],
+    database_name: Optional[str],
+    schema_name: Optional[str],
+    service_name: Optional[str],
+    timeout_seconds: int,
+    column_lineage: dict,
+    es_result_entities: List[StoredProcedure],
+    procedure: Optional[StoredProcedure] = None,
+):
+    for entity in es_result_entities:
+        if (
+            entity.storedProcedureType == StoredProcedureType.UDF
+            and entity.storedProcedureCode
+            and entity.storedProcedureCode.language == Language.SQL
+        ):
+            expected_table_name = str(source_table).replace(
+                f"{DEFAULT_SCHEMA_NAME}.", ""
+            )
+            lineage_parser = LineageParser(
+                f"create table {str(expected_table_name)} as {entity.storedProcedureCode.code}",
+                dialect=dialect,
+                timeout_seconds=timeout_seconds,
+            )
+            handle_udf_column_lineage(column_lineage, lineage_parser.column_lineage)
+            for source in lineage_parser.source_tables or []:
+                yield from get_source_table_names(
+                    metadata,
+                    dialect,
+                    source,
+                    database_name,
+                    schema_name,
+                    service_name,
+                    timeout_seconds,
+                    column_lineage,
+                    procedure or entity,
+                )
+
+
+def __process_udf_table_names(
+    metadata: OpenMetadata,
+    dialect: Dialect,
+    source_table: Union[DataFunction, LineageTable],
+    database_name: Optional[str],
+    schema_name: Optional[str],
+    service_name: Optional[str],
+    timeout_seconds: int,
+    column_lineage: dict,
+    procedure: Optional[StoredProcedure] = None,
+):
+    database_query, schema_query, table = get_table_fqn_from_query_name(
+        str(source_table)
+    )
+    function_fqn_string = build_es_fqn_search_string(
+        database_query or database_name,
+        schema_query or schema_name,
+        service_name,
+        table,
+    )
+    es_result_entities: Optional[List[StoredProcedure]] = metadata.es_search_from_fqn(
+        entity_type=StoredProcedure,
+        fqn_search_string=function_fqn_string,
+    )
+    if es_result_entities:
+        yield from __process_udf_es_results(
+            metadata,
+            dialect,
+            source_table,
+            database_name,
+            schema_name,
+            service_name,
+            timeout_seconds,
+            column_lineage,
+            es_result_entities,
+            procedure,
+        )
+
+
+def get_source_table_names(
+    metadata: OpenMetadata,
+    dialect: Dialect,
+    source_table: Union[DataFunction, LineageTable],
+    database_name: Optional[str],
+    schema_name: Optional[str],
+    service_name: Optional[str],
+    timeout_seconds: int,
+    column_lineage: dict,
+    procedure: Optional[StoredProcedure] = None,
+) -> Iterable[Tuple[Optional[EntityReference], str]]:
+    """
+    Get source table names from DataFunction
+    """
+    try:
+        if not isinstance(source_table, DataFunction):
+            yield EntityReference(
+                id=procedure.id.root, type="storedProcedure"
+            ) if procedure else None, str(source_table)
+        else:
+            yield from __process_udf_table_names(
+                metadata,
+                dialect,
+                source_table,
+                database_name,
+                schema_name,
+                service_name,
+                timeout_seconds,
+                column_lineage,
+                procedure,
+            )
+
+    except Exception as exc:
+        logger.debug(traceback.format_exc())
+        logger.error(
+            f"Error getting source table names for table [{source_table}]: {exc}"
+        )
+
+
 def get_table_entities_from_query(
     metadata: OpenMetadata,
     service_name: str,
@@ -244,6 +434,7 @@ def get_column_lineage(
     return column_lineage
 
 
+# pylint: disable=too-many-arguments
 def _build_table_lineage(
     from_entity: Table,
     to_entity: Table,
@@ -252,6 +443,7 @@ def _build_table_lineage(
     masked_query: str,
     column_lineage_map: dict,
     lineage_source: LineageSource = LineageSource.QueryLineage,
+    procedure: Optional[EntityReference] = None,
 ) -> Either[AddLineageRequest]:
     """
     Prepare the lineage request generator
@@ -276,7 +468,9 @@ def _build_table_lineage(
             from_table_raw_name=str(from_table_raw_name),
             column_lineage_map=column_lineage_map,
         )
-        lineage_details = LineageDetails(sqlQuery=masked_query, source=lineage_source)
+        lineage_details = LineageDetails(
+            sqlQuery=masked_query, source=lineage_source, pipeline=procedure
+        )
         if col_lineage:
             lineage_details.columnsLineage = col_lineage
         lineage = AddLineageRequest(
@@ -315,6 +509,7 @@ def _create_lineage_by_table_name(
     masked_query: str,
     column_lineage_map: dict,
     lineage_source: LineageSource = LineageSource.QueryLineage,
+    procedure: Optional[EntityReference] = None,
 ) -> Iterable[Either[AddLineageRequest]]:
     """
     This method is to create a lineage between two tables
@@ -358,6 +553,7 @@ def _create_lineage_by_table_name(
                     masked_query=masked_query,
                     column_lineage_map=column_lineage_map,
                     lineage_source=lineage_source,
+                    procedure=procedure,
                 )
 
     except Exception as exc:
@@ -429,17 +625,28 @@ def get_lineage_by_query(
 
         for intermediate_table in lineage_parser.intermediate_tables:
             for source_table in lineage_parser.source_tables:
-                yield from _create_lineage_by_table_name(
-                    metadata,
-                    from_table=str(source_table),
-                    to_table=str(intermediate_table),
-                    service_name=service_name,
+                for procedure, from_table_name in get_source_table_names(
+                    metadata=metadata,
+                    dialect=dialect,
+                    source_table=source_table,
                     database_name=database_name,
                     schema_name=schema_name,
-                    masked_query=masked_query,
-                    column_lineage_map=column_lineage,
-                    lineage_source=lineage_source,
-                )
+                    service_name=service_name,
+                    timeout_seconds=timeout_seconds,
+                    column_lineage=column_lineage,
+                ):
+                    yield from _create_lineage_by_table_name(
+                        metadata,
+                        from_table=str(from_table_name),
+                        to_table=str(intermediate_table),
+                        service_name=service_name,
+                        database_name=database_name,
+                        schema_name=schema_name,
+                        masked_query=masked_query,
+                        column_lineage_map=column_lineage,
+                        lineage_source=lineage_source,
+                        procedure=procedure,
+                    )
             for target_table in lineage_parser.target_tables:
                 yield from _create_lineage_by_table_name(
                     metadata,
@@ -455,17 +662,28 @@ def get_lineage_by_query(
         if not lineage_parser.intermediate_tables:
             for target_table in lineage_parser.target_tables:
                 for source_table in lineage_parser.source_tables:
-                    yield from _create_lineage_by_table_name(
-                        metadata,
-                        from_table=str(source_table),
-                        to_table=str(target_table),
-                        service_name=service_name,
+                    for procedure, from_table_name in get_source_table_names(
+                        metadata=metadata,
+                        dialect=dialect,
+                        source_table=source_table,
                         database_name=database_name,
                         schema_name=schema_name,
-                        masked_query=masked_query,
-                        column_lineage_map=column_lineage,
-                        lineage_source=lineage_source,
-                    )
+                        service_name=service_name,
+                        timeout_seconds=timeout_seconds,
+                        column_lineage=column_lineage,
+                    ):
+                        yield from _create_lineage_by_table_name(
+                            metadata,
+                            from_table=str(from_table_name),
+                            to_table=str(target_table),
+                            service_name=service_name,
+                            database_name=database_name,
+                            schema_name=schema_name,
+                            masked_query=masked_query,
+                            column_lineage_map=column_lineage,
+                            lineage_source=lineage_source,
+                            procedure=procedure,
+                        )
         if not lineage_parser.query_parsing_success:
             query_parsing_failures.add(
                 QueryParsingError(
@@ -505,17 +723,28 @@ def get_lineage_via_table_entity(
         to_table_name = table_entity.name.root
 
         for from_table_name in lineage_parser.source_tables:
-            yield from _create_lineage_by_table_name(
-                metadata,
-                from_table=str(from_table_name),
-                to_table=f"{schema_name}.{to_table_name}",
-                service_name=service_name,
+            for procedure, source_table in get_source_table_names(
+                metadata=metadata,
+                dialect=dialect,
+                source_table=from_table_name,
                 database_name=database_name,
                 schema_name=schema_name,
-                masked_query=masked_query,
-                column_lineage_map=column_lineage,
-                lineage_source=lineage_source,
-            ) or []
+                service_name=service_name,
+                timeout_seconds=timeout_seconds,
+                column_lineage=column_lineage,
+            ):
+                yield from _create_lineage_by_table_name(
+                    metadata,
+                    from_table=str(source_table),
+                    to_table=f"{schema_name}.{to_table_name}",
+                    service_name=service_name,
+                    database_name=database_name,
+                    schema_name=schema_name,
+                    masked_query=masked_query,
+                    column_lineage_map=column_lineage,
+                    lineage_source=lineage_source,
+                    procedure=procedure,
+                ) or []
         if not lineage_parser.query_parsing_success:
             query_parsing_failures.add(
                 QueryParsingError(

metadata/ingestion/models/patch_request.py

@@ -100,6 +100,7 @@ ALLOWED_COMMON_PATCH_FIELDS = {
     "fileFormat": True,
     # Stored Procedure Fields
     "storedProcedureCode": True,
+    "storedProcedureType": True,
     "code": True,
     # Dashboard Entity Fields
     "chartType": True,

metadata/ingestion/ometa/mixins/data_insight_mixin.py

@@ -31,6 +31,7 @@ from metadata.generated.schema.dataInsight.dataInsightChartResult import (
 )
 from metadata.generated.schema.dataInsight.kpi.basic import KpiResult
 from metadata.generated.schema.dataInsight.kpi.kpi import Kpi
+from metadata.ingestion.ometa.utils import quote
 
 
 class DataInsightMixin:
@@ -58,7 +59,7 @@ class DataInsightMixin:
             record (ReportData): report data
         """
 
-        resp = self.client.put(f"/kpi/{fqn}/kpiResult", record.model_dump_json())
+        resp = self.client.put(f"/kpi/{quote(fqn)}/kpiResult", record.model_dump_json())
 
         return resp
 
@@ -143,7 +144,7 @@ class DataInsightMixin:
         params = {"startTs": start_ts, "endTs": end_ts}
 
         resp = self.client.get(
-            f"/kpi/{fqn}/kpiResult",
+            f"/kpi/{quote(fqn)}/kpiResult",
             params,
         )
 

metadata/ingestion/ometa/mixins/ingestion_pipeline_mixin.py

@@ -22,6 +22,7 @@ from metadata.generated.schema.entity.services.ingestionPipelines.ingestionPipel
 )
 from metadata.ingestion.api.parser import parse_ingestion_pipeline_config_gracefully
 from metadata.ingestion.ometa.client import REST
+from metadata.ingestion.ometa.utils import quote
 from metadata.utils.logger import ometa_logger
 
 logger = ometa_logger()
@@ -46,7 +47,7 @@ class OMetaIngestionPipelineMixin:
         :param pipeline_status: Pipeline Status data to add
         """
         resp = self.client.put(
-            f"{self.get_suffix(IngestionPipeline)}/{ingestion_pipeline_fqn}/pipelineStatus",
+            f"{self.get_suffix(IngestionPipeline)}/{quote(ingestion_pipeline_fqn)}/pipelineStatus",
             data=pipeline_status.model_dump_json(),
         )
         logger.debug(
@@ -64,7 +65,8 @@ class OMetaIngestionPipelineMixin:
         :param pipeline_status_run_id: Pipeline Status run id
         """
         resp = self.client.get(
-            f"{self.get_suffix(IngestionPipeline)}/{ingestion_pipeline_fqn}/pipelineStatus/{pipeline_status_run_id}"
+            f"{self.get_suffix(IngestionPipeline)}/"
+            f"{quote(ingestion_pipeline_fqn)}/pipelineStatus/{pipeline_status_run_id}"
         )
         if resp:
             return PipelineStatus(**resp)
@@ -99,7 +101,7 @@ class OMetaIngestionPipelineMixin:
         params = {"startTs": start_ts, "endTs": end_ts}
 
         resp = self.client.get(
-            f"{self.get_suffix(IngestionPipeline)}/{ingestion_pipeline_fqn}/pipelineStatus",
+            f"{self.get_suffix(IngestionPipeline)}/{quote(ingestion_pipeline_fqn)}/pipelineStatus",
             data=params,
         )