openmetadata-ingestion 1.5.6.0__py3-none-any.whl → 1.5.7.0__py3-none-any.whl

metadata/generated/schema/type/tableUsageCount.py

@@ -1,6 +1,6 @@
 # generated by datamodel-codegen:
 #   filename:  type/tableUsageCount.json
-#   timestamp: 2024-10-03T13:02:21+00:00
+#   timestamp: 2024-10-17T16:44:45+00:00
 
 from __future__ import annotations
 

metadata/generated/schema/type/tagLabel.py

@@ -1,6 +1,6 @@
 # generated by datamodel-codegen:
 #   filename:  type/tagLabel.json
-#   timestamp: 2024-10-03T13:02:21+00:00
+#   timestamp: 2024-10-17T16:44:45+00:00
 
 from __future__ import annotations
 

metadata/generated/schema/type/usageDetails.py

@@ -1,6 +1,6 @@
 # generated by datamodel-codegen:
 #   filename:  type/usageDetails.json
-#   timestamp: 2024-10-03T13:02:21+00:00
+#   timestamp: 2024-10-17T16:44:45+00:00
 
 from __future__ import annotations
 

metadata/generated/schema/type/usageRequest.py

@@ -1,6 +1,6 @@
 # generated by datamodel-codegen:
 #   filename:  type/usageRequest.json
-#   timestamp: 2024-10-03T13:02:21+00:00
+#   timestamp: 2024-10-17T16:44:45+00:00
 
 from __future__ import annotations
 

metadata/generated/schema/type/votes.py

@@ -1,6 +1,6 @@
 # generated by datamodel-codegen:
 #   filename:  type/votes.json
-#   timestamp: 2024-10-03T13:02:21+00:00
+#   timestamp: 2024-10-17T16:44:45+00:00
 
 from __future__ import annotations
 

metadata/ingestion/source/dashboard/quicksight/metadata.py

@@ -247,6 +247,7 @@ class QuicksightSource(DashboardServiceSource):
                             **data_source["RelationalTable"]
                         )
                     except (KeyError, ValidationError) as err:
+                        data_source_resp = None
                         yield Either(
                             left=StackTraceError(
                                 name="Lineage",
@@ -257,62 +258,70 @@ class QuicksightSource(DashboardServiceSource):
                                 stackTrace=traceback.format_exc(),
                             )
                         )
+                    if data_source_resp:
+                        schema_name = data_source_resp.schema_name
+                        table_name = data_source_resp.table_name
 
-                    schema_name = data_source_resp.schema_name
-                    table_name = data_source_resp.table_name
-
-                    list_data_source_func = lambda kwargs: self.client.list_data_sources(  # pylint: disable=unnecessary-lambda-assignment
-                        **kwargs
-                    )
+                        list_data_source_func = lambda kwargs: self.client.list_data_sources(  # pylint: disable=unnecessary-lambda-assignment
+                            **kwargs
+                        )
 
-                    data_source_summary_list = self._check_pagination(
-                        listing_method=list_data_source_func,
-                        entity_key="DataSources",
-                    )
+                        data_source_summary_list = self._check_pagination(
+                            listing_method=list_data_source_func,
+                            entity_key="DataSources",
+                        )
 
-                    data_source_ids = [
-                        data_source_arn["DataSourceId"]
-                        for data_source_arn in data_source_summary_list or []
-                        if data_source_arn["Arn"] in data_source_resp.datasource_arn
-                    ]
+                        data_source_ids = [
+                            data_source_arn["DataSourceId"]
+                            for data_source_arn in data_source_summary_list or []
+                            if data_source_arn["Arn"] in data_source_resp.datasource_arn
+                        ]
 
-                    for data_source_id in data_source_ids or []:
-                        data_source_resp = DescribeDataSourceResponse(
-                            **self.client.describe_data_source(
-                                AwsAccountId=self.aws_account_id,
-                                DataSourceId=data_source_id,
-                            )
-                        ).DataSource
-                        if data_source_resp and data_source_resp.DataSourceParameters:
-                            data_source_dict = data_source_resp.DataSourceParameters
-                            for db in data_source_dict.keys() or []:
-                                from_fqn = fqn.build(
-                                    self.metadata,
-                                    entity_type=Table,
-                                    service_name=db_service_name,
-                                    database_name=data_source_dict[db].get("Database"),
-                                    schema_name=schema_name,
-                                    table_name=table_name,
-                                    skip_es_search=True,
+                        for data_source_id in data_source_ids or []:
+                            data_source_resp = DescribeDataSourceResponse(
+                                **self.client.describe_data_source(
+                                    AwsAccountId=self.aws_account_id,
+                                    DataSourceId=data_source_id,
                                 )
-                                from_entity = self.metadata.get_by_name(
-                                    entity=Table,
-                                    fqn=from_fqn,
-                                )
-                                to_fqn = fqn.build(
-                                    self.metadata,
-                                    entity_type=Dashboard,
-                                    service_name=self.config.serviceName,
-                                    dashboard_name=dashboard_details.DashboardId,
-                                )
-                                to_entity = self.metadata.get_by_name(
-                                    entity=Dashboard,
-                                    fqn=to_fqn,
-                                )
-                                if from_entity is not None and to_entity is not None:
-                                    yield self._get_add_lineage_request(
-                                        to_entity=to_entity, from_entity=from_entity
+                            ).DataSource
+                            if (
+                                data_source_resp
+                                and data_source_resp.DataSourceParameters
+                            ):
+                                data_source_dict = data_source_resp.DataSourceParameters
+                                for db in data_source_dict.keys() or []:
+                                    from_fqn = fqn.build(
+                                        self.metadata,
+                                        entity_type=Table,
+                                        service_name=db_service_name,
+                                        database_name=data_source_dict[db].get(
+                                            "Database"
+                                        ),
+                                        schema_name=schema_name,
+                                        table_name=table_name,
+                                        skip_es_search=True,
+                                    )
+                                    from_entity = self.metadata.get_by_name(
+                                        entity=Table,
+                                        fqn=from_fqn,
+                                    )
+                                    to_fqn = fqn.build(
+                                        self.metadata,
+                                        entity_type=Dashboard,
+                                        service_name=self.config.serviceName,
+                                        dashboard_name=dashboard_details.DashboardId,
+                                    )
+                                    to_entity = self.metadata.get_by_name(
+                                        entity=Dashboard,
+                                        fqn=to_fqn,
                                     )
+                                    if (
+                                        from_entity is not None
+                                        and to_entity is not None
+                                    ):
+                                        yield self._get_add_lineage_request(
+                                            to_entity=to_entity, from_entity=from_entity
+                                        )
         except Exception as exc:  # pylint: disable=broad-except
             yield Either(
                 left=StackTraceError(

metadata/ingestion/source/dashboard/quicksight/models.py

@@ -50,12 +50,12 @@ class DashboardResp(BaseModel):
     RequestId: Optional[str] = None
 
 
-class DataSource(BaseModel):
+class DataSourceModel(BaseModel):
     DataSourceId: str
     DataSourceParameters: Optional[dict] = None
 
 
 class DescribeDataSourceResponse(BaseModel):
-    DataSource: Optional[DataSource] = None
+    DataSource: Optional[DataSourceModel] = None
     RequestId: Optional[str] = None
     Status: Optional[int] = None

metadata/ingestion/source/dashboard/superset/connection.py

@@ -84,7 +84,12 @@ def test_connection(
     else:
         test_fn["CheckAccess"] = partial(test_connection_engine_step, client)
         test_fn["GetDashboards"] = partial(test_query, client, FETCH_DASHBOARDS_TEST)
-        test_fn["GetCharts"] = partial(test_query, client, FETCH_ALL_CHARTS_TEST)
+        if isinstance(service_connection.connection, MysqlConnection):
+            test_fn["GetCharts"] = partial(
+                test_query, client, FETCH_ALL_CHARTS_TEST.replace('"', "`")
+            )
+        else:
+            test_fn["GetCharts"] = partial(test_query, client, FETCH_ALL_CHARTS_TEST)
 
     test_connection_steps(
         metadata=metadata,

metadata/ingestion/source/dashboard/superset/db_source.py

@@ -26,6 +26,9 @@ from metadata.generated.schema.api.data.createDashboardDataModel import (
 from metadata.generated.schema.entity.data.chart import Chart
 from metadata.generated.schema.entity.data.dashboardDataModel import DataModelType
 from metadata.generated.schema.entity.data.table import Table
+from metadata.generated.schema.entity.services.connections.database.mysqlConnection import (
+    MysqlConnection,
+)
 from metadata.generated.schema.entity.services.databaseService import DatabaseService
 from metadata.generated.schema.entity.services.ingestionPipelines.status import (
     StackTraceError,
@@ -81,7 +84,10 @@ class SupersetDBSource(SupersetSourceMixin):
         the required information which is not available in fetch_charts_with_id api
         """
         try:
-            charts = self.engine.execute(FETCH_ALL_CHARTS)
+            if isinstance(self.service_connection.connection, MysqlConnection):
+                charts = self.engine.execute(FETCH_ALL_CHARTS.replace('"', "`"))
+            else:
+                charts = self.engine.execute(FETCH_ALL_CHARTS)
             for chart in charts:
                 chart_detail = FetchChart(**chart)
                 self.all_charts[chart_detail.id] = chart_detail

metadata/ingestion/source/dashboard/superset/mixin.py

@@ -17,7 +17,7 @@ from typing import Iterable, List, Optional, Union
 
 from metadata.generated.schema.api.lineage.addLineage import AddLineageRequest
 from metadata.generated.schema.entity.data.dashboardDataModel import DashboardDataModel
-from metadata.generated.schema.entity.data.table import Column, Table
+from metadata.generated.schema.entity.data.table import Column, DataType, Table
 from metadata.generated.schema.entity.services.connections.dashboard.supersetConnection import (
     SupersetConnection,
 )
@@ -247,6 +247,12 @@ class SupersetSourceMixin(DashboardServiceSource):
                     parsed_fields = Column(
                         dataTypeDisplay=field.type,
                         dataType=col_parse["dataType"],
+                        arrayDataType=DataType(col_parse["arrayDataType"])
+                        if col_parse.get("arrayDataType")
+                        else None,
+                        children=list(col_parse["children"])
+                        if col_parse.get("children")
+                        else None,
                         name=str(field.id),
                         displayName=field.column_name,
                         description=field.description,

metadata/ingestion/source/database/couchbase/metadata.py

@@ -27,10 +27,7 @@ from metadata.ingestion.source.database.common_nosql_source import (
     SAMPLE_SIZE,
     CommonNoSQLSource,
 )
-from metadata.ingestion.source.database.couchbase.queries import (
-    COUCHBASE_GET_DATA,
-    COUCHBASE_SQL_STATEMENT,
-)
+from metadata.ingestion.source.database.couchbase.queries import COUCHBASE_GET_DATA
 from metadata.utils.logger import ingestion_logger
 
 logger = ingestion_logger()
@@ -111,18 +108,14 @@ class CouchbaseSource(CommonNoSQLSource):
         """
         try:
             database_name = self.context.get().database
-            query = COUCHBASE_SQL_STATEMENT.format(table_name=table_name)
-            result = self.couchbase.query(query)
-            for row in result.rows():
-                if len(row) > 0:
-                    query_coln = COUCHBASE_GET_DATA.format(
-                        database_name=database_name,
-                        schema_name=schema_name,
-                        table_name=table_name,
-                        sample_size=SAMPLE_SIZE,
-                    )
-                    query_iter = self.couchbase.query(query_coln)
-                    return list(query_iter.rows())
+            query_coln = COUCHBASE_GET_DATA.format(
+                database_name=database_name,
+                schema_name=schema_name,
+                table_name=table_name,
+                sample_size=SAMPLE_SIZE,
+            )
+            query_iter = self.couchbase.query(query_coln)
+            return list(query_iter.rows())
         except Exception as exp:
             logger.debug(f"Failed to list column names for table [{table_name}]: {exp}")
             logger.debug(traceback.format_exc())

metadata/ingestion/source/database/couchbase/queries.py

@@ -14,10 +14,6 @@ SQL Queries used during ingestion
 
 import textwrap
 
-COUCHBASE_SQL_STATEMENT = textwrap.dedent(
-    """ SELECT * FROM system:indexes WHERE keyspace_id = '{table_name}' AND is_primary = TRUE """
-)
-
 COUCHBASE_GET_DATA = textwrap.dedent(
     """ select crc.* from `{database_name}`.`{schema_name}`.`{table_name}` crc limit {sample_size} """
 )

metadata/ingestion/source/database/dbt/constants.py

@@ -22,6 +22,61 @@ REQUIRED_MANIFEST_KEYS = ["name", "schema", "resource_type"]
 # Based on https://schemas.getdbt.com/dbt/catalog/v1.json
 REQUIRED_CATALOG_KEYS = ["name", "type", "index"]
 
+REQUIRED_RESULTS_KEYS = {
+    "status",
+    "timing",
+    "thread_id",
+    "execution_time",
+    "message",
+    "adapter_response",
+    "unique_id",
+}
+
+REQUIRED_NODE_KEYS = {
+    "schema_",
+    "schema",
+    "freshness",
+    "name",
+    "resource_type",
+    "path",
+    "unique_id",
+    "source_name",
+    "source_description",
+    "source_meta",
+    "loader",
+    "identifier",
+    "relation_name",
+    "fqn",
+    "alias",
+    "checksum",
+    "config",
+    "column_name",
+    "test_metadata",
+    "original_file_path",
+    "root_path",
+    "database",
+    "tags",
+    "description",
+    "columns",
+    "meta",
+    "owner",
+    "created_at",
+    "group",
+    "sources",
+    "compiled",
+    "docs",
+    "version",
+    "latest_version",
+    "package_name",
+    "depends_on",
+    "compiled_code",
+    "compiled_sql",
+    "raw_code",
+    "raw_sql",
+    "language",
+}
+
+
 NONE_KEYWORDS_LIST = ["none", "null"]
 
 DBT_CATALOG_FILE_NAME = "catalog.json"

metadata/ingestion/source/database/dbt/dbt_service.py

@@ -13,7 +13,7 @@ DBT service Topology.
 """
 
 from abc import ABC, abstractmethod
-from typing import Iterable
+from typing import Iterable, List
 
 from dbt_artifacts_parser.parser import parse_catalog, parse_manifest, parse_run_results
 from pydantic import Field
@@ -37,6 +37,10 @@ from metadata.ingestion.models.topology import (
     TopologyNode,
 )
 from metadata.ingestion.source.database.database_service import DataModelLink
+from metadata.ingestion.source.database.dbt.constants import (
+    REQUIRED_NODE_KEYS,
+    REQUIRED_RESULTS_KEYS,
+)
 from metadata.ingestion.source.database.dbt.dbt_config import get_dbt_details
 from metadata.ingestion.source.database.dbt.models import (
     DbtFiles,
@@ -169,51 +173,27 @@ class DbtServiceSource(TopologyRunnerMixin, Source, ABC):
             }
         )
 
-        required_nodes_keys = {
-            "schema_",
-            "schema",
-            "name",
-            "resource_type",
-            "path",
-            "unique_id",
-            "fqn",
-            "alias",
-            "checksum",
-            "config",
-            "column_name",
-            "test_metadata",
-            "original_file_path",
-            "root_path",
-            "database",
-            "tags",
-            "description",
-            "columns",
-            "meta",
-            "owner",
-            "created_at",
-            "group",
-            "sources",
-            "compiled",
-            "docs",
-            "version",
-            "latest_version",
-            "package_name",
-            "depends_on",
-            "compiled_code",
-            "compiled_sql",
-            "raw_code",
-            "raw_sql",
-            "language",
-        }
+        for field in ["nodes", "sources"]:
+            for node, value in manifest_dict.get(  # pylint: disable=unused-variable
+                field
+            ).items():
+                keys_to_delete = [
+                    key for key in value if key.lower() not in REQUIRED_NODE_KEYS
+                ]
+                for key in keys_to_delete:
+                    del value[key]
 
-        for node, value in manifest_dict.get(  # pylint: disable=unused-variable
-            "nodes"
-        ).items():
-            keys_to_delete = [
-                key for key in value if key.lower() not in required_nodes_keys
-            ]
-            for key in keys_to_delete:
-                del value[key]
+    def remove_run_result_non_required_keys(self, run_results: List[dict]):
+        """
+        Method to remove the non required keys from run results file
+        """
+        for run_result in run_results:
+            for result in run_result.get("results"):
+                keys_to_delete = [
+                    key for key in result if key.lower() not in REQUIRED_RESULTS_KEYS
+                ]
+                for key in keys_to_delete:
+                    del result[key]
 
     def get_dbt_files(self) -> Iterable[DbtFiles]:
         dbt_files = get_dbt_details(self.source_config.dbtConfigSource)
@@ -225,6 +205,10 @@ class DbtServiceSource(TopologyRunnerMixin, Source, ABC):
         self.remove_manifest_non_required_keys(
             manifest_dict=self.context.get().dbt_file.dbt_manifest
         )
+        if self.context.get().dbt_file.dbt_run_results:
+            self.remove_run_result_non_required_keys(
+                run_results=self.context.get().dbt_file.dbt_run_results
+            )
         dbt_objects = DbtObjects(
             dbt_catalog=parse_catalog(self.context.get().dbt_file.dbt_catalog)
             if self.context.get().dbt_file.dbt_catalog

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

@@ -12,6 +12,7 @@
 """
 Source connection handler
 """
+from copy import deepcopy
 from typing import Optional
 
 from sqlalchemy.engine import Engine
@@ -26,6 +27,7 @@ from metadata.ingestion.connections.builders import (
     create_generic_db_connection,
     get_connection_args_common,
     get_connection_url_common,
+    init_empty_connection_arguments,
 )
 from metadata.ingestion.connections.test_connections import test_connection_db_common
 from metadata.ingestion.ometa.ometa_api import OpenMetadata
@@ -41,8 +43,16 @@ def get_connection(connection: PinotDBConnection) -> Engine:
     """
     Create connection
     """
+    # TODO: Rename database field to DatabaseSchema
+    # Pinot does not support multi database concept
+    if connection.database is not None:
+        if not connection.connectionArguments:
+            connection.connectionArguments = init_empty_connection_arguments()
+        connection.connectionArguments.root["database"] = connection.database
+    connection_copy = deepcopy(connection)
+    connection_copy.database = None
     return create_generic_db_connection(
-        connection=connection,
+        connection=connection_copy,
         get_connection_url_fn=get_connection_url,
         get_connection_args_fn=get_connection_args_common,
     )

metadata/ingestion/source/database/pinotdb/metadata.py

@@ -9,7 +9,10 @@
 #  See the License for the specific language governing permissions and
 #  limitations under the License.
 """PinotDb source module"""
-from typing import Optional
+from typing import Iterable, Optional
+
+from pinotdb import sqlalchemy as pinot_sqlalchemy
+from sqlalchemy import types
 
 from metadata.generated.schema.entity.services.connections.database.pinotDBConnection import (
     PinotDBConnection,
@@ -22,6 +25,31 @@ from metadata.ingestion.ometa.ometa_api import OpenMetadata
 from metadata.ingestion.source.database.common_db_source import CommonDbSourceService
 
 
+def get_type_custom(data_type, field_size):
+    type_map = {
+        "int": types.BigInteger,
+        "long": types.BigInteger,
+        "float": types.Float,
+        "double": types.Numeric,
+        # BOOLEAN, is added after release 0.7.1.
+        # In release 0.7.1 and older releases, BOOLEAN is equivalent to STRING.
+        "boolean": types.Boolean,
+        "timestamp": types.TIMESTAMP,
+        "string": types.String,
+        "json": types.JSON,
+        "bytes": types.LargeBinary,
+        "big_decimal": types.DECIMAL,
+        # Complex types
+        "struct": types.BLOB,
+        "map": types.BLOB,
+        "array": types.ARRAY,
+    }
+    return type_map.get(data_type.lower())
+
+
+pinot_sqlalchemy.get_type = get_type_custom
+
+
 class PinotdbSource(CommonDbSourceService):
     """
     Implements the necessary methods to extract
@@ -39,3 +67,15 @@ class PinotdbSource(CommonDbSourceService):
                 f"Expected PinotdbConnection, but got {connection}"
             )
         return cls(config, metadata)
+
+    def get_database_names(self) -> Iterable[str]:
+        """
+        Default case with a single database.
+
+        It might come informed - or not - from the source.
+
+        Sources with multiple databases should overwrite this and
+        apply the necessary filters.
+        """
+        # TODO: Add databaseDisplayName field in PinotDBConnection
+        yield "default"