openmetadata-ingestion 1.5.13.0__py3-none-any.whl → 1.6.0.0rc2__py3-none-any.whl

metadata/ingestion/source/database/common_db_source.py

@@ -11,13 +11,11 @@
 """
 Generic source to build SQL connectors.
 """
-import math
-import time
+import copy
 import traceback
 from abc import ABC
-from concurrent.futures import ThreadPoolExecutor
 from copy import deepcopy
-from typing import Any, Iterable, List, Optional, Tuple, Union, cast
+from typing import Any, Dict, Iterable, List, Optional, Tuple, cast
 
 from pydantic import BaseModel
 from sqlalchemy.engine import Connection
@@ -29,15 +27,14 @@ from metadata.generated.schema.api.data.createDatabase import CreateDatabaseRequ
 from metadata.generated.schema.api.data.createDatabaseSchema import (
     CreateDatabaseSchemaRequest,
 )
-from metadata.generated.schema.api.data.createQuery import CreateQueryRequest
 from metadata.generated.schema.api.data.createStoredProcedure import (
     CreateStoredProcedureRequest,
 )
 from metadata.generated.schema.api.data.createTable import CreateTableRequest
-from metadata.generated.schema.api.lineage.addLineage import AddLineageRequest
 from metadata.generated.schema.entity.data.database import Database
 from metadata.generated.schema.entity.data.databaseSchema import DatabaseSchema
 from metadata.generated.schema.entity.data.table import (
+    Column,
     ConstraintType,
     Table,
     TableConstraint,
@@ -62,8 +59,7 @@ from metadata.generated.schema.type.basic import (
 from metadata.ingestion.api.models import Either
 from metadata.ingestion.connections.session import create_and_bind_thread_safe_session
 from metadata.ingestion.models.ometa_classification import OMetaTagAndClassification
-from metadata.ingestion.models.ometa_lineage import OMetaLineageRequest
-from metadata.ingestion.models.topology import Queue
+from metadata.ingestion.models.patch_request import PatchedEntity, PatchRequest
 from metadata.ingestion.ometa.ometa_api import OpenMetadata
 from metadata.ingestion.source.connections import (
     get_connection,
@@ -75,9 +71,8 @@ from metadata.ingestion.source.database.sqlalchemy_source import SqlAlchemySourc
 from metadata.ingestion.source.database.stored_procedures_mixin import QueryByProcedure
 from metadata.ingestion.source.models import TableView
 from metadata.utils import fqn
-from metadata.utils.db_utils import get_view_lineage
+from metadata.utils.constraints import get_relationship_type
 from metadata.utils.execution_time_tracker import (
-    ExecutionTimeTrackerContextMap,
     calculate_execution_time,
     calculate_execution_time_generator,
 )
@@ -88,6 +83,13 @@ from metadata.utils.ssl_manager import SSLManager, check_ssl_and_init
 logger = ingestion_logger()
 
 
+class ColumnAndReferredColumn(BaseModel):
+    table_name: str
+    schema_name: str
+    db_name: Optional[str]
+    column: Dict
+
+
 class TableNameAndType(BaseModel):
     """
     Helper model for passing down
@@ -141,6 +143,7 @@ class CommonDbSourceService(
         self.database_source_state = set()
         self.context.get_global().table_views = []
         self.context.get_global().table_constrains = []
+        self.context.get_global().foreign_tables = []
         self.context.set_threads(self.source_config.threads)
         super().__init__()
 
@@ -471,13 +474,6 @@ class CommonDbSourceService(
     def get_stored_procedure_queries(self) -> Iterable[QueryByProcedure]:
         """Not Implemented"""
 
-    @calculate_execution_time_generator()
-    def yield_procedure_lineage_and_queries(
-        self,
-    ) -> Iterable[Either[Union[AddLineageRequest, CreateQueryRequest]]]:
-        """Not Implemented"""
-        yield from []
-
     def get_location_path(self, table_name: str, schema_name: str) -> Optional[str]:
         """
         Method to fetch the location path of the table
@@ -518,7 +514,12 @@ class CommonDbSourceService(
             )
 
             table_constraints = self.update_table_constraints(
-                table_constraints, foreign_columns
+                schema_name=schema_name,
+                table_name=table_name,
+                db_name=self.context.get().database,
+                table_constraints=table_constraints,
+                foreign_columns=foreign_columns,
+                columns=columns,
             )
 
             description = (
@@ -601,107 +602,74 @@ class CommonDbSourceService(
                 )
             )
 
-    def multithread_process_view_lineage(self) -> Iterable[Either[OMetaLineageRequest]]:
-        """Multithread Processing of a Node"""
-
-        views_list = list(self.context.get().table_views or [])
-        views_length = len(views_list)
-
-        if views_length != 0:
-            chunksize = int(math.ceil(views_length / self.source_config.threads))
-            chunks = [
-                views_list[i : i + chunksize] for i in range(0, views_length, chunksize)
-            ]
-
-            thread_pool = ThreadPoolExecutor(max_workers=self.source_config.threads)
-            queue = Queue()
-
-            futures = [
-                thread_pool.submit(
-                    self._process_view_def_chunk,
-                    chunk,
-                    queue,
-                    self.context.get_current_thread_id(),
-                )
-                for chunk in chunks
-            ]
-
-            while True:
-                if queue.has_tasks():
-                    yield from queue.process()
-
-                else:
-                    if not futures:
-                        break
-
-                    for i, future in enumerate(futures):
-                        if future.done():
-                            future.result()
-                            futures.pop(i)
-
-                time.sleep(0.01)
-
-    def _process_view_def_chunk(
-        self, chunk: List[TableView], queue: Queue, thread_id: int
-    ) -> None:
+    def _prepare_foreign_constraints(  # pylint: disable=too-many-arguments, too-many-locals
+        self,
+        supports_database: bool,
+        column: Dict,
+        table_name: str,
+        schema_name: str,
+        db_name: str,
+        columns: List[Column],
+        add_to_global: bool = True,
+    ):
         """
-        Process a chunk of view definitions
-        """
-        self.context.copy_from(thread_id)
-        ExecutionTimeTrackerContextMap().copy_from_parent(thread_id)
-        for view in [v for v in chunk if v.view_definition is not None]:
-            for lineage in get_view_lineage(
-                view=view,
-                metadata=self.metadata,
-                service_name=self.context.get().database_service,
-                connection_type=self.service_connection.type.value,
-                timeout_seconds=self.source_config.queryParsingTimeoutLimit,
-            ):
-                if lineage.right is not None:
-                    queue.put(
-                        Either(
-                            right=OMetaLineageRequest(
-                                lineage_request=lineage.right,
-                                override_lineage=self.source_config.overrideViewLineage,
-                            )
-                        )
-                    )
-                else:
-                    queue.put(lineage)
-
-    def _process_view_def_serial(self) -> Iterable[Either[OMetaLineageRequest]]:
-        """
-        Process view definitions serially
-        """
-        for view in [
-            v for v in self.context.get().table_views if v.view_definition is not None
-        ]:
-            for lineage in get_view_lineage(
-                view=view,
-                metadata=self.metadata,
-                service_name=self.context.get().database_service,
-                connection_type=self.service_connection.type.value,
-                timeout_seconds=self.source_config.queryParsingTimeoutLimit,
-            ):
-                if lineage.right is not None:
-                    yield Either(
-                        right=OMetaLineageRequest(
-                            lineage_request=lineage.right,
-                            override_lineage=self.source_config.overrideViewLineage,
-                        )
-                    )
-                else:
-                    yield lineage
-
-    @calculate_execution_time_generator()
-    def yield_view_lineage(self) -> Iterable[Either[OMetaLineageRequest]]:
-        logger.info("Processing Lineage for Views")
-        if self.source_config.threads > 1:
-            yield from self.multithread_process_view_lineage()
+        Method to prepare the foreign constraints
+        """
+        referred_column_fqns = []
+        if supports_database:
+            database_name = column.get("referred_database")
+        else:
+            database_name = self.context.get().database
+        referred_table_fqn = fqn.build(
+            metadata=self.metadata,
+            entity_type=Table,
+            table_name=column.get("referred_table"),
+            schema_name=column.get("referred_schema"),
+            database_name=database_name,
+            service_name=self.context.get().database_service,
+        )
+        referred_table = self.metadata.get_by_name(entity=Table, fqn=referred_table_fqn)
+        if referred_table:
+            for referred_column in column.get("referred_columns"):
+                col_fqn = fqn._build(  # pylint: disable=protected-access
+                    referred_table_fqn, referred_column, quote=False
+                )
+                if col_fqn:
+                    referred_column_fqns.append(FullyQualifiedEntityName(col_fqn))
         else:
-            yield from self._process_view_def_serial()
+            if add_to_global:
+                column_and_referred_columns = ColumnAndReferredColumn(
+                    table_name=table_name,
+                    schema_name=schema_name,
+                    db_name=db_name,
+                    column=column,
+                )
+                self.context.get_global().foreign_tables.append(
+                    column_and_referred_columns
+                )
+            return None
+        relationship_type = None
+        if referred_table:
+            relationship_type = get_relationship_type(
+                column,  # sqlalchemy foreign column
+                referred_table.columns,  # referred table columns
+                columns,  # current table om columns
+            )
+        return TableConstraint(
+            constraintType=ConstraintType.FOREIGN_KEY,
+            columns=column.get("constrained_columns"),
+            referredColumns=referred_column_fqns,
+            relationshipType=relationship_type,
+        )
 
-    def _get_foreign_constraints(self, foreign_columns) -> List[TableConstraint]:
+    def _get_foreign_constraints(
+        self,
+        table_name,
+        schema_name,
+        db_name,
+        foreign_columns: List[Dict],
+        columns: List[Column],
+    ) -> List[TableConstraint]:
         """
         Search the referred table for foreign constraints
         and get referred column fqn
@@ -710,48 +678,31 @@ class CommonDbSourceService(
 
         foreign_constraints = []
         for column in foreign_columns:
-            referred_column_fqns = []
-            if supports_database:
-                database_name = column.get("referred_database")
-            else:
-                database_name = self.context.get().database
-            referred_table_fqn = fqn.build(
-                metadata=self.metadata,
-                entity_type=Table,
-                table_name=column.get("referred_table"),
-                schema_name=column.get("referred_schema"),
-                database_name=database_name,
-                service_name=self.context.get().database_service,
-            )
-            if referred_table_fqn:
-                for referred_column in column.get("referred_columns"):
-                    col_fqn = fqn._build(
-                        referred_table_fqn, referred_column, quote=False
-                    )
-                    if col_fqn:
-                        referred_column_fqns.append(FullyQualifiedEntityName(col_fqn))
-            else:
-                # do not build partial foreign constraint. It will updated in next run.
-                continue
-            foreign_constraints.append(
-                TableConstraint(
-                    constraintType=ConstraintType.FOREIGN_KEY,
-                    columns=column.get("constrained_columns"),
-                    referredColumns=referred_column_fqns,
-                )
+            foreign_constraint = self._prepare_foreign_constraints(
+                supports_database, column, table_name, schema_name, db_name, columns
             )
+            if foreign_constraint:
+                foreign_constraints.append(foreign_constraint)
 
         return foreign_constraints
 
     @calculate_execution_time()
     def update_table_constraints(
-        self, table_constraints, foreign_columns
+        self,
+        table_name,
+        schema_name,
+        db_name,
+        table_constraints,
+        foreign_columns,
+        columns,
     ) -> List[TableConstraint]:
         """
         From topology.
         process the table constraints of all tables
         """
-        foreign_table_constraints = self._get_foreign_constraints(foreign_columns)
+        foreign_table_constraints = self._get_foreign_constraints(
+            table_name, schema_name, db_name, foreign_columns, columns
+        )
         if foreign_table_constraints:
             if table_constraints:
                 table_constraints.extend(foreign_table_constraints)
@@ -820,3 +771,55 @@ class CommonDbSourceService(
         """
         By default the source url is not supported for
         """
+
+    def yield_table_constraints(self) -> Iterable[Either[PatchedEntity]]:
+        """
+        Process remaining table constraints by patching the table
+        """
+        supports_database = hasattr(self.service_connection, "supportsDatabase")
+
+        for foreign_table in self.context.get_global().foreign_tables or []:
+            try:
+                foreign_constraints = []
+                table_fqn = fqn.build(
+                    metadata=self.metadata,
+                    entity_type=Table,
+                    service_name=self.context.get().database_service,
+                    database_name=foreign_table.db_name,
+                    schema_name=foreign_table.schema_name,
+                    table_name=foreign_table.table_name,
+                )
+                table = self.metadata.get_by_name(entity=Table, fqn=table_fqn)
+                if table:
+                    foreign_constraint = self._prepare_foreign_constraints(
+                        supports_database,
+                        foreign_table.column,
+                        foreign_table.table_name,
+                        foreign_table.schema_name,
+                        foreign_table.db_name,
+                        table.columns,
+                        False,
+                    )
+                    if foreign_constraint:
+                        foreign_constraints.append(foreign_constraint)
+
+                # send the patch request
+                if foreign_constraints:
+                    new_entity = copy.deepcopy(table)
+                    new_entity.tableConstraints = (
+                        new_entity.tableConstraints or []
+                    ) + foreign_constraints
+                    patch_request = PatchRequest(
+                        original_entity=table,
+                        new_entity=new_entity,
+                        override_metadata=True,
+                    )
+                    yield Either(right=patch_request)
+            except Exception as exc:
+                yield Either(
+                    left=StackTraceError(
+                        name=str(foreign_table.table_name),
+                        error=f"Error to yield tableConstraints for {str(foreign_table.table_name)}: {exc}",
+                        stackTrace=traceback.format_exc(),
+                    )
+                )

metadata/ingestion/source/database/common_nosql_source.py

@@ -20,12 +20,10 @@ from metadata.generated.schema.api.data.createDatabase import CreateDatabaseRequ
 from metadata.generated.schema.api.data.createDatabaseSchema import (
     CreateDatabaseSchemaRequest,
 )
-from metadata.generated.schema.api.data.createQuery import CreateQueryRequest
 from metadata.generated.schema.api.data.createStoredProcedure import (
     CreateStoredProcedureRequest,
 )
 from metadata.generated.schema.api.data.createTable import CreateTableRequest
-from metadata.generated.schema.api.lineage.addLineage import AddLineageRequest
 from metadata.generated.schema.entity.data.database import Database
 from metadata.generated.schema.entity.data.databaseSchema import DatabaseSchema
 from metadata.generated.schema.entity.data.table import (
@@ -279,12 +277,6 @@ class CommonNoSQLSource(DatabaseServiceSource, ABC):
                 )
             )
 
-    def yield_view_lineage(self) -> Iterable[Either[AddLineageRequest]]:
-        """
-        views are not supported with NoSQL
-        """
-        yield from []
-
     def yield_tag(
         self, schema_name: str
     ) -> Iterable[Either[OMetaTagAndClassification]]:
@@ -303,12 +295,6 @@ class CommonNoSQLSource(DatabaseServiceSource, ABC):
     def get_stored_procedure_queries(self) -> Iterable[QueryByProcedure]:
         """Not Implemented"""
 
-    def yield_procedure_lineage_and_queries(
-        self,
-    ) -> Iterable[Either[Union[AddLineageRequest, CreateQueryRequest]]]:
-        """Not Implemented"""
-        yield from []
-
     def get_source_url(
         self,
         database_name: Optional[str] = None,

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/{data_quality/builders/sqa_validator_builder.py → ingestion/source/database/couchbase/models.py}

@@ -1,4 +1,4 @@
-#  Copyright 2021 Collate
+#  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
@@ -8,18 +8,24 @@
 #  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.
-
 """
-Builder defining the structure of builders for validators for SQA sources
+Couchbase source models.
 """
 
+from typing import List, Optional
+
+from pydantic import BaseModel
+
+
+class IndexKey(BaseModel):
+    """A Bigtable index key."""
 
-from metadata.data_quality.builders.i_validator_builder import IValidatorBuilder
+    index_key: List[str] = []
+    condition: Optional[str] = None
+    is_primary: Optional[bool] = False
 
 
-class SQAValidatorBuilder(IValidatorBuilder):
-    """Builder for SQA validators"""
+class IndexObject(BaseModel):
+    """A Bigtable cell value."""
 
-    def _get_source_type(self) -> str:
-        """Return the test case"""
-        return "sqlalchemy"
+    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/couchbase/service_spec.py

@@ -0,0 +1,4 @@
+from metadata.ingestion.source.database.couchbase.metadata import CouchbaseSource
+from metadata.utils.service_spec.default import DefaultDatabaseSpec
+
+ServiceSpec = DefaultDatabaseSpec(metadata_source_class=CouchbaseSource)