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

metadata/ingestion/source/database/snowflake/utils.py

@@ -10,12 +10,14 @@
 #  limitations under the License.
 
 """
-Module to define overriden dialect methods
+Module to define overridden dialect methods
 """
-
+import operator
+from functools import reduce
 from typing import Dict, Optional
 
 import sqlalchemy.types as sqltypes
+from snowflake.sqlalchemy.snowdialect import SnowflakeDialect
 from sqlalchemy import exc as sa_exc
 from sqlalchemy import util as sa_util
 from sqlalchemy.engine import reflection
@@ -35,13 +37,17 @@ from metadata.ingestion.source.database.snowflake.queries import (
     SNOWFLAKE_GET_EXTERNAL_TABLE_NAMES,
     SNOWFLAKE_GET_MVIEW_NAMES,
     SNOWFLAKE_GET_SCHEMA_COLUMNS,
+    SNOWFLAKE_GET_STREAM_DEFINITION,
+    SNOWFLAKE_GET_STREAM_NAMES,
     SNOWFLAKE_GET_TABLE_DDL,
     SNOWFLAKE_GET_TRANSIENT_NAMES,
+    SNOWFLAKE_GET_VIEW_DEFINITION,
     SNOWFLAKE_GET_VIEW_NAMES,
     SNOWFLAKE_GET_WITHOUT_TRANSIENT_TABLE_NAMES,
     SNOWFLAKE_INCREMENTAL_GET_DYNAMIC_TABLE_NAMES,
     SNOWFLAKE_INCREMENTAL_GET_EXTERNAL_TABLE_NAMES,
     SNOWFLAKE_INCREMENTAL_GET_MVIEW_NAMES,
+    SNOWFLAKE_INCREMENTAL_GET_STREAM_NAMES,
     SNOWFLAKE_INCREMENTAL_GET_TRANSIENT_NAMES,
     SNOWFLAKE_INCREMENTAL_GET_VIEW_NAMES,
     SNOWFLAKE_INCREMENTAL_GET_WITHOUT_TRANSIENT_TABLE_NAMES,
@@ -52,6 +58,7 @@ from metadata.utils.sqlalchemy_utils import (
     get_table_comment_wrapper,
 )
 
+dialect = SnowflakeDialect()
 Query = str
 QueryMap = Dict[str, Query]
 
@@ -82,6 +89,29 @@ VIEW_QUERY_MAPS = {
     },
 }
 
+STREAM_QUERY_MAPS = {
+    "full": {
+        "default": SNOWFLAKE_GET_STREAM_NAMES,
+    },
+    "incremental": {
+        "default": SNOWFLAKE_INCREMENTAL_GET_STREAM_NAMES,
+    },
+}
+
+
+def _denormalize_quote_join(*idents):
+    ip = dialect.identifier_preparer
+    split_idents = reduce(
+        operator.add,
+        [ip._split_schema_by_dot(ids) for ids in idents if ids is not None],
+    )
+    quoted_identifiers = ip._quote_free_identifiers(*split_idents)
+    normalized_identifiers = (
+        item if item.startswith('"') and item.endswith('"') else f'"{item}"'
+        for item in quoted_identifiers
+    )
+    return ".".join(normalized_identifiers)
+
 
 def _quoted_name(entity_name: Optional[str]) -> Optional[str]:
     if entity_name:
@@ -133,6 +163,20 @@ def get_view_names_reflection(self, schema=None, **kw):
         )
 
 
+def get_stream_names_reflection(self, schema=None, **kw):
+    """Return all stream names in `schema`.
+
+    :param schema: Optional, retrieve names from a non-default schema.
+        For special quoting, use :class:`.quoted_name`.
+
+    """
+
+    with self._operation_context() as conn:  # pylint: disable=protected-access
+        return self.dialect.get_stream_names(
+            conn, schema, info_cache=self.info_cache, **kw
+        )
+
+
 def _get_query_map(
     incremental: Optional[IncrementalConfig], query_maps: Dict[str, QueryMap]
 ):
