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

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

metadata/ingestion/ometa/mixins/domain_mixin.py

@@ -0,0 +1,80 @@
+#  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.
+"""Domain and Data Product specific operations"""
+from typing import Dict, List
+
+from metadata.generated.schema.type.entityReference import EntityReference
+from metadata.ingestion.models.custom_pydantic import BaseModel
+from metadata.ingestion.ometa.client import REST
+
+
+class AssetsRequest(BaseModel):
+    """Request to add assets to a data product"""
+
+    assets: List[EntityReference]
+
+
+class OMetaDomainMixin:
+    """Mixin class containing Data Product specific methods"""
+
+    client: REST
+
+    def add_assets_to_data_product(
+        self, name: str, assets: List[EntityReference]
+    ) -> Dict:
+        """
+        Add assets to a data product
+
+        Args:
+            name: Name of the data product
+            assets: List of entity references to add as assets
+
+        Returns:
+            API response as a dictionary
+        """
+        return self._handle_data_product_assets(name, assets, "add")
+
+    def remove_assets_from_data_product(
+        self, name: str, assets: List[EntityReference]
+    ) -> Dict:
+        """
+        Remove assets from a data product
+
+        Args:
+            name: Name of the data product
+            assets: List of entity references to remove from assets
+
+        Returns:
+            API response as a dictionary
+        """
+        return self._handle_data_product_assets(name, assets, "remove")
+
+    def _handle_data_product_assets(
+        self,
+        name: str,
+        assets: List[EntityReference],
+        operation: str,
+    ) -> Dict:
+        """
+        Handle adding or removing assets from a data product
+
+        Args:
+            name: Name of the data product
+            assets: List of entity references to add/remove
+            operation: Operation to perform ("add" or "remove")
+
+        Returns:
+            API response as a dictionary
+        """
+        path = f"/dataProducts/{name}/assets/{operation}"
+        payload = AssetsRequest(assets=assets)
+
+        return self.client.put(path, payload.model_dump_json())

metadata/ingestion/ometa/mixins/es_mixin.py

@@ -16,7 +16,17 @@ To be used by OpenMetadata class
 import functools
 import json
 import traceback
-from typing import Generic, Iterable, Iterator, List, Optional, Set, Type, TypeVar
+from typing import (
+    Generic,
+    Iterable,
+    Iterator,
+    List,
+    Optional,
+    Set,
+    Type,
+    TypeVar,
+    Union,
+)
 from urllib.parse import quote_plus
 
 from pydantic import Field
@@ -29,7 +39,7 @@ 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.elasticsearch import ES_INDEX_MAP, get_entity_from_es_result
 from metadata.utils.logger import ometa_logger
 
 logger = ometa_logger()
@@ -436,6 +446,11 @@ class ESMixin(Generic[T]):
                                                         "tableType": TableType.Dynamic.value
                                                     }
                                                 },
+                                                {
+                                                    "term": {
+                                                        "tableType": TableType.Stream.value
+                                                    }
+                                                },
                                             ]
                                         }
                                     }
@@ -487,3 +502,28 @@ class ESMixin(Generic[T]):
                         schema_name=schema_name,
                         table_name=table_name,
                     )
+
+    def search_in_any_service(
+        self,
+        entity_type: Type[T],
+        fqn_search_string: str,
+        fetch_multiple_entities: bool = False,
+    ) -> Optional[Union[List[Table], Table]]:
+        """
+        fetch table from es when with/without `db_service_name`
+        """
+        try:
+            entity_result = get_entity_from_es_result(
+                entity_list=self.es_search_from_fqn(
+                    entity_type=entity_type,
+                    fqn_search_string=fqn_search_string,
+                ),
+                fetch_multiple_entities=fetch_multiple_entities,
+            )
+            return entity_result
+        except Exception as exc:
+            logger.debug(
+                f"Error to fetch entity: fqn={fqn_search_string} from es: {exc}"
+            )
+            logger.debug(traceback.format_exc())
+        return None

metadata/ingestion/ometa/mixins/lineage_mixin.py

@@ -171,13 +171,9 @@ class OMetaLineageMixin(Generic[T]):
 
         except APIError as err:
             logger.debug(traceback.format_exc())
