openmetadata-ingestion 1.6.7.1__py3-none-any.whl → 1.7.0.0rc1__py3-none-any.whl

metadata/generated/schema/type/databaseConnectionConfig.py

@@ -1,6 +1,6 @@
 # generated by datamodel-codegen:
 #   filename:  type/databaseConnectionConfig.json
-#   timestamp: 2025-04-02T12:35:09+00:00
+#   timestamp: 2025-04-02T05:35:35+00:00
 
 from __future__ import annotations
 

metadata/generated/schema/type/entityHierarchy.py

@@ -1,6 +1,6 @@
 # generated by datamodel-codegen:
 #   filename:  type/entityHierarchy.json
-#   timestamp: 2025-04-02T12:35:09+00:00
+#   timestamp: 2025-04-02T05:35:35+00:00
 
 from __future__ import annotations
 

metadata/generated/schema/type/entityHistory.py

@@ -1,6 +1,6 @@
 # generated by datamodel-codegen:
 #   filename:  type/entityHistory.json
-#   timestamp: 2025-04-02T12:35:09+00:00
+#   timestamp: 2025-04-02T05:35:35+00:00
 
 from __future__ import annotations
 
@@ -11,6 +11,8 @@ from typing_extensions import Annotated
 
 from metadata.ingestion.models.custom_pydantic import BaseModel
 
+from . import changeSummaryMap
+
 
 class EntityVersionHistory(BaseModel):
     model_config = ConfigDict(
@@ -64,7 +66,7 @@ class FieldChange(BaseModel):
     ]
 
 
-class ChangeDescription(BaseModel):
+class IncrementalChangeDescription(BaseModel):
     model_config = ConfigDict(
         extra='forbid',
     )
@@ -95,7 +97,7 @@ class ChangeDescription(BaseModel):
     ]
 
 
-class IncrementalChangeDescription(BaseModel):
+class ChangeDescription(BaseModel):
     model_config = ConfigDict(
         extra='forbid',
     )
@@ -124,3 +126,4 @@ class IncrementalChangeDescription(BaseModel):
             description='When a change did not result in change, this could be same as the current version.',
         ),
     ]
+    changeSummary: Optional[changeSummaryMap.ChangeSummaryModel] = None

metadata/generated/schema/type/entityLineage.py

@@ -1,6 +1,6 @@
 # generated by datamodel-codegen:
 #   filename:  type/entityLineage.json
-#   timestamp: 2025-04-02T12:35:09+00:00
+#   timestamp: 2025-04-02T05:35:35+00:00
 
 from __future__ import annotations
 
@@ -49,6 +49,8 @@ class Source(Enum):
     SparkLineage = 'SparkLineage'
     OpenLineage = 'OpenLineage'
     ExternalTableLineage = 'ExternalTableLineage'
+    CrossDatabaseLineage = 'CrossDatabaseLineage'
+    ChildAssets = 'ChildAssets'
 
 
 class LineageDetails(BaseModel):
@@ -77,6 +79,30 @@ class LineageDetails(BaseModel):
             description='Lineage type describes how a lineage was created.',
         ),
     ]
+    createdAt: Annotated[
+        Optional[basic.Timestamp],
+        Field(
+            None,
+            description='Last update time corresponding to the new version of the entity in Unix epoch time milliseconds.',
+        ),
+    ]
+    createdBy: Annotated[
+        Optional[str], Field(None, description='User who created the node.')
+    ]
+    updatedAt: Annotated[
+        Optional[basic.Timestamp],
+        Field(
+            None,
+            description='Last update time corresponding to the new version of the entity in Unix epoch time milliseconds.',
+        ),
+    ]
+    updatedBy: Annotated[
+        Optional[str], Field(None, description='User who made the update.')
+    ]
+    assetEdges: Annotated[
+        Optional[int],
+        Field(None, description='Asset count in case of child assets lineage.'),
+    ]
 
 
 class Edge(BaseModel):

metadata/generated/schema/type/entityReference.py

@@ -1,6 +1,6 @@
 # generated by datamodel-codegen:
 #   filename:  type/entityReference.json
-#   timestamp: 2025-04-02T12:35:09+00:00
+#   timestamp: 2025-04-02T05:35:35+00:00
 
 from __future__ import annotations
 

