openmetadata-ingestion 1.5.13.0__py3-none-any.whl → 1.6.0.0rc2__py3-none-any.whl

metadata/ingestion/lineage/sql_lineage.py

@@ -13,17 +13,9 @@ Helper functions to handle SQL lineage operations
 """
 import itertools
 import traceback
-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 typing import Any, Iterable, List, Optional, Tuple
 
 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,
@@ -171,116 +163,6 @@ def get_table_fqn_from_query_name(
     return database_query, schema_query, table
 
 
-def handle_udf_column_lineage(
-    column_lineage_original: dict,
-    column_lineage_generated: List[Tuple[Column, Column]],
-):
-    """
-    Handle UDF column lineage
-    """
-    try:
-        result = {}
-        intermediate_column_lineage = populate_column_lineage_map(
-            column_lineage_generated
-        )
-        # Iterate through the original dictionary
-        for source_table, mappings in column_lineage_original.items():
-            result[source_table] = {}
-            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:
-                        # 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)
-                                    )
-        column_lineage_original.update(result)
-    except Exception as exc:
-        logger.debug(traceback.format_exc())
-        logger.error(f"Error handling UDF column lineage: {exc}")
-
-
-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:
-            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:
-                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,
-                            )
-    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,
@@ -370,7 +252,6 @@ 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
@@ -395,9 +276,7 @@ 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, pipeline=procedure
-        )
+        lineage_details = LineageDetails(sqlQuery=masked_query, source=lineage_source)
         if col_lineage:
             lineage_details.columnsLineage = col_lineage
         lineage = AddLineageRequest(
@@ -436,7 +315,6 @@ 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
@@ -480,7 +358,6 @@ 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:
@@ -552,28 +429,17 @@ def get_lineage_by_query(
 
         for intermediate_table in lineage_parser.intermediate_tables:
             for source_table in lineage_parser.source_tables:
-                for procedure, from_table_name in get_source_table_names(
-                    metadata=metadata,
-                    dialect=dialect,
-                    source_table=source_table,
+                yield from _create_lineage_by_table_name(
+                    metadata,
+                    from_table=str(source_table),
+                    to_table=str(intermediate_table),
+                    service_name=service_name,
                     database_name=database_name,
                     schema_name=schema_name,
-                    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,
-                    )
+                    masked_query=masked_query,
+                    column_lineage_map=column_lineage,
+                    lineage_source=lineage_source,
+                )
             for target_table in lineage_parser.target_tables:
                 yield from _create_lineage_by_table_name(
                     metadata,
@@ -589,28 +455,17 @@ 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:
-                    for procedure, from_table_name in get_source_table_names(
-                        metadata=metadata,
-                        dialect=dialect,
-                        source_table=source_table,
+                    yield from _create_lineage_by_table_name(
+                        metadata,
+                        from_table=str(source_table),
+                        to_table=str(target_table),
+                        service_name=service_name,
                         database_name=database_name,
                         schema_name=schema_name,
-                        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,
-                        )
+                        masked_query=masked_query,
+                        column_lineage_map=column_lineage,
+                        lineage_source=lineage_source,
+                    )
         if not lineage_parser.query_parsing_success:
             query_parsing_failures.add(
                 QueryParsingError(
@@ -650,28 +505,17 @@ def get_lineage_via_table_entity(
         to_table_name = table_entity.name.root
 
         for from_table_name in lineage_parser.source_tables:
-            for procedure, source_table in get_source_table_names(
-                metadata=metadata,
-                dialect=dialect,
-                source_table=from_table_name,
+            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,
                 database_name=database_name,
                 schema_name=schema_name,
-                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 []
+                masked_query=masked_query,
+                column_lineage_map=column_lineage,
+                lineage_source=lineage_source,
+            ) or []
         if not lineage_parser.query_parsing_success:
             query_parsing_failures.add(
                 QueryParsingError(

metadata/ingestion/models/custom_basemodel_validation.py

@@ -0,0 +1,80 @@
+#  Copyright 2022 Collate
+#  Licensed under the Apache License, Version 2.0 (the "License");
+#  you may not use this file except in compliance with the License.
+#  You may obtain a copy of the License at
+#  http://www.apache.org/licenses/LICENSE-2.0
+#  Unless required by applicable law or agreed to in writing, software
+#  distributed under the License is distributed on an "AS IS" BASIS,
+#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+#  See the License for the specific language governing permissions and
+#  limitations under the License.
+"""
+Validation logic for Custom Pydantic BaseModel
+"""
+
+import logging
+
+logger = logging.getLogger("metadata")
+
+
+RESTRICTED_KEYWORDS = ["::", ">"]
+RESERVED_COLON_KEYWORD = "__reserved__colon__"
+RESERVED_ARROW_KEYWORD = "__reserved__arrow__"
+
+CREATE_ADJACENT_MODELS = {"ProfilerResponse", "SampleData"}
+NAME_FIELDS = {"EntityName", "str", "ColumnName", "TableData"}
+FETCH_MODELS = {"Table", "CustomColumnName"}
+FIELD_NAMES = {"name", "columns", "root"}
+
+
+def revert_separators(value):
+    return value.replace(RESERVED_COLON_KEYWORD, "::").replace(
+        RESERVED_ARROW_KEYWORD, ">"
+    )
+
+
+def replace_separators(value):
+    return value.replace("::", RESERVED_COLON_KEYWORD).replace(
+        ">", RESERVED_ARROW_KEYWORD
+    )
+
+
+def validate_name_and_transform(values, modification_method, field_name: str = None):
+    """
+    Validate the name and transform it if needed.
+    """
+    if isinstance(values, str) and field_name in FIELD_NAMES:
+        values = modification_method(values)
+    elif (
+        hasattr(values, "root")
+        and isinstance(values.root, str)
+        and field_name in FIELD_NAMES
+    ):
+        values.root = modification_method(values.root)
+    elif hasattr(values, "model_fields"):
+        for key in type(values).model_fields.keys():
+            if getattr(values, key):
+                if getattr(values, key).__class__.__name__ in NAME_FIELDS:
+                    setattr(
+                        values,
+                        key,
+                        validate_name_and_transform(
+                            getattr(values, key),
+                            modification_method=modification_method,
+                            field_name=key,
+                        ),
+                    )
+                elif isinstance(getattr(values, key), list):
+                    setattr(
+                        values,
+                        key,
+                        [
+                            validate_name_and_transform(
+                                item,
+                                modification_method=modification_method,
+                                field_name=key,
+                            )
+                            for item in getattr(values, key)
+                        ],
+                    )
+    return values

metadata/ingestion/models/custom_pydantic.py

@@ -20,11 +20,19 @@ import logging
 from typing import Any, Dict, Literal, Optional, Union
 
 from pydantic import BaseModel as PydanticBaseModel
-from pydantic import PlainSerializer
+from pydantic import PlainSerializer, model_validator
 from pydantic.main import IncEx
 from pydantic.types import SecretStr
 from typing_extensions import Annotated
 
+from metadata.ingestion.models.custom_basemodel_validation import (
+    CREATE_ADJACENT_MODELS,
+    FETCH_MODELS,
+    replace_separators,
+    revert_separators,
+    validate_name_and_transform,
+)
+
 logger = logging.getLogger("metadata")
 
 SECRET = "secret:"
@@ -37,6 +45,30 @@ class BaseModel(PydanticBaseModel):
     Specified as `--base-class BASE_CLASS` in the generator.
     """
 