@@ -145,7 +189,7 @@ def _get_query_map(
 def _get_query_parameters(
     self, connection, schema: str, incremental: Optional[IncrementalConfig]
 ):
-    """Returns the proper query parameters depending if the extraciton is Incremental or Full"""
+    """Returns the proper query parameters depending if the extraction is Incremental or Full"""
     parameters = {"schema": fqn.unquote_name(schema)}
 
     if incremental and incremental.enabled:
@@ -208,6 +252,24 @@ def get_view_names(self, connection, schema, **kw):
     return result
 
 
+def get_stream_names(self, connection, schema, **kw):
+    incremental = kw.get("incremental")
+
+    queries = _get_query_map(incremental, STREAM_QUERY_MAPS)
+    parameters = _get_query_parameters(self, connection, schema, incremental)
+
+    query = queries["default"]
+
+    cursor = connection.execute(query.format(**parameters))
+    result = SnowflakeTableList(
+        tables=[
+            SnowflakeTable(name=self.normalize_name(row[1]), deleted=None)
+            for row in cursor
+        ]
+    )
+    return result
+
+
 @reflection.cache
 def get_view_definition(  # pylint: disable=unused-argument
     self, connection, view_name, schema=None, **kw
@@ -216,12 +278,10 @@ def get_view_definition(  # pylint: disable=unused-argument
     Gets the view definition
     """
     schema = schema or self.default_schema_name
-    if schema:
-        cursor = connection.execute(
-            f"SELECT GET_DDL('VIEW','{schema}.{view_name}') AS \"text\""
-        )
-    else:
-        cursor = connection.execute(f"SELECT GET_DDL('VIEW','{view_name}') AS \"text\"")
+    view_name = f"{schema}.{view_name}" if schema else view_name
+    cursor = connection.execute(
+        SNOWFLAKE_GET_VIEW_DEFINITION.format(view_name=view_name)
+    )
     n2i = self.__class__._map_name_to_idx(cursor)  # pylint: disable=protected-access
     try:
         ret = cursor.fetchone()
@@ -232,6 +292,27 @@ def get_view_definition(  # pylint: disable=unused-argument
     return None
 
 
+@reflection.cache
+def get_stream_definition(  # pylint: disable=unused-argument
+    self, connection, stream_name, schema=None, **kw
+):
+    """
+    Gets the stream definition
+    """
+    schema = schema or self.default_schema_name
+    stream_name = f"{schema}.{stream_name}" if schema else stream_name
+    cursor = connection.execute(
+        SNOWFLAKE_GET_STREAM_DEFINITION.format(stream_name=stream_name)
+    )
+    try:
+        result = cursor.fetchone()
+        if result:
+            return result[0]
+    except Exception:
+        pass
+    return None
+
+
 @reflection.cache
 def get_table_comment(
     self, connection, table_name, schema=None, **kw
@@ -256,17 +337,16 @@ def get_schema_columns(self, connection, schema, **kw):
     None, as it is cacheable and is an unexpected return type for this function"""
     ans = {}
     current_database, _ = self._current_database_schema(connection, **kw)
-    full_schema_name = self._denormalize_quote_join(
-        current_database, fqn.quote_name(schema)
-    )
+    full_schema_name = _denormalize_quote_join(current_database, fqn.quote_name(schema))
     try:
         schema_primary_keys = self._get_schema_primary_keys(
             connection, full_schema_name, **kw
         )
+        # removing " " from schema name because schema name is in the WHERE clause of a query
+        table_schema = self.denormalize_name(fqn.unquote_name(schema))
+        table_schema = table_schema.lower() if schema.islower() else table_schema
         result = connection.execute(
-            text(SNOWFLAKE_GET_SCHEMA_COLUMNS),
-            {"table_schema": self.denormalize_name(fqn.unquote_name(schema))}
-            # removing " " from schema name because schema name is in the WHERE clause of a query
+            text(SNOWFLAKE_GET_SCHEMA_COLUMNS), {"table_schema": table_schema}
         )
 
     except sa_exc.ProgrammingError as p_err:
@@ -300,7 +380,8 @@ def get_schema_columns(self, connection, schema, **kw):
             sa_util.warn(
                 f"Did not recognize type '{coltype}' of column '{column_name}'"
             )
-            col_type = sqltypes.NULLTYPE
+            col_type = sqltypes.NullType
+            type_instance = col_type()
         else:
             if issubclass(col_type, FLOAT):
                 col_type_kw["precision"] = numeric_precision
@@ -310,8 +391,7 @@ def get_schema_columns(self, connection, schema, **kw):
                 col_type_kw["scale"] = numeric_scale
             elif issubclass(col_type, (sqltypes.String, sqltypes.BINARY)):
                 col_type_kw["length"] = character_maximum_length
-
-        type_instance = col_type(**col_type_kw)
+            type_instance = col_type(**col_type_kw)
 
         current_table_pks = schema_primary_keys.get(table_name)
 
@@ -330,11 +410,13 @@ def get_schema_columns(self, connection, schema, **kw):
                 ),
                 "comment": comment,
                 "primary_key": (
-                    column_name
-                    in schema_primary_keys[table_name]["constrained_columns"]
-                )
-                if current_table_pks
-                else False,
+                    (
+                        column_name
+                        in schema_primary_keys[table_name]["constrained_columns"]
+                    )
+                    if current_table_pks
+                    else False
+                ),
             }
         )
         if is_identity == "YES":