metadata/generated/schema/type/entityReferenceList.py

@@ -1,6 +1,6 @@
 # generated by datamodel-codegen:
 #   filename:  type/entityReferenceList.json
-#   timestamp: 2025-04-02T12:35:09+00:00
+#   timestamp: 2025-04-02T05:35:35+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: 2025-04-02T12:35:09+00:00
+#   timestamp: 2025-04-02T05:35:35+00:00
 
 from __future__ import annotations
 

metadata/generated/schema/type/entityUsage.py

@@ -1,6 +1,6 @@
 # generated by datamodel-codegen:
 #   filename:  type/entityUsage.json
-#   timestamp: 2025-04-02T12:35:09+00:00
+#   timestamp: 2025-04-02T05:35:35+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: 2025-04-02T12:35:09+00:00
+#   timestamp: 2025-04-02T05:35:35+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: 2025-04-02T12:35:09+00:00
+#   timestamp: 2025-04-02T05:35:35+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: 2025-04-02T12:35:09+00:00
+#   timestamp: 2025-04-02T05:35:35+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: 2025-04-02T12:35:09+00:00
+#   timestamp: 2025-04-02T05:35:35+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: 2025-04-02T12:35:09+00:00
+#   timestamp: 2025-04-02T05:35:35+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: 2025-04-02T12:35:09+00:00
+#   timestamp: 2025-04-02T05:35:35+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: 2025-04-02T12:35:09+00:00
+#   timestamp: 2025-04-02T05:35:35+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: 2025-04-02T12:35:09+00:00
+#   timestamp: 2025-04-02T05:35:35+00:00
 
 from __future__ import annotations
 
@@ -53,6 +53,9 @@ class ParsedData(BaseModel):
         Optional[float],
         Field(None, description='How long did the query took to run in milliseconds.'),
     ]
+    cost: Annotated[
+        Optional[float], Field(None, description='Cost of the query execution')
+    ]
 
 
 class QueryParserData(BaseModel):

metadata/generated/schema/type/reaction.py

@@ -1,6 +1,6 @@
 # generated by datamodel-codegen:
 #   filename:  type/reaction.json
-#   timestamp: 2025-04-02T12:35:09+00:00
+#   timestamp: 2025-04-02T05:35:35+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: 2025-04-02T12:35:09+00:00
+#   timestamp: 2025-04-02T05:35:35+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: 2025-04-02T12:35:09+00:00
+#   timestamp: 2025-04-02T05:35:35+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: 2025-04-02T12:35:09+00:00
+#   timestamp: 2025-04-02T05:35:35+00:00
 
 from __future__ import annotations
 
@@ -58,6 +58,9 @@ class TableQuery(BaseModel):
         Optional[float],
         Field(None, description='How long did the query took to run in milliseconds.'),
     ]
+    cost: Annotated[
+        Optional[float], Field(None, description='Cost of the query execution')
+    ]
 
 
 class TableQueries(BaseModel):

metadata/generated/schema/type/tableUsageCount.py

@@ -1,6 +1,6 @@
 # generated by datamodel-codegen:
 #   filename:  type/tableUsageCount.json
-#   timestamp: 2025-04-02T12:35:09+00:00
+#   timestamp: 2025-04-02T05:35:35+00:00
 
 from __future__ import annotations
 
@@ -38,6 +38,23 @@ class TableColumnJoin(BaseModel):
     ]
 
 
