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

metadata/generated/schema/type/jdbcConnection.py

@@ -1,6 +1,6 @@
 # generated by datamodel-codegen:
 #   filename:  type/jdbcConnection.json
-#   timestamp: 2024-10-17T16:44:45+00:00
+#   timestamp: 2024-10-24T15:21:08+00:00
 
 from __future__ import annotations
 

metadata/generated/schema/type/lifeCycle.py

@@ -1,6 +1,6 @@
 # generated by datamodel-codegen:
 #   filename:  type/lifeCycle.json
-#   timestamp: 2024-10-17T16:44:45+00:00
+#   timestamp: 2024-10-24T15:21:08+00:00
 
 from __future__ import annotations
 

metadata/generated/schema/type/paging.py

@@ -1,6 +1,6 @@
 # generated by datamodel-codegen:
 #   filename:  type/paging.json
-#   timestamp: 2024-10-17T16:44:45+00:00
+#   timestamp: 2024-10-24T15:21:08+00:00
 
 from __future__ import annotations
 

metadata/generated/schema/type/profile.py

@@ -1,6 +1,6 @@
 # generated by datamodel-codegen:
 #   filename:  type/profile.json
-#   timestamp: 2024-10-17T16:44:45+00:00
+#   timestamp: 2024-10-24T15:21:08+00:00
 
 from __future__ import annotations
 

metadata/generated/schema/type/queryParserData.py

@@ -1,6 +1,6 @@
 # generated by datamodel-codegen:
 #   filename:  type/queryParserData.json
-#   timestamp: 2024-10-17T16:44:45+00:00
+#   timestamp: 2024-10-24T15:21:08+00:00
 
 from __future__ import annotations
 

metadata/generated/schema/type/reaction.py

@@ -1,6 +1,6 @@
 # generated by datamodel-codegen:
 #   filename:  type/reaction.json
-#   timestamp: 2024-10-17T16:44:45+00:00
+#   timestamp: 2024-10-24T15:21:08+00:00
 
 from __future__ import annotations
 

metadata/generated/schema/type/schedule.py

@@ -1,6 +1,6 @@
 # generated by datamodel-codegen:
 #   filename:  type/schedule.json
-#   timestamp: 2024-10-17T16:44:45+00:00
+#   timestamp: 2024-10-24T15:21:08+00:00
 
 from __future__ import annotations
 

metadata/generated/schema/type/schema.py

@@ -1,6 +1,6 @@
 # generated by datamodel-codegen:
 #   filename:  type/schema.json
-#   timestamp: 2024-10-17T16:44:45+00:00
+#   timestamp: 2024-10-24T15:21:08+00:00
 
 from __future__ import annotations
 

metadata/generated/schema/type/tableQuery.py

@@ -1,6 +1,6 @@
 # generated by datamodel-codegen:
 #   filename:  type/tableQuery.json
-#   timestamp: 2024-10-17T16:44:45+00:00
+#   timestamp: 2024-10-24T15:21:08+00:00
 
 from __future__ import annotations
 

metadata/generated/schema/type/tableUsageCount.py

@@ -1,6 +1,6 @@
 # generated by datamodel-codegen:
 #   filename:  type/tableUsageCount.json
-#   timestamp: 2024-10-17T16:44:45+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-17T16:44:45+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-17T16:44:45+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-17T16:44:45+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-17T16:44:45+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/database/couchbase/metadata.py

@@ -11,7 +11,7 @@
 """
 Couchbase source methods.
 """
-
+import re
 import traceback
 from typing import Dict, Iterable, List, Optional
 
@@ -27,11 +27,17 @@ from metadata.ingestion.source.database.common_nosql_source import (
     SAMPLE_SIZE,
     CommonNoSQLSource,
 )
-from metadata.ingestion.source.database.couchbase.queries import COUCHBASE_GET_DATA
+from metadata.ingestion.source.database.couchbase.models import IndexObject as Index
+from metadata.ingestion.source.database.couchbase.queries import (
+    COUCHBASE_GET_DATA,
+    COUCHBASE_GET_INDEX_KEYS,
+)
 from metadata.utils.logger import ingestion_logger
 
 logger = ingestion_logger()
 
+DEFAULT_SCHEMA_NAME = "_default"
+
 
 class CouchbaseSource(CommonNoSQLSource):
     """
@@ -42,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(
@@ -58,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}")
@@ -101,21 +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_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,6 +14,11 @@ SQL Queries used during ingestion
 
 import textwrap
 
+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/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,

metadata/ingestion/source/database/glue/models.py

@@ -20,6 +20,7 @@ from pydantic import BaseModel
 class GlueSchema(BaseModel):
     CatalogId: Optional[str] = None
     Name: str
+    Description: Optional[str] = None
 
 
 class DatabasePage(BaseModel):