-            logger.error(
-                "Error %s trying to PUT lineage for %s: %s",
-                err.status_code,
-                data.model_dump_json(),
-                str(err),
-            )
-            raise err
+            error = f"Error {err.status_code} trying to PUT lineage for {data.model_dump_json()}: {str(err)}"
+            logger.error(error)
+            return {"error": error}
 
         from_entity_lineage = self.get_lineage_by_id(
             data.edge.fromEntity.type, str(data.edge.fromEntity.id.root)
@@ -397,6 +393,10 @@ class OMetaLineageMixin(Generic[T]):
                     resp = self.add_lineage(
                         lineage_request.right, check_patch=check_patch
                     )
+                    if resp.get("error"):
+                        logger.error(resp["error"])
+                        continue
+
                     entity_name = resp.get("entity", {}).get("name")
                     for node in resp.get("nodes", []):
                         logger.info(

metadata/ingestion/ometa/mixins/patch_mixin.py

@@ -29,6 +29,9 @@ from metadata.generated.schema.entity.domains.domain import Domain
 from metadata.generated.schema.entity.services.connections.testConnectionResult import (
     TestConnectionResult,
 )
+from metadata.generated.schema.entity.services.ingestionPipelines.reverseIngestionResponse import (
+    ReverseIngestionResponse,
+)
 from metadata.generated.schema.tests.testCase import TestCase, TestCaseParameterValue
 from metadata.generated.schema.type.basic import EntityLink, Markdown
 from metadata.generated.schema.type.entityReference import EntityReference
@@ -488,7 +491,7 @@ class OMetaPatchMixin(OMetaPatchMixinBase):
     def patch_automation_workflow_response(
         self,
         automation_workflow: AutomationWorkflow,
-        test_connection_result: TestConnectionResult,
+        result: Union[TestConnectionResult, ReverseIngestionResponse],
         workflow_status: WorkflowStatus,
     ) -> None:
         """
@@ -496,14 +499,21 @@ class OMetaPatchMixin(OMetaPatchMixinBase):
         """
         result_data: Dict = {
             PatchField.PATH: PatchPath.RESPONSE,
-            PatchField.VALUE: test_connection_result.model_dump(),
+            PatchField.VALUE: result.model_dump(),
             PatchField.OPERATION: PatchOperation.ADD,
         }
 
         # for deserializing into json convert enum object to string
-        result_data[PatchField.VALUE]["status"] = result_data[PatchField.VALUE][
-            "status"
-        ].value
+        if isinstance(result, TestConnectionResult):
+            result_data[PatchField.VALUE]["status"] = result_data[PatchField.VALUE][
+                "status"
+            ].value
+        else:
+            # Convert UUID in string
+            data = result_data[PatchField.VALUE]
+            data["serviceId"] = str(data["serviceId"])
+            for operation_result in data["results"]:
+                operation_result["id"] = str(operation_result["id"])
 
         status_data: Dict = {
             PatchField.PATH: PatchPath.STATUS,

metadata/ingestion/ometa/mixins/pipeline_mixin.py

@@ -21,6 +21,7 @@ from metadata.generated.schema.entity.data.pipeline import (
     PipelineStatus,
     Task,
 )
+from metadata.generated.schema.type.usageRequest import UsageRequest
 from metadata.ingestion.ometa.client import REST
 from metadata.utils.logger import ometa_logger
 
@@ -83,6 +84,7 @@ class OMetaPipelineMixin:
             sourceUrl=pipeline.sourceUrl,
             concurrency=pipeline.concurrency,
             pipelineLocation=pipeline.pipelineLocation,
+            state=pipeline.state,
             startDate=pipeline.startDate,
             service=pipeline.service.fullyQualifiedName,
             tasks=all_tasks,
@@ -110,6 +112,7 @@ class OMetaPipelineMixin:
             sourceUrl=pipeline.sourceUrl,
             concurrency=pipeline.concurrency,
             pipelineLocation=pipeline.pipelineLocation,
+            state=pipeline.state,
             startDate=pipeline.startDate,
             service=pipeline.service.fullyQualifiedName,
             tasks=[task for task in pipeline.tasks if task.name in task_ids],
@@ -118,3 +121,18 @@ class OMetaPipelineMixin:
         )
 
         return self.create_or_update(updated_pipeline)
+
+    def publish_pipeline_usage(
+        self, pipeline: Pipeline, pipeline_usage_request: UsageRequest
+    ) -> None:
+        """
+        POST usage details for a Pipeline
+
+        :param pipeline: Pipeline Entity to update
+        :param pipeline_usage_request: Usage data to add
+        """
+        resp = self.client.put(
+            f"/usage/pipeline/{pipeline.id.root}",
+            data=pipeline_usage_request.model_dump_json(),
+        )
+        logger.debug("Published pipeline usage %s", resp)

metadata/ingestion/ometa/mixins/query_mixin.py

@@ -15,14 +15,20 @@ To be used by OpenMetadata class
 """
 import hashlib
 import json
+from functools import lru_cache
 from typing import List, Optional, Union
 
 from metadata.generated.schema.api.data.createQuery import CreateQueryRequest
+from metadata.generated.schema.api.data.createQueryCostRecord import (
+    CreateQueryCostRecordRequest,
+)
 from metadata.generated.schema.entity.data.dashboard import Dashboard
 from metadata.generated.schema.entity.data.query import Query
+from metadata.generated.schema.entity.data.queryCostRecord import QueryCostRecord
 from metadata.generated.schema.entity.data.table import Table
 from metadata.generated.schema.type.basic import Uuid
 from metadata.generated.schema.type.entityReference import EntityReference
+from metadata.generated.schema.type.tableUsageCount import QueryCostWrapper
 from metadata.ingestion.lineage.masker import mask_query
 from metadata.ingestion.ometa.client import REST
 from metadata.ingestion.ometa.utils import model_str
@@ -115,3 +121,40 @@ class OMetaQueryMixin:
         if res and res.get("data"):
             return [Query(**query) for query in res.get("data")]
         return None
+
+    @lru_cache(maxsize=5000)
+    def __get_query_by_hash(
+        self, query_hash: str, service_name: str
+    ) -> Optional[Query]:
+        return self.get_by_name(entity=Query, fqn=f"{service_name}.{query_hash}")
+
+    def publish_query_cost(self, query_cost_data: QueryCostWrapper, service_name: str):
+        """
+        Create Query Cost Record
+
+        Args:
+            query_cost_record: QueryCostWrapper
+        """
+
+        masked_query = mask_query(query_cost_data.query, query_cost_data.dialect)
+
+        query_hash = self._get_query_hash(masked_query)
+
+        query = self.__get_query_by_hash(
+            query_hash=query_hash, service_name=service_name
+        )
+        if not query:
+            return None
+
+        create_request = CreateQueryCostRecordRequest(
+            timestamp=int(query_cost_data.date),
+            jsonSchema="queryCostRecord",
+            queryReference=EntityReference(id=query.id.root, type="query"),
+            cost=query_cost_data.cost,
+            count=query_cost_data.count,
+            totalDuration=query_cost_data.totalDuration,
+        )
+
+        return self.client.post(
+            self.get_suffix(QueryCostRecord), data=create_request.model_dump_json()
+        )

metadata/ingestion/ometa/ometa_api.py

@@ -17,7 +17,7 @@ working with OpenMetadata entities.
 import traceback
 from typing import Dict, Generic, Iterable, List, Optional, Type, TypeVar, Union
 
-from pydantic import BaseModel
+from pydantic_settings import BaseSettings, SettingsConfigDict
 
 from metadata.generated.schema.api.createBot import CreateBot
 from metadata.generated.schema.api.services.ingestionPipelines.createIngestionPipeline import (
@@ -30,6 +30,7 @@ from metadata.generated.schema.type import basic
 from metadata.generated.schema.type.basic import FullyQualifiedEntityName
 from metadata.generated.schema.type.entityHistory import EntityVersionHistory
 from metadata.generated.schema.type.entityReference import EntityReference
+from metadata.ingestion.models.custom_pydantic import BaseModel
 from metadata.ingestion.ometa.auth_provider import OpenMetadataAuthenticationProvider
 from metadata.ingestion.ometa.client import REST, APIError, ClientConfig
 from metadata.ingestion.ometa.mixins.custom_property_mixin import (
@@ -37,6 +38,7 @@ from metadata.ingestion.ometa.mixins.custom_property_mixin import (
 )
 from metadata.ingestion.ometa.mixins.dashboard_mixin import OMetaDashboardMixin
 from metadata.ingestion.ometa.mixins.data_insight_mixin import DataInsightMixin
+from metadata.ingestion.ometa.mixins.domain_mixin import OMetaDomainMixin
 from metadata.ingestion.ometa.mixins.es_mixin import ESMixin
 from metadata.ingestion.ometa.mixins.ingestion_pipeline_mixin import (
     OMetaIngestionPipelineMixin,
@@ -89,6 +91,16 @@ class EmptyPayloadException(Exception):
     """
 
 
+class OpenMetadataSettings(BaseSettings):
+    """OpenMetadataConnection settings wrapper"""
+
+    model_config = SettingsConfigDict(
+        env_prefix="OPENMETADATA__", env_nested_delimiter="__", case_sensitive=True
+    )
+
+    connection: OpenMetadataConnection
+
+
 class OpenMetadata(
     OMetaPipelineMixin,
     OMetaMlModelMixin,
@@ -109,6 +121,7 @@ class OpenMetadata(
     OMetaSearchIndexMixin,
     OMetaCustomPropertyMixin,
     OMetaSuggestionsMixin,
+    OMetaDomainMixin,
     Generic[T, C],
 ):
     """
@@ -145,11 +158,14 @@ class OpenMetadata(
 
         get_verify_ssl = get_verify_ssl_fn(self.config.verifySSL)
 
+        extra_headers: Optional[dict[str, str]] = None
+        if self.config.extraHeaders:
+            extra_headers = self.config.extraHeaders.root
         client_config: ClientConfig = ClientConfig(
             base_url=self.config.hostPort,
             api_version=self.config.apiVersion,
             auth_header="Authorization",
-            extra_headers=self.config.extraHeaders,
+            extra_headers=extra_headers,
             auth_token=self._auth_provider.get_access_token,
             verify=get_verify_ssl(self.config.sslConfig),
         )
@@ -158,6 +174,11 @@ class OpenMetadata(
         if self.config.enableVersionValidation:
             self.validate_versions()
 
+    @classmethod
+    def from_env(cls) -> "OpenMetadata":
+        settings = OpenMetadataSettings()
+        return cls(settings.connection)
+
     @staticmethod
     def get_suffix(entity: Type[T]) -> str:
         """
@@ -268,7 +289,11 @@ class OpenMetadata(
             )
 
         fn = getattr(self.client, method)
-        resp = fn(self.get_suffix(entity), data=data.model_dump_json())
+        resp = fn(
+            # this might be a regular pydantic model so we build the context manually
+            self.get_suffix(entity),
+            data=data.model_dump_json(context={"mask_secrets": False}),
+        )
         if not resp:
             raise EmptyPayloadException(
                 f"Got an empty response when trying to PUT to {self.get_suffix(entity)}, {data.model_dump_json()}"

metadata/ingestion/ometa/routes.py

@@ -46,6 +46,9 @@ from metadata.generated.schema.api.data.createMetric import CreateMetricRequest
 from metadata.generated.schema.api.data.createMlModel import CreateMlModelRequest
 from metadata.generated.schema.api.data.createPipeline import CreatePipelineRequest
 from metadata.generated.schema.api.data.createQuery import CreateQueryRequest
+from metadata.generated.schema.api.data.createQueryCostRecord import (
+    CreateQueryCostRecordRequest,
+)
 from metadata.generated.schema.api.data.createSearchIndex import (
     CreateSearchIndexRequest,
 )
@@ -133,6 +136,7 @@ from metadata.generated.schema.entity.data.metric import Metric
 from metadata.generated.schema.entity.data.mlmodel import MlModel
 from metadata.generated.schema.entity.data.pipeline import Pipeline
 from metadata.generated.schema.entity.data.query import Query
+from metadata.generated.schema.entity.data.queryCostRecord import QueryCostRecord
 from metadata.generated.schema.entity.data.report import Report
 from metadata.generated.schema.entity.data.searchIndex import SearchIndex
 from metadata.generated.schema.entity.data.storedProcedure import StoredProcedure
@@ -280,4 +284,7 @@ ROUTES = {
     CreateAppMarketPlaceDefinitionRequest.__name__: "/apps/marketplace",
     # Settings
     Settings.__name__: "/system/settings",
+    # Query Cost
+    QueryCostRecord.__name__: "/queryCostRecord",
+    CreateQueryCostRecordRequest.__name__: "/queryCostRecord",
 }

metadata/ingestion/ometa/utils.py

@@ -20,6 +20,7 @@ from pydantic import BaseModel
 from requests.utils import quote as url_quote
 
 from metadata.generated.schema.type.basic import FullyQualifiedEntityName
+from metadata.generated.schema.type.entityReference import EntityReference
 
 T = TypeVar("T", bound=BaseModel)
 
@@ -85,3 +86,15 @@ def quote(fqn: Union[FullyQualifiedEntityName, str]) -> str:
     E.g., `"foo.bar/baz"` -> `%22foo.bar%2Fbaz%22`
     """
     return url_quote(model_str(fqn), safe="")
+
+
+def build_entity_reference(entity: T) -> EntityReference:
+    """Get the EntityReference from the Entity itself"""
+    return EntityReference(
+        id=entity.id,
+        type=get_entity_type(type(entity)),
+        name=model_str(entity.name),
+        fullyQualifiedName=model_str(entity.fullyQualifiedName),
+        description=entity.description,
+        href=entity.href,
+    )

metadata/ingestion/processor/query_parser.py

@@ -66,6 +66,7 @@ def parse_sql_statement(record: TableQuery, dialect: Dialect) -> Optional[Parsed
         date=str(start_time),
         serviceName=record.serviceName,
         duration=record.duration,
+        cost=record.cost,
     )