+class QueryCostWrapper(BaseModel):
+    model_config = ConfigDict(
+        extra='forbid',
+    )
+    cost: Annotated[float, Field(description='Avg query cost per execution')]
+    count: Annotated[float, Field(description='Number of times the query was executed')]
+    date: Annotated[
+        Optional[str], Field(None, description='Date of execution of SQL query')
+    ]
+    queryHash: Annotated[Optional[str], Field(None, description='Hash of the query')]
+    query: Annotated[Optional[str], Field(None, description='SQL query')]
+    dialect: Annotated[Optional[str], Field(None, description='SQL dialect')]
+    totalDuration: Annotated[
+        Optional[float], Field(None, description='Total duration of the query')
+    ]
+
+
 class TableUsageCount(BaseModel):
     model_config = ConfigDict(
         extra='forbid',

metadata/generated/schema/type/tagLabel.py

@@ -1,6 +1,6 @@
 # generated by datamodel-codegen:
 #   filename:  type/tagLabel.json
-#   timestamp: 2025-04-02T12:35:09+00:00
+#   timestamp: 2025-04-02T05:35:35+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: 2025-04-02T12:35:09+00:00
+#   timestamp: 2025-04-02T05:35:35+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: 2025-04-02T12:35:09+00:00
+#   timestamp: 2025-04-02T05:35:35+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: 2025-04-02T12:35:09+00:00
+#   timestamp: 2025-04-02T05:35:35+00:00
 
 from __future__ import annotations
 

metadata/ingestion/api/topology_runner.py

@@ -418,9 +418,9 @@ class TopologyRunnerMixin(Generic[C]):
             if not entity:
                 # Safe access to Entity Request name
                 raise MissingExpectedEntityAckException(
-                    f"Missing ack back from [{stage.type_.__name__}: {entity_fqn}] - "
-                    "Possible causes are changes in the server Fernet key or mismatched JSON Schemas "
-                    "for the service connection."
+                    f"We are trying to create a [{stage.type_.__name__}] with FQN [{entity_fqn}],"
+                    " but we got no Entity back from the API. Checking for errors in the OpenMetadata Sink could help"
+                    " validate if the Entity was properly created or not."
                 )
 
         self.context.get().update_context_name(stage=stage, right=right)

metadata/ingestion/bulksink/metadata_usage.py

@@ -41,7 +41,11 @@ from metadata.generated.schema.entity.services.ingestionPipelines.status import
 from metadata.generated.schema.entity.teams.user import User
 from metadata.generated.schema.type.basic import Timestamp
 from metadata.generated.schema.type.lifeCycle import AccessDetails, LifeCycle
-from metadata.generated.schema.type.tableUsageCount import TableColumn, TableUsageCount
+from metadata.generated.schema.type.tableUsageCount import (
+    QueryCostWrapper,
+    TableColumn,
+    TableUsageCount,
+)
 from metadata.generated.schema.type.usageRequest import UsageRequest
 from metadata.ingestion.api.steps import BulkSink
 from metadata.ingestion.lineage.sql_lineage import (
@@ -163,7 +167,7 @@ class MetadataUsageBulkSink(BulkSink):
                     )
                 )
 
-    def iterate_files(self):
+    def iterate_files(self, usage_files: bool = True):
         """
         Iterate through files in the given directory
         """
@@ -173,11 +177,16 @@ class MetadataUsageBulkSink(BulkSink):
                 full_file_name = os.path.join(self.config.filename, filename)
                 if not os.path.isfile(full_file_name):
                     continue
-                with open(full_file_name, encoding=UTF_8) as file:
-                    yield file
+                # if usage_files is True, then we want to iterate through files does not end with query
+                # if usage_files is False, then we want to iterate through files that end with query
+                if filename.endswith("query") ^ usage_files:
+                    with open(full_file_name, encoding=UTF_8) as file:
+                        yield file
 
-    # Check here how to properly pick up ES and/or table query data
-    def run(self) -> None:
+    def handle_table_usage(self) -> None:
+        """
+        Handle table usage.
+        """
         for file_handler in self.iterate_files():
             self.table_usage_map = {}
             for usage_record in file_handler.readlines():
@@ -210,6 +219,18 @@ class MetadataUsageBulkSink(BulkSink):
 
             self.__publish_usage_records()
 
+    def handle_query_cost(self) -> None:
+        for file_handler in self.iterate_files(usage_files=False):
+            for usage_record in file_handler.readlines():
+                record = json.loads(usage_record)
+                cost_record = QueryCostWrapper(**record)
+                self.metadata.publish_query_cost(cost_record, self.service_name)
+
+    # Check here how to properly pick up ES and/or table query data
+    def run(self) -> None:
+        self.handle_table_usage()
+        self.handle_query_cost()
+
     def get_table_usage_and_joins(
         self, table_entities: List[Table], table_usage: TableUsageCount
     ):

metadata/ingestion/lineage/masker.py

@@ -14,6 +14,7 @@ Query masking utilities
 
 import traceback
 
