openmetadata-ingestion 1.5.6.0__py3-none-any.whl → 1.5.7.1__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-24T15:21:08+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-24T15:21:08+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-24T15:21:08+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-24T15:21:08+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-24T15:21:08+00:00
 
 from __future__ import annotations
 

metadata/ingestion/source/dashboard/mstr/client.py

@@ -30,6 +30,7 @@ from metadata.ingestion.source.dashboard.mstr.models import (
     MstrSearchResult,
     MstrSearchResultList,
 )
+from metadata.utils.helpers import clean_uri
 from metadata.utils.logger import ingestion_logger
 
 logger = ingestion_logger()
@@ -44,8 +45,8 @@ class MSTRClient:
 
     def _get_base_url(self, path=None):
         if not path:
-            return f"{self.config.hostPort}/{API_VERSION}/"
-        return f"{self.config.hostPort}/{API_VERSION}/{path}"
+            return f"{clean_uri(self.config.hostPort)}/{API_VERSION}/"
+        return f"{clean_uri(self.config.hostPort)}/{API_VERSION}/{path}"
 
     def _get_mstr_session(self) -> MSTRRESTSession:
         try:

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

@@ -11,7 +11,7 @@
 """
 Couchbase source methods.
 """
-
+import re
 import traceback
 from typing import Dict, Iterable, List, Optional
 
@@ -27,14 +27,17 @@ from metadata.ingestion.source.database.common_nosql_source import (
     SAMPLE_SIZE,
     CommonNoSQLSource,
 )
+from metadata.ingestion.source.database.couchbase.models import IndexObject as Index
 from metadata.ingestion.source.database.couchbase.queries import (
     COUCHBASE_GET_DATA,
-    COUCHBASE_SQL_STATEMENT,
+    COUCHBASE_GET_INDEX_KEYS,
 )
 from metadata.utils.logger import ingestion_logger
 
 logger = ingestion_logger()
 
+DEFAULT_SCHEMA_NAME = "_default"
+
 
 class CouchbaseSource(CommonNoSQLSource):
     """
@@ -45,6 +48,7 @@ class CouchbaseSource(CommonNoSQLSource):
     def __init__(self, config: WorkflowSource, metadata: OpenMetadata):
         super().__init__(config, metadata)
         self.couchbase = self.connection_obj
+        self.index_condition_map = {}
 
     @classmethod
     def create(
@@ -61,10 +65,12 @@ class CouchbaseSource(CommonNoSQLSource):
     def get_database_names(self) -> Iterable[str]:
         try:
             if self.service_connection.bucket:
+                self.index_condition_map.clear()
                 yield self.service_connection.__dict__.get("bucket")
             else:
                 buckets = self.couchbase.buckets()
                 for bucket_name in buckets.get_all_buckets():
+                    self.index_condition_map.clear()
                     yield bucket_name.name
         except Exception as exp:
             logger.debug(f"Failed to fetch bucket name: {exp}")
@@ -104,25 +110,72 @@ class CouchbaseSource(CommonNoSQLSource):
             logger.debug(traceback.format_exc())
         return []
 
+    def _is_valid_key(self, key: str) -> bool:
+        return bool(re.fullmatch(r"[a-zA-Z0-9._`\%]+", key))
+
+    def get_index_condition(self, schema_name: str) -> str:
+        """
+        Method to prepare query condition based on index
+        """
+        bucket_name = self.context.get().database
+        if self.index_condition_map.get((bucket_name, schema_name)):
+            return self.index_condition_map.get((bucket_name, schema_name))
+
+        index_condition = set()
+
+        if schema_name == DEFAULT_SCHEMA_NAME:
+            condition = f"keyspace_id = '{bucket_name}' "
+        else:
+            condition = f"bucket_id = '{bucket_name}' AND scope_id = '{schema_name}'"
+
+        query = COUCHBASE_GET_INDEX_KEYS.format(condition=condition)
+        result = self.couchbase.query(query)
+        for row in result.rows():
+            index_obj = Index(**dict(row))
+            if index_obj.indexes:
+                if index_obj.indexes.is_primary:
+                    self.index_condition_map[(bucket_name, schema_name)] = ""
+                    return ""
+                for key in index_obj.indexes.index_key or []:
+                    if self._is_valid_key(key):
+                        condition = ""
+                        if index_obj.indexes.condition:
+                            condition = f"AND {index_obj.indexes.condition}"
+                        index_condition.add(f"({key} is not missing {condition})")
+        if index_condition:
+            self.index_condition_map[
+                (bucket_name, schema_name)
+            ] = "WHERE " + " OR ".join(index_condition)
+            return self.index_condition_map[(bucket_name, schema_name)]
+
+        self.index_condition_map[(bucket_name, schema_name)] = ""
+        return ""
+
     def get_table_columns_dict(self, schema_name: str, table_name: str) -> List[Dict]:
         """
         Method to get actual data available within table
         need to be overridden by sources
         """
+        from couchbase.exceptions import QueryIndexNotFoundException
+
         try:
+            condition = self.get_index_condition(schema_name)
             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,
+                condition=condition,
+            )
+            query_iter = self.couchbase.query(query_coln)
+            return list(query_iter.rows())
+        except QueryIndexNotFoundException as exp:
+            logger.warning(
+                f"Fetching columns failed for [`{database_name}`.`{schema_name}`.`{table_name}`],"
+                " check if the index is created for the table or data exists in the table"
+            )
+            logger.debug(traceback.format_exc())
         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/models.py

@@ -0,0 +1,31 @@
+#  Copyright 2024 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.
+"""
+Couchbase source models.
+"""
+
+from typing import List, Optional
+
+from pydantic import BaseModel
+
+
+class IndexKey(BaseModel):
+    """A Bigtable index key."""
+
+    index_key: List[str] = []
+    condition: Optional[str] = None
+    is_primary: Optional[bool] = False
+
+
+class IndexObject(BaseModel):
+    """A Bigtable cell value."""
+
+    indexes: Optional[IndexKey] = None

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

@@ -14,10 +14,11 @@ 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_INDEX_KEYS = textwrap.dedent(
+    """ select * from system:indexes where {condition}"""
 )
 
+
 COUCHBASE_GET_DATA = textwrap.dedent(
-    """ select crc.* from `{database_name}`.`{schema_name}`.`{table_name}` crc limit {sample_size} """
+    """ select * from `{database_name}`.`{schema_name}`.`{table_name}` {condition} 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/glue/metadata.py

@@ -39,7 +39,11 @@ from metadata.generated.schema.metadataIngestion.databaseServiceMetadataPipeline
 from metadata.generated.schema.metadataIngestion.workflow import (
     Source as WorkflowSource,
 )
-from metadata.generated.schema.type.basic import EntityName, FullyQualifiedEntityName
+from metadata.generated.schema.type.basic import (
+    EntityName,
+    FullyQualifiedEntityName,
+    Markdown,
+)
 from metadata.ingestion.api.models import Either
 from metadata.ingestion.api.steps import InvalidSourceException
 from metadata.ingestion.models.ometa_classification import OMetaTagAndClassification
@@ -79,6 +83,7 @@ class GlueSource(DatabaseServiceSource):
         self.glue = get_connection(self.service_connection)
 
         self.connection_obj = self.glue
+        self.schema_description_map = {}
         self.test_connection()
 
     @classmethod
@@ -191,6 +196,10 @@ class GlueSource(DatabaseServiceSource):
                     ):
                         self.status.filter(schema_fqn, "Schema Filtered Out")
                         continue
+                    if schema.Description:
+                        self.schema_description_map[schema.Name] = Markdown(
+                            schema.Description
+                        )
                     yield schema.Name
                 except Exception as exc:
                     self.status.failed(
@@ -211,6 +220,7 @@ class GlueSource(DatabaseServiceSource):
         yield Either(
             right=CreateDatabaseSchemaRequest(
                 name=EntityName(schema_name),
+                description=self.schema_description_map.get(schema_name),
                 database=FullyQualifiedEntityName(
                     fqn.build(
                         metadata=self.metadata,