+    @model_validator(mode="after")
+    @classmethod
+    def parse_name(cls, values):  # pylint: disable=inconsistent-return-statements
+        """
+        Primary entry point to process values based on their class.
+        """
+
+        if not values:
+            return
+
+        try:
+
+            if cls.__name__ in CREATE_ADJACENT_MODELS or cls.__name__.startswith(
+                "Create"
+            ):
+                values = validate_name_and_transform(values, replace_separators)
+            elif cls.__name__ in FETCH_MODELS:
+                values = validate_name_and_transform(values, revert_separators)
+
+        except Exception as exc:
+            logger.warning("Exception while parsing Basemodel: %s", exc)
+            raise exc
+        return values
+
     def model_dump_json(  # pylint: disable=too-many-arguments
         self,
         *,

metadata/ingestion/models/custom_types.py

@@ -14,6 +14,7 @@ Custom defined types
 
 from typing import NewType, Union
 
+from metadata.generated.schema.entity.services.apiService import ApiService
 from metadata.generated.schema.entity.services.dashboardService import DashboardService
 from metadata.generated.schema.entity.services.databaseService import DatabaseService
 from metadata.generated.schema.entity.services.messagingService import MessagingService
@@ -25,6 +26,7 @@ from metadata.generated.schema.entity.services.pipelineService import PipelineSe
 ServiceWithConnectionType = NewType(
     "ServiceWithConnectionType",
     Union[
+        ApiService,
         DashboardService,
         DatabaseService,
         MessagingService,

metadata/ingestion/models/entity_interface.py

@@ -0,0 +1,71 @@
+#  Copyright 2021 Collate
+#  Licensed under the Apache License, Version 2.0 (the "License");
+#  you may not use this file except in compliance with the License.
+#  You may obtain a copy of the License at
+#  http://www.apache.org/licenses/LICENSE-2.0
+#  Unless required by applicable law or agreed to in writing, software
+#  distributed under the License is distributed on an "AS IS" BASIS,
+#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+#  See the License for the specific language governing permissions and
+#  limitations under the License.
+#  pylint: disable=invalid-name
+"""
+Entity interface model
+"""
+
+from typing import List, Optional, Protocol, runtime_checkable
+
+from metadata.generated.schema.type import basic, entityHistory, tagLabel
+
+
+@runtime_checkable
+class EntityInterface(Protocol):
+    """Entity interface model use where entity classes are used for structural typing"""
+
+    @property
+    def id(self) -> basic.Uuid:
+        ...
+
+    @property
+    def description(self) -> Optional[basic.Markdown]:
+        ...
+
+    @property
+    def displayName(self) -> Optional[str]:
+        ...
+
+    @property
+    def name(self) -> basic.EntityName:
+        ...
+
+    @property
+    def version(self) -> Optional[entityHistory.EntityVersion]:
+        ...
+
+    @property
+    def updatedBy(self) -> Optional[str]:
+        ...
+
+    @property
+    def updatedAt(self) -> Optional[basic.Timestamp]:
+        ...
+
+    @property
+    def href(self) -> Optional[basic.Href]:
+        ...
+
+    @property
+    def changeDescription(self) -> Optional[entityHistory.ChangeDescription]:
+        ...
+
+    @property
+    def fullyQualifiedName(self) -> Optional[basic.FullyQualifiedEntityName]:
+        ...
+
+
+class EntityInterfaceWithTags(EntityInterface, Protocol):
+    """Entity interface model with tags"""
+
+    @property
+    def tags(self) -> Optional[List[tagLabel.TagLabel]]:
+        ...

metadata/ingestion/models/patch_request.py

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

metadata/ingestion/ometa/mixins/es_mixin.py

@@ -24,9 +24,11 @@ from typing_extensions import Annotated
 
 from metadata.generated.schema.entity.data.container import Container
 from metadata.generated.schema.entity.data.query import Query
+from metadata.generated.schema.entity.data.table import Table, TableType
 from metadata.ingestion.models.custom_pydantic import BaseModel
 from metadata.ingestion.ometa.client import REST, APIError
 from metadata.ingestion.ometa.utils import quote
+from metadata.ingestion.source.models import TableView
 from metadata.utils.elasticsearch import ES_INDEX_MAP
 from metadata.utils.logger import ometa_logger
 
@@ -90,7 +92,7 @@ class ESMixin(Generic[T]):
 
     # sort_field needs to be unique for the pagination to work, so we can use the FQN
     paginate_query = (
-        "/search/query?q=&size={size}&deleted=false{filter}&index={index}"
+        "/search/query?q=&size={size}&deleted=false{filter}&index={index}{include_fields}"
         "&sort_field=fullyQualifiedName{after}"
     )
 
@@ -301,13 +303,19 @@ class ESMixin(Generic[T]):
             logger.warning(f"Unknown error extracting results from ES query [{err}]")
             return None
 
-    def paginate_es(
+    def _get_include_fields_query(self, fields: Optional[List[str]]) -> str:
+        """Get the include fields query"""
+        if fields:
+            return "&include_source_fields=" + "&include_source_fields=".join(fields)
+        return ""
+
+    def _paginate_es_internal(
         self,
         entity: Type[T],
         query_filter: Optional[str] = None,
         size: int = 100,
-        fields: Optional[List[str]] = None,
-    ) -> Iterator[T]:
+        include_fields: Optional[List[str]] = None,
+    ) -> Iterator[ESResponse]:
         """Paginate through the ES results, ignoring individual errors"""
         after: Optional[str] = None
         error_pages = 0
@@ -316,6 +324,7 @@ class ESMixin(Generic[T]):
             index=ES_INDEX_MAP[entity.__name__],
             filter="&query_filter=" + quote_plus(query_filter) if query_filter else "",
             size=size,
+            include_fields=self._get_include_fields_query(include_fields),
         )
         while True:
             query_string = query(
@@ -330,10 +339,7 @@ class ESMixin(Generic[T]):
                     continue
                 else:
                     break
-
-            yield from self._yield_hits_from_api(
-                response=response, entity=entity, fields=fields
-            )
+            yield response
 
             # Get next page
             last_hit = response.hits.hits[-1] if response.hits.hits else None
@@ -343,6 +349,18 @@ class ESMixin(Generic[T]):
 
             after = ",".join(last_hit.sort)
 
+    def paginate_es(
+        self,
+        entity: Type[T],
+        query_filter: Optional[str] = None,
+        size: int = 100,
+        fields: Optional[List[str]] = None,
+    ) -> Iterator[T]:
+        for response in self._paginate_es_internal(entity, query_filter, size):
+            yield from self._yield_hits_from_api(
+                response=response, entity=entity, fields=fields
+            )
+
     def _get_es_response(self, query_string: str) -> Optional[ESResponse]:
         """Get the Elasticsearch response"""
         try:
@@ -369,3 +387,43 @@ class ESMixin(Generic[T]):
                 logger.warning(
                     f"Error while getting {hit.source['fullyQualifiedName']} - {exc}"
                 )
+
+    def yield_es_view_def(
+        self,
+        service_name: str,
+    ) -> Iterable[TableView]:
+        """
+        Get the view definition from ES
+        """
+
+        from metadata.utils import fqn
+
+        query = {
+            "query": {
+                "bool": {
+                    "must": [
+                        {"term": {"service.name.keyword": service_name}},
+                        {"term": {"tableType": TableType.View.value}},
+                        {"term": {"deleted": False}},
+                        {"exists": {"field": "schemaDefinition"}},
+                    ]
+                }
+            }
+        }
+        query = json.dumps(query)
+        for response in self._paginate_es_internal(
+            entity=Table,
+            query_filter=query,
+            include_fields=["schemaDefinition", "fullyQualifiedName"],
+        ):
+            for hit in response.hits.hits:
+                _, database_name, schema_name, table_name = fqn.split(
+                    hit.source["fullyQualifiedName"]
+                )
+                yield TableView(
+                    view_definition=hit.source["schemaDefinition"],
+                    service_name=service_name,
+                    db_name=database_name,
+                    schema_name=schema_name,
+                    table_name=table_name,
+                )

metadata/ingestion/ometa/models.py

@@ -30,3 +30,4 @@ class EntityList(BaseModel, Generic[T]):
     entities: List[T]
     total: int
     after: Optional[str] = None
+    before: Optional[str] = None

metadata/ingestion/ometa/ometa_api.py

@@ -378,12 +378,13 @@ class OpenMetadata(
         logger.debug("Cannot find the Entity %s", fqn)
         return None
 
-    # pylint: disable=too-many-locals
+    # pylint: disable=too-many-locals, too-many-arguments
     def list_entities(
         self,
         entity: Type[T],
         fields: Optional[List[str]] = None,
         after: Optional[str] = None,
+        before: Optional[str] = None,
         limit: int = 100,
         params: Optional[Dict[str, str]] = None,
         skip_on_failure: bool = False,
@@ -395,9 +396,10 @@ class OpenMetadata(
         suffix = self.get_suffix(entity)
         url_limit = f"?limit={limit}"
         url_after = f"&after={after}" if after else ""
+        url_before = f"&before={before}" if before else ""
         url_fields = f"&fields={','.join(fields)}" if fields else ""
         resp = self.client.get(
-            path=f"{suffix}{url_limit}{url_after}{url_fields}", data=params
+            path=f"{suffix}{url_limit}{url_after}{url_before}{url_fields}", data=params
         )
 
         if self._use_raw_data:
@@ -421,7 +423,8 @@ class OpenMetadata(
 
         total = resp["paging"]["total"]
         after = resp["paging"]["after"] if "after" in resp["paging"] else None
-        return EntityList(entities=entities, total=total, after=after)
+        before = resp["paging"]["before"] if "before" in resp["paging"] else None
+        return EntityList(entities=entities, total=total, after=after, before=before)
 
     def list_all_entities(
         self,