+from cachetools import LRUCache
 from collate_sqllineage.runner import SQLPARSE_DIALECT, LineageRunner
 from sqlparse.sql import Comparison
 from sqlparse.tokens import Literal, Number, String
@@ -22,6 +23,8 @@ from metadata.ingestion.lineage.models import Dialect
 
 MASK_TOKEN = "?"
 
+# Cache size is 128 to avoid memory issues
+masked_query_cache = LRUCache(maxsize=128)
 
 # pylint: disable=protected-access
 def get_logger():
@@ -112,8 +115,13 @@ def mask_literals_with_sqlfluff(query: str, parser: LineageRunner) -> str:
 def mask_query(
     query: str, dialect: str = Dialect.ANSI.value, parser: LineageRunner = None
 ) -> str:
+    """
+    Mask a query using sqlparse or sqlfluff.
+    """
     logger = get_logger()
     try:
+        if masked_query_cache.get((query, dialect)):
+            return masked_query_cache.get((query, dialect))
         if not parser:
             try:
                 parser = LineageRunner(query, dialect=dialect)
@@ -122,8 +130,11 @@ def mask_query(
                 parser = LineageRunner(query)
                 len(parser.source_tables)
         if parser._dialect == SQLPARSE_DIALECT:
-            return mask_literals_with_sqlparse(query, parser)
-        return mask_literals_with_sqlfluff(query, parser)
+            masked_query = mask_literals_with_sqlparse(query, parser)
+        else:
+            masked_query = mask_literals_with_sqlfluff(query, parser)
+        masked_query_cache[(query, dialect)] = masked_query
+        return masked_query
     except Exception as exc:
         logger.debug(f"Failed to mask query with sqlfluff: {exc}")
         logger.debug(traceback.format_exc())

metadata/ingestion/models/custom_pydantic.py

@@ -17,12 +17,13 @@ be self-sufficient with only pydantic at import time.
 """
 import json
 import logging
-from typing import Any, Dict, Literal, Optional, Union
+from typing import Any, Callable, Dict, Literal, Optional, Union
 
 from pydantic import BaseModel as PydanticBaseModel
-from pydantic import PlainSerializer, model_validator
+from pydantic import WrapSerializer, model_validator
 from pydantic.main import IncEx
 from pydantic.types import SecretStr
+from pydantic_core.core_schema import SerializationInfo
 from typing_extensions import Annotated
 
 from metadata.ingestion.models.custom_basemodel_validation import (
@@ -69,9 +70,10 @@ class BaseModel(PydanticBaseModel):
             raise exc
         return values
 
-    def model_dump_json(  # pylint: disable=too-many-arguments
+    def model_dump_json(  # pylint: disable=too-many-arguments,unused-argument
         self,
         *,
+        mask_secrets: Optional[bool] = None,
         indent: Optional[int] = None,
         include: IncEx = None,
         exclude: IncEx = None,
@@ -82,6 +84,7 @@ class BaseModel(PydanticBaseModel):
         exclude_none: bool = True,
         round_trip: bool = False,
         warnings: Union[bool, Literal["none", "warn", "error"]] = True,
+        fallback: Optional[Callable[[Any], Any]] = None,
         serialize_as_any: bool = False,
     ) -> str:
         """
@@ -92,10 +95,20 @@ class BaseModel(PydanticBaseModel):
 
         This solution is covered in the `test_pydantic_v2` test comparing the
         dump results from V1 vs. V2.
+
+        mask_secrets: bool - Can be overridedn by either passing it as an argument or setting it in the context.
+        With the following rules:
+            - if mask_secrets is not None, it will be used as is
+            - if mask_secrets is None and context is not None, it will be set to context.get("mask_secrets", True)
+            - if mask_secrets is None and context is None, it will be set to True
+
         """
+        if mask_secrets is None:
+            mask_secrets = context.get("mask_secrets", True) if context else True
         return json.dumps(
             self.model_dump(
                 mode="json",
+                mask_secrets=mask_secrets,
                 include=include,
                 exclude=exclude,
                 context=context,
@@ -110,6 +123,18 @@ class BaseModel(PydanticBaseModel):
             ensure_ascii=True,
         )
 
+    def model_dump(
+        self,
+        *,
+        mask_secrets: bool = False,
+        **kwargs,
+    ) -> Dict[str, Any]:
+        if mask_secrets:
+            context = kwargs.pop("context", None) or {}
+            context["mask_secrets"] = True
+            kwargs["context"] = context
+        return super().model_dump(**kwargs)
+
 
 class _CustomSecretStr(SecretStr):
     """
@@ -117,6 +142,9 @@ class _CustomSecretStr(SecretStr):
 
     If the secret string value starts with `config:` it will use the rest of the string as secret id to search for it
     in the secrets store.
+
+    By default the secrets will be unmasked when dumping ot python objects and masked when dumping to json unless
+    explicitly set otherwise using the `mask_secrets` or `context` arguments.
     """
 
     def __repr__(self) -> str:
@@ -154,9 +182,19 @@ class _CustomSecretStr(SecretStr):
         return self._secret_value
 
 
-CustomSecretStr = Annotated[
-    _CustomSecretStr, PlainSerializer(lambda secret: secret.get_secret_value())
-]
+def handle_secret(value: Any, handler, info: SerializationInfo) -> str:
+    """
+    Handle the secret value in the model.
+    """
+    if not (info.context is not None and info.context.get("mask_secrets", False)):
+        if info.mode == "json":
+            # short circuit the json serialization and return the actual value
+            return value.get_secret_value()
+        return handler(value.get_secret_value())
+    return str(value)  # use pydantic's logic to mask the secret
+
+
+CustomSecretStr = Annotated[_CustomSecretStr, WrapSerializer(handle_secret)]
 
 
 def ignore_type_decoder(type_: Any) -> None:

metadata/ingestion/ometa/client.py

@@ -14,6 +14,7 @@ Python API REST wrapper and helpers
 import time
 import traceback
 from datetime import datetime, timezone
+from json import JSONDecodeError
 from typing import Any, Callable, Dict, List, Optional, Union
 
 import requests
@@ -117,6 +118,8 @@ class ClientConfig(ConfigModel):
     verify: Optional[Union[bool, str]] = None
     cookies: Optional[Any] = None
     ttl_cache: int = 60
+    timeout: Optional[int] = None
+    cert: Optional[Union[str, tuple]] = None
 
 
 # pylint: disable=too-many-instance-attributes
@@ -140,6 +143,8 @@ class REST:
         self._auth_token_mode = self.config.auth_token_mode
         self._verify = self.config.verify
         self._cookies = self.config.cookies
+        self._cert = self.config.cert
+        self._timeout = self.config.timeout
 
         self._limits_reached = TTLCache(config.ttl_cache)
 
@@ -177,6 +182,7 @@ class REST:
                     self.config.expires_in = (
                         datetime.now(timezone.utc).timestamp() + expiry - 120
                     )
+
         if self.config.auth_header:
             headers[self.config.auth_header] = (
                 f"{self._auth_token_mode} {self.config.access_token}"
@@ -189,6 +195,7 @@ class REST:
         # the value will be set to that value.
         # Example: "Proxy-Authorization": "%(Authorization)s"
         # This will result in the Authorization value being set for the Proxy-Authorization Extra Header
+        # Any header which is comming as extra header from client will overwrite the header with same name in headers
         if self.config.extra_headers:
             extra_headers: Dict[str, str] = self.config.extra_headers
             extra_headers = {k: (v % headers) for k, v in extra_headers.items()}
@@ -210,6 +217,12 @@ class REST:
         if json:
             opts["json"] = json
 
+        if self._cert:
+            opts["cert"] = self._cert
+
+        if self._timeout:
+            opts["timeout"] = self._timeout
+
         total_retries = self._retry if self._retry > 0 else 0
         retry = total_retries
         while retry >= 0:
@@ -250,6 +263,12 @@ class REST:
             if resp.text != "":
                 try:
                     return resp.json()
+                except JSONDecodeError as json_decode_error:
+                    logger.error(
+                        f"Json decoding error while returning response {resp} in json format - {json_decode_error}."
+                        f"The Response still returned to be handled by client..."
+                    )
+                    return resp
                 except Exception as exc:
                     logger.debug(traceback.format_exc())
                     logger.warning(
@@ -308,6 +327,7 @@ class REST:
         Parameters:
             path (str):
             data ():
+            json ():
 
         Returns:
             Response