@@ -362,9 +444,10 @@ def get_pk_constraint(self, connection, table_name, schema=None, **kw):
     schema = schema or self.default_schema_name
     schema = _quoted_name(entity_name=schema)
     current_database, current_schema = self._current_database_schema(connection, **kw)
-    full_schema_name = self._denormalize_quote_join(
+    full_schema_name = _denormalize_quote_join(
         current_database, schema if schema else current_schema
     )
+
     return self._get_schema_primary_keys(
         connection, self.denormalize_name(full_schema_name), **kw
     ).get(table_name, {"constrained_columns": [], "name": None})
@@ -378,7 +461,7 @@ def get_foreign_keys(self, connection, table_name, schema=None, **kw):
     schema = schema or self.default_schema_name
     schema = _quoted_name(entity_name=schema)
     current_database, current_schema = self._current_database_schema(connection, **kw)
-    full_schema_name = self._denormalize_quote_join(
+    full_schema_name = _denormalize_quote_join(
         current_database, schema if schema else current_schema
     )
 
@@ -452,9 +535,10 @@ def get_unique_constraints(self, connection, table_name, schema, **kw):
     schema = schema or self.default_schema_name
     schema = _quoted_name(entity_name=schema)
     current_database, current_schema = self._current_database_schema(connection, **kw)
-    full_schema_name = self._denormalize_quote_join(
+    full_schema_name = _denormalize_quote_join(
         current_database, schema if schema else current_schema
     )
+
     return self._get_schema_unique_constraints(
         connection, self.denormalize_name(full_schema_name), **kw
     ).get(table_name, [])

metadata/ingestion/source/database/sql_column_handler.py

@@ -24,6 +24,7 @@ from metadata.generated.schema.entity.data.table import (
     ConstraintType,
     DataType,
     TableConstraint,
+    TableType,
 )
 from metadata.ingestion.source.database.column_type_parser import ColumnTypeParser
 from metadata.utils.execution_time_tracker import calculate_execution_time
@@ -202,9 +203,30 @@ class SqlColumnHandlerMixin:
             ]
         return Column(**parsed_string)
 
+    def _get_columns_internal(
+        self,
+        schema_name: str,
+        table_name: str,
+        db_name: str,
+        inspector: Inspector,
+        table_type: TableType = None,
+    ):
+        """
+        Get columns list
+        """
+
+        return inspector.get_columns(
+            table_name, schema_name, table_type=table_type, db_name=db_name
+        )
+
     @calculate_execution_time()
     def get_columns_and_constraints(  # pylint: disable=too-many-locals
-        self, schema_name: str, table_name: str, db_name: str, inspector: Inspector
+        self,
+        schema_name: str,
+        table_name: str,
+        db_name: str,
+        inspector: Inspector,
+        table_type: TableType = None,
     ) -> Tuple[
         Optional[List[Column]], Optional[List[TableConstraint]], Optional[List[Dict]]
     ]:
@@ -246,74 +268,87 @@ class SqlColumnHandlerMixin:
 
         table_columns = []
 
-        columns = inspector.get_columns(table_name, schema_name, db_name=db_name)
-        for column in columns:
-            try:
-                children = None
-                (
+        columns = self._get_columns_internal(
+            schema_name, table_name, db_name, inspector, table_type
+        )
+
+        def process_column(column: dict):
+            (
+                data_type_display,
+                arr_data_type,
+                parsed_string,
+            ) = self._process_col_type(column, schema_name)
+            self.process_additional_table_constraints(
+                column=column, table_constraints=table_constraints
+            )
+            if parsed_string is None:
+                col_type = ColumnTypeParser.get_column_type(column["type"])
+                col_constraint = self._get_column_constraints(
+                    column, pk_columns, column_level_unique_constraints
+                )
+                col_data_length = self._check_col_length(col_type, column["type"])
+                precision = ColumnTypeParser.check_col_precision(
+                    col_type, column["type"]
+                )
+                if col_type is None:
+                    col_type = DataType.UNKNOWN.name
+                    data_type_display = col_type.lower()
+                    logger.warning(
+                        f"Unknown type {repr(column['type'])}: {column['name']}"
+                    )
+                data_type_display = self._get_display_datatype(
                     data_type_display,
+                    col_type,
+                    col_data_length,
                     arr_data_type,
-                    parsed_string,
-                ) = self._process_col_type(column, schema_name)
-                self.process_additional_table_constraints(
-                    column=column, table_constraints=table_constraints
+                    precision,
                 )
-                if parsed_string is None:
-                    col_type = ColumnTypeParser.get_column_type(column["type"])
-                    col_constraint = self._get_column_constraints(
-                        column, pk_columns, column_level_unique_constraints
-                    )
-                    col_data_length = self._check_col_length(col_type, column["type"])
-                    precision = ColumnTypeParser.check_col_precision(
-                        col_type, column["type"]
-                    )
-                    if col_type is None:
-                        col_type = DataType.UNKNOWN.name
-                        data_type_display = col_type.lower()
-                        logger.warning(
-                            f"Unknown type {repr(column['type'])}: {column['name']}"
-                        )
-                    data_type_display = self._get_display_datatype(
-                        data_type_display,
-                        col_type,
-                        col_data_length,
-                        arr_data_type,
-                        precision,
-                    )
-                    col_data_length = 1 if col_data_length is None else col_data_length
+                col_data_length = 1 if col_data_length is None else col_data_length
 
-                    om_column = Column(
-                        name=ColumnName(
-                            root=column["name"]
+                om_column = Column(
+                    name=ColumnName(
+                        root=(
+                            column["name"]
                             # Passing whitespace if column name is an empty string
                             # since pydantic doesn't accept empty string
                             if column["name"]
                             else " "
-                        ),
-                        description=column.get("comment"),
-                        dataType=col_type,
-                        dataTypeDisplay=column.get(
-                            "system_data_type", data_type_display
-                        ),
-                        dataLength=col_data_length,
-                        constraint=col_constraint,
-                        children=children,
-                        arrayDataType=arr_data_type,
-                        ordinalPosition=column.get("ordinalPosition"),
-                    )
-                    if precision:
-                        # Precision and scale must be integer values
-                        om_column.precision = int(precision[0])
-                        om_column.scale = int(precision[1])
-
-                else:
-                    col_obj = self._process_complex_col_type(
-                        column=column, parsed_string=parsed_string
-                    )
-                    om_column = col_obj
-                om_column.tags = self.get_column_tag_labels(
-                    table_name=table_name, column=column
+                        )
+                    ),
+                    description=column.get("comment"),
+                    dataType=col_type,
+                    dataTypeDisplay=column.get("system_data_type", data_type_display),
+                    dataLength=col_data_length,
+                    constraint=col_constraint,
+                    arrayDataType=arr_data_type,
+                    ordinalPosition=column.get("ordinalPosition"),
                 )
+                if precision:
+                    # Precision and scale must be integer values
+                    om_column.precision = int(precision[0])
+                    om_column.scale = int(precision[1])
+            else:
+                col_obj = self._process_complex_col_type(
+                    column=column, parsed_string=parsed_string
+                )
+                om_column = col_obj
+
+                if column.get("children"):
+                    # Prioritize source-provided children for column processing.
+                    # If 'children' are directly provided in the source metadata,
+                    # process and assign them to the output column, overriding any derived children.
+                    # Currently, this is only used for BigQuery.
+                    om_column.children = [
+                        process_column(children) for children in column.get("children")
+                    ]
+            om_column.tags = self.get_column_tag_labels(
+                table_name=table_name, column=column
+            )
+            return om_column
+
+        for column in columns:
+            try:
+                om_column = process_column(column)
             except Exception as exc:
                 logger.debug(traceback.format_exc())
                 logger.warning(

metadata/ingestion/source/database/sqlite/service_spec.py

@@ -1,4 +1,8 @@
+from metadata.ingestion.source.database.sqlite.lineage import SqliteLineageSource
 from metadata.ingestion.source.database.sqlite.metadata import SqliteSource
 from metadata.utils.service_spec.default import DefaultDatabaseSpec
 
-ServiceSpec = DefaultDatabaseSpec(metadata_source_class=SqliteSource)
+ServiceSpec = DefaultDatabaseSpec(
+    metadata_source_class=SqliteSource,
+    lineage_source_class=SqliteLineageSource,
+)

metadata/ingestion/source/database/teradata/service_spec.py

@@ -1,4 +1,8 @@
+from metadata.ingestion.source.database.teradata.lineage import TeradataLineageSource
 from metadata.ingestion.source.database.teradata.metadata import TeradataSource
 from metadata.utils.service_spec.default import DefaultDatabaseSpec
 
-ServiceSpec = DefaultDatabaseSpec(metadata_source_class=TeradataSource)
+ServiceSpec = DefaultDatabaseSpec(
+    metadata_source_class=TeradataSource,
+    lineage_source_class=TeradataLineageSource,
+)

metadata/ingestion/source/database/trino/lineage.py

@@ -11,7 +11,16 @@
 """
 Trino lineage module
 """
+import traceback
+from typing import Iterable, List
 
+from metadata.generated.schema.api.lineage.addLineage import AddLineageRequest
+from metadata.generated.schema.entity.data.database import Database
+from metadata.generated.schema.entity.data.table import Table
+from metadata.generated.schema.entity.services.ingestionPipelines.status import (
+    StackTraceError,
+)
+from metadata.ingestion.api.models import Either
 from metadata.ingestion.source.database.lineage_source import LineageSource
 from metadata.ingestion.source.database.trino.queries import TRINO_SQL_STATEMENT
 from metadata.ingestion.source.database.trino.query_parser import TrinoQueryParserSource
@@ -32,3 +41,89 @@ class TrinoLineageSource(TrinoQueryParserSource, LineageSource):
             OR lower("query") LIKE '%%merge%%'
         )
     """
+
+    def get_cross_database_fqn_from_service_names(self) -> List[str]:
+        database_service_names = self.source_config.crossDatabaseServiceNames
+        return [
+            database.fullyQualifiedName.root
+            for service in database_service_names
+            for database in self.metadata.list_all_entities(
+                entity=Database, params={"service": service}
+            )
+        ]
+
+    def check_same_table(self, table1: Table, table2: Table) -> bool:
+        """
+        Method to check whether the table1 and table2 are same
+        """
+        return table1.name.root == table2.name.root and {
+            column.name.root for column in table1.columns
+        } == {column.name.root for column in table2.columns}
+
+    def get_cross_database_lineage(
+        self, from_table: Table, to_table: Table
+    ) -> Either[AddLineageRequest]:
+        """
+        Method to return cross database lineage request object
+        """
+        column_lineage = None
+        if from_table and from_table.columns and to_table and to_table.columns:
+            column_lineage = self.get_column_lineage(
+                from_table=from_table, to_table=to_table
+            )
+        return self.get_add_cross_database_lineage_request(
+            from_entity=from_table, to_entity=to_table, column_lineage=column_lineage
+        )
+
+    def yield_cross_database_lineage(self) -> Iterable[Either[AddLineageRequest]]:
+        try:
+            all_cross_database_fqns = self.get_cross_database_fqn_from_service_names()
+            cross_database_table_fqn_mapping = {}
+
+            # Get all databases for the specified Trino service
+            trino_databases = self.metadata.list_all_entities(
+                entity=Database, params={"service": self.config.serviceName}
+            )
+            for trino_database in trino_databases:
+                trino_database_fqn = trino_database.fullyQualifiedName.root
+
+                # Get all tables for the specified Trino database schema
+                trino_tables = self.metadata.list_all_entities(
+                    entity=Table, params={"database": trino_database_fqn}
+                )
+                # NOTE: Currently, tables in system-defined schemas will also be checked for lineage.
+                for trino_table in trino_tables:
+                    trino_table_fqn = trino_table.fullyQualifiedName.root
+                    for cross_database_fqn in all_cross_database_fqns:
+                        # Construct the FQN for cross-database tables
+                        cross_database_table_fqn = trino_table_fqn.replace(
+                            trino_database_fqn, cross_database_fqn
+                        )
+                        # Cache cross-database table against its FQN to avoid repeated API calls
+                        cross_database_table = cross_database_table_fqn_mapping[
+                            cross_database_table_fqn
+                        ] = cross_database_table_fqn_mapping.get(
+                            cross_database_table_fqn,
+                            self.metadata.get_by_name(
+                                Table, fqn=cross_database_table_fqn
+                            ),
+                        )
+                        # Create cross database lineage request if both tables are same
+                        if cross_database_table and self.check_same_table(
+                            trino_table, cross_database_table
+                        ):
+                            yield self.get_cross_database_lineage(
+                                cross_database_table, trino_table
+                            )
+                            break
+        except Exception as exc:
+            yield Either(
+                left=StackTraceError(
+                    name=f"{self.config.serviceName} Cross Database Lineage",
+                    error=(
+                        "Error to yield cross database lineage details "
+                        f"service name [{self.config.serviceName}]: {exc}"
+                    ),
+                    stackTrace=traceback.format_exc(),
+                )
+            )

metadata/ingestion/source/database/unitycatalog/connection.py

@@ -16,6 +16,7 @@ from functools import partial
 from typing import Optional
 
 from databricks.sdk import WorkspaceClient
+from sqlalchemy.engine import Engine
 
 from metadata.generated.schema.entity.automations.workflow import (
     Workflow as AutomationWorkflow,
@@ -26,6 +27,11 @@ from metadata.generated.schema.entity.services.connections.database.unityCatalog
 from metadata.generated.schema.entity.services.connections.testConnectionResult import (
     TestConnectionResult,
 )
+from metadata.ingestion.connections.builders import (
+    create_generic_db_connection,
+    get_connection_args_common,
+    init_empty_connection_arguments,
+)
 from metadata.ingestion.connections.test_connections import test_connection_steps
 from metadata.ingestion.ometa.ometa_api import OpenMetadata
 from metadata.ingestion.source.database.unitycatalog.client import UnityCatalogClient
@@ -53,6 +59,23 @@ def get_connection(connection: UnityCatalogConnection) -> WorkspaceClient:
     )
 
 
+def get_sqlalchemy_connection(connection: UnityCatalogConnection) -> Engine:
+    """
+    Create sqlalchemy connection
+    """
+
+    if connection.httpPath:
+        if not connection.connectionArguments:
+            connection.connectionArguments = init_empty_connection_arguments()
+        connection.connectionArguments.root["http_path"] = connection.httpPath
+
+    return create_generic_db_connection(
+        connection=connection,
+        get_connection_url_fn=get_connection_url,
+        get_connection_args_fn=get_connection_args_common,
+    )
+
+
 def test_connection(
     metadata: OpenMetadata,
     connection: WorkspaceClient,

metadata/ingestion/source/database/unitycatalog/lineage.py

@@ -42,6 +42,7 @@ from metadata.ingestion.source.database.unitycatalog.client import UnityCatalogC
 from metadata.ingestion.source.database.unitycatalog.connection import get_connection
 from metadata.ingestion.source.database.unitycatalog.models import LineageTableStreams
 from metadata.utils import fqn
+from metadata.utils.helpers import retry_with_docker_host
 from metadata.utils.logger import ingestion_logger
 
 logger = ingestion_logger()
@@ -52,6 +53,7 @@ class UnitycatalogLineageSource(Source):
     Lineage Unity Catalog Source
     """
 
+    @retry_with_docker_host()
     def __init__(
         self,
         config: WorkflowSource,

metadata/ingestion/source/database/unitycatalog/metadata.py

@@ -75,6 +75,7 @@ from metadata.ingestion.source.database.unitycatalog.models import (
 from metadata.ingestion.source.models import TableView
 from metadata.utils import fqn
 from metadata.utils.filters import filter_by_database, filter_by_schema, filter_by_table
+from metadata.utils.helpers import retry_with_docker_host
 from metadata.utils.logger import ingestion_logger
 
 logger = ingestion_logger()
@@ -89,6 +90,7 @@ class UnitycatalogSource(
     the unity catalog source
     """
 
+    @retry_with_docker_host()
     def __init__(self, config: WorkflowSource, metadata: OpenMetadata):
         super().__init__()
         self.config = config

metadata/ingestion/source/database/usage_source.py

@@ -72,6 +72,7 @@ class UsageSource(QueryParserSource, ABC):
                                 startTime=query_dict.get("start_time", ""),
                                 endTime=query_dict.get("end_time", ""),
                                 duration=query_dict.get("duration"),
+                                cost=query_dict.get("cost"),
                                 analysisDate=DateTime(analysis_date),
                                 aborted=self.get_aborted_status(query_dict),
                                 databaseName=self.get_database_name(query_dict),
@@ -142,6 +143,7 @@ class UsageSource(QueryParserSource, ABC):
                                         duration=row.get("duration"),
                                         serviceName=self.config.serviceName,
                                         databaseSchema=self.get_schema_name(row),
+                                        cost=row.get("cost"),
                                     )
                                 )
                             except Exception as exc: