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

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

@@ -0,0 +1,71 @@
+#  Copyright 2021 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.
+
+"""
+Source connection handler
+"""
+from typing import Optional
+
+from sqlalchemy.engine import Engine
+
+from metadata.generated.schema.entity.automations.workflow import (
+    Workflow as AutomationWorkflow,
+)
+from metadata.generated.schema.entity.services.connections.database.cockroachConnection import (
+    CockroachConnection,
+)
+from metadata.generated.schema.entity.services.connections.testConnectionResult import (
+    TestConnectionResult,
+)
+from metadata.ingestion.connections.builders import (
+    create_generic_db_connection,
+    get_connection_args_common,
+    get_connection_url_common,
+)
+from metadata.ingestion.connections.test_connections import test_connection_db_common
+from metadata.ingestion.ometa.ometa_api import OpenMetadata
+from metadata.ingestion.source.database.cockroach.queries import COCKROACH_GET_DATABASE
+from metadata.utils.constants import THREE_MIN
+
+
+def get_connection(connection: CockroachConnection) -> Engine:
+    """
+    Create connection
+    """
+    return create_generic_db_connection(
+        connection=connection,
+        get_connection_url_fn=get_connection_url_common,
+        get_connection_args_fn=get_connection_args_common,
+    )
+
+
+def test_connection(
+    metadata: OpenMetadata,
+    engine: Engine,
+    service_connection: CockroachConnection,
+    automation_workflow: Optional[AutomationWorkflow] = None,
+    timeout_seconds: Optional[int] = THREE_MIN,
+) -> TestConnectionResult:
+    """
+    Test connection. This can be executed either as part
+    of a metadata workflow or during an Automation Workflow
+    """
+    queries = {
+        "GetDatabases": COCKROACH_GET_DATABASE,
+    }
+    return test_connection_db_common(
+        metadata=metadata,
+        engine=engine,
+        service_connection=service_connection,
+        automation_workflow=automation_workflow,
+        queries=queries,
+        timeout_seconds=timeout_seconds,
+    )

metadata/ingestion/source/database/cockroach/metadata.py

@@ -0,0 +1,206 @@
+#  Copyright 2021 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.
+"""
+Cockroach source module
+"""
+import traceback
+from collections import namedtuple
+from typing import Iterable, Optional, Tuple
+
+from sqlalchemy import sql
+from sqlalchemy.dialects.postgresql.base import PGDialect
+
+from metadata.generated.schema.entity.data.database import Database
+from metadata.generated.schema.entity.data.table import (
+    PartitionColumnDetails,
+    PartitionIntervalTypes,
+    TablePartition,
+    TableType,
+)
+from metadata.generated.schema.entity.services.connections.database.cockroachConnection import (
+    CockroachConnection,
+)
+from metadata.generated.schema.entity.services.connections.metadata.openMetadataConnection import (
+    OpenMetadataConnection,
+)
+from metadata.generated.schema.metadataIngestion.workflow import (
+    Source as WorkflowSource,
+)
+from metadata.ingestion.api.steps import InvalidSourceException
+from metadata.ingestion.ometa.ometa_api import OpenMetadata
+from metadata.ingestion.source.database.cockroach.queries import (
+    COCKROACH_GET_DB_NAMES,
+    COCKROACH_GET_PARTITION_DETAILS,
+    COCKROACH_GET_TABLE_NAMES,
+    COCKROACH_GET_VIEW_NAMES,
+    COCKROACH_SCHEMA_COMMENTS,
+)
+from metadata.ingestion.source.database.common_db_source import (
+    CommonDbSourceService,
+    TableNameAndType,
+)
+from metadata.ingestion.source.database.common_pg_mappings import (
+    INTERVAL_TYPE_MAP,
+    RELKIND_MAP,
+    ischema_names,
+)
+from metadata.ingestion.source.database.multi_db_source import MultiDBSource
+from metadata.utils import fqn
+from metadata.utils.filters import filter_by_database
+from metadata.utils.importer import import_side_effects
+from metadata.utils.logger import ingestion_logger
+
+import_side_effects(
+    "metadata.ingestion.source.database.postgres.converter_orm",
+    "metadata.ingestion.source.database.postgres.metrics",
+)
+
+TableKey = namedtuple("TableKey", ["schema", "table_name"])
+
+logger = ingestion_logger()
+
+PGDialect.ischema_names = ischema_names
+
+
+class CockroachSource(CommonDbSourceService, MultiDBSource):
+    """
+    Implements the necessary methods to extract
+    Database metadata from Cockroach Source
+    """
+
+    def __init__(self, config: WorkflowSource, metadata: OpenMetadata):
+        super().__init__(config, metadata)
+        self.schema_desc_map = {}
+
+    @classmethod
+    def create(
+        cls,
+        config_dict,
+        metadata: OpenMetadataConnection,
+        pipeline_name: Optional[str] = None,
+    ):
+        config: WorkflowSource = WorkflowSource.model_validate(config_dict)
+        connection: CockroachConnection = config.serviceConnection.root.config
+        if not isinstance(connection, CockroachConnection):
+            raise InvalidSourceException(
+                f"Expected CockroachConnection, but got {connection}"
+            )
+        return cls(config, metadata)
+
+    def set_schema_description_map(self) -> None:
+        self.schema_desc_map.clear()
+        results = self.engine.execute(COCKROACH_SCHEMA_COMMENTS).all()
+        for row in results:
+            self.schema_desc_map[(row.database_name, row.schema_name)] = row.comment
+
+    def get_schema_description(self, schema_name: str) -> Optional[str]:
+        """
+        Method to fetch the schema description
+        """
+        return self.schema_desc_map.get((self.context.get().database, schema_name))
+
+    def query_table_names_and_types(
+        self, schema_name: str
+    ) -> Iterable[TableNameAndType]:
+        """
+        Overwrite the inspector implementation to handle partitioned
+        and foreign types
+        """
+        result = self.connection.execute(
+            sql.text(COCKROACH_GET_TABLE_NAMES),
+            {"schema": schema_name},
+        )
+        return [
+            TableNameAndType(
+                name=name, type_=RELKIND_MAP.get(relkind, TableType.Regular)
+            )
+            for name, relkind in result
+        ]
+
+    def query_view_names_and_types(
+        self, schema_name: str
+    ) -> Iterable[TableNameAndType]:
+        result = self.connection.execute(
+            sql.text(COCKROACH_GET_VIEW_NAMES),
+            {"schema": schema_name},
+        )
+        return [
+            TableNameAndType(
+                name=name, type_=RELKIND_MAP.get(relkind, TableType.Regular)
+            )
+            for name, relkind in result
+        ]
+
+    def get_configured_database(self) -> Optional[str]:
+        if not self.service_connection.ingestAllDatabases:
+            return self.service_connection.database
+        return None
+
+    def get_database_names_raw(self) -> Iterable[str]:
+        yield from self._execute_database_query(COCKROACH_GET_DB_NAMES)
+
+    def get_database_names(self) -> Iterable[str]:
+        if not self.config.serviceConnection.root.config.ingestAllDatabases:
+            configured_db = self.config.serviceConnection.root.config.database
+            self.set_inspector(database_name=configured_db)
+            self.set_schema_description_map()
+            yield configured_db
+        else:
+            for new_database in self.get_database_names_raw():
+                database_fqn = fqn.build(
+                    self.metadata,
+                    entity_type=Database,
+                    service_name=self.context.get().database_service,
+                    database_name=new_database,
+                )
+
+                if filter_by_database(
+                    self.source_config.databaseFilterPattern,
+                    (
+                        database_fqn
+                        if self.source_config.useFqnForFiltering
+                        else new_database
+                    ),
+                ):
+                    self.status.filter(database_fqn, "Database Filtered Out")
+                    continue
+
+                try:
+                    self.set_inspector(database_name=new_database)
+                    self.set_schema_description_map()
+                    yield new_database
+                except Exception as exc:
+                    logger.debug(traceback.format_exc())
+                    logger.error(
+                        f"Error trying to connect to database {new_database}: {exc}"
+                    )
+
+    def get_table_partition_details(
+        self, table_name: str, schema_name: str, inspector
+    ) -> Tuple[bool, TablePartition]:
+        result = self.engine.execute(
+            COCKROACH_GET_PARTITION_DETAILS, table_name=table_name
+        ).all()
+        if result:
+            partition_details = TablePartition(
+                columns=[
+                    PartitionColumnDetails(
+                        columnName=row[1],
+                        intervalType=INTERVAL_TYPE_MAP.get(
+                            row[2], PartitionIntervalTypes.COLUMN_VALUE
+                        ),
+                        interval=None,
+                    )
+                    for row in result
+                ]
+            )
+            return True, partition_details
+        return False, None

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

@@ -0,0 +1,90 @@
+#  Copyright 2021 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.
+"""
+SQL Queries used during ingestion
+"""
+
+from sqlalchemy import text
+
+COCKROACH_GET_TABLE_NAMES = """
+    SELECT 
+    c.relname AS table_name,
+    c.relkind AS relkind
+    FROM 
+        pg_class c
+    JOIN 
+        pg_namespace n ON n.oid = c.relnamespace
+    WHERE 
+        n.nspname = :schema 
+        AND c.relkind IN ('r', 'p', 'f')
+    ORDER BY 
+        c.relname
+"""
+
+COCKROACH_GET_VIEW_NAMES = """
+    SELECT 
+    c.relname AS table_name,
+    c.relkind AS relkind
+    FROM 
+        pg_class c
+    JOIN 
+        pg_namespace n ON n.oid = c.relnamespace
+    WHERE 
+        n.nspname = :schema 
+        AND c.relkind IN ('v')
+    ORDER BY 
+        c.relname
+"""
+
+
+COCKROACH_SCHEMA_COMMENTS = """
+    SELECT
+    current_database() AS database_name,
+    n.nspname AS schema_name,
+    d.description AS comment
+FROM 
+    pg_namespace n
+LEFT JOIN 
+    pg_description d 
+ON 
+    n.oid = d.objoid
+WHERE 
+    d.objsubid = 0;
+"""
+
+
+COCKROACH_GET_DATABASE = text(
+    """
+    select datname FROM pg_catalog.pg_database
+"""
+)
+
+COCKROACH_GET_DB_NAMES = """
+    select datname from pg_catalog.pg_database
+"""
+
+COCKROACH_GET_PARTITION_DETAILS = """
+    SELECT
+    partitions.name AS partition_name,
+    column_names,
+    CASE 
+        WHEN list_value IS NOT NULL THEN 'list'
+        ELSE 'range'
+    END AS partition_type,
+    tables.name AS table_name,
+    database_name
+FROM
+    crdb_internal.partitions
+JOIN
+    crdb_internal.tables ON partitions.table_id = tables.table_id
+WHERE
+    tables.name = %(table_name)s;
+"""

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

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

metadata/ingestion/source/database/common_db_source.py

@@ -77,6 +77,7 @@ from metadata.utils.execution_time_tracker import (
     calculate_execution_time_generator,
 )
 from metadata.utils.filters import filter_by_table
+from metadata.utils.helpers import retry_with_docker_host
 from metadata.utils.logger import ingestion_logger
 from metadata.utils.ssl_manager import SSLManager, check_ssl_and_init
 
@@ -108,6 +109,7 @@ class CommonDbSourceService(
     - fetch_column_tags implemented at SqlColumnHandler. Sources should override this when needed
     """
 
+    @retry_with_docker_host()
     def __init__(
         self,
         config: WorkflowSource,
@@ -193,6 +195,12 @@ class CommonDbSourceService(
         by default there will be no schema description
         """
 
+    def get_stored_procedure_description(self, stored_procedure: str) -> Optional[str]:
+        """
+        Method to fetch the stored procedure description
+        by default there will be no stored procedure description
+        """
+
     @calculate_execution_time_generator()
     def yield_database(
         self, database_name: str
@@ -501,6 +509,7 @@ class CommonDbSourceService(
                 foreign_columns,
             ) = self.get_columns_and_constraints(
                 schema_name=schema_name,
+                table_type=table_type,
                 table_name=table_name,
                 db_name=self.context.get().database,
                 inspector=self.inspector,

metadata/ingestion/source/database/common_nosql_source.py

@@ -16,6 +16,8 @@ import traceback
 from abc import ABC, abstractmethod
 from typing import Any, Dict, Iterable, List, Optional, Tuple, Union
 
+from pydantic import BaseModel
+
 from metadata.generated.schema.api.data.createDatabase import CreateDatabaseRequest
 from metadata.generated.schema.api.data.createDatabaseSchema import (
     CreateDatabaseSchemaRequest,
@@ -27,6 +29,7 @@ from metadata.generated.schema.api.data.createTable import CreateTableRequest
 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,
     Table,
     TableConstraint,
     TableType,
@@ -51,6 +54,7 @@ from metadata.utils import fqn
 from metadata.utils.constants import DEFAULT_DATABASE
 from metadata.utils.datalake.datalake_utils import DataFrameColumnParser
 from metadata.utils.filters import filter_by_schema, filter_by_table
+from metadata.utils.helpers import retry_with_docker_host
 from metadata.utils.logger import ingestion_logger
 from metadata.utils.ssl_manager import check_ssl_and_init
 
@@ -60,12 +64,23 @@ logger = ingestion_logger()
 SAMPLE_SIZE = 1000
 
 
+class TableNameAndType(BaseModel):
+    """
+    Helper model for passing down
+    names and types of tables
+    """
+
+    name: str
+    type_: TableType = TableType.Regular
+
+
 class CommonNoSQLSource(DatabaseServiceSource, ABC):
     """
     Implements the necessary methods to extract
     Database metadata from NoSQL source
     """
 
+    @retry_with_docker_host()
     def __init__(self, config: WorkflowSource, metadata: OpenMetadata):
         super().__init__()
         self.config = config
@@ -168,12 +183,22 @@ class CommonNoSQLSource(DatabaseServiceSource, ABC):
         )
 
     @abstractmethod
-    def get_table_name_list(self, schema_name: str) -> List[str]:
+    def query_table_names_and_types(
+        self, schema_name: str
+    ) -> Iterable[TableNameAndType]:
         """
-        Method to get list of table names available within schema db
+        Method to get list of table names and types available within schema db
         need to be overridden by sources
         """
 
+    def query_view_names_and_types(
+        self, schema_name: str
+    ) -> Iterable[TableNameAndType]:
+        """
+        Method to get list of materialized view names and types available within schema db
+        need to be overridden by sources if views are supported by the database.
+        """
+
     def get_tables_name_and_type(self) -> Optional[Iterable[Tuple[str, TableType]]]:
         """
         Handle table and views.
@@ -184,29 +209,46 @@ class CommonNoSQLSource(DatabaseServiceSource, ABC):
         :return: tables or views, depending on config
         """
         schema_name = self.context.get().database_schema
-        if self.source_config.includeTables:
-            for collection in self.get_table_name_list(schema_name):
-                table_name = collection
-                table_fqn = fqn.build(
-                    self.metadata,
-                    entity_type=Table,
-                    service_name=self.context.get().database_service,
-                    database_name=self.context.get().database,
-                    schema_name=self.context.get().database_schema,
-                    table_name=table_name,
-                )
-                if filter_by_table(
-                    self.source_config.tableFilterPattern,
-                    table_fqn if self.source_config.useFqnForFiltering else table_name,
-                ):
-                    self.status.filter(
-                        table_fqn,
-                        "Table Filtered Out",
-                    )
+
+        try:
+            for table_and_type_fn, flag in {
+                self.query_table_names_and_types: self.source_config.includeTables,
+                self.query_view_names_and_types: self.source_config.includeViews,
+            }.items():
+                if not flag:
                     continue
-                yield table_name, TableType.Regular
 
-    @abstractmethod
+                for table_and_type in table_and_type_fn(schema_name=schema_name) or []:
+                    table_name, table_type = table_and_type.name, table_and_type.type_
+
+                    table_fqn = fqn.build(
+                        self.metadata,
+                        entity_type=Table,
+                        service_name=self.context.get().database_service,
+                        database_name=self.context.get().database,
+                        schema_name=self.context.get().database_schema,
+                        table_name=table_name,
+                    )
+                    if filter_by_table(
+                        self.source_config.tableFilterPattern,
+                        (
+                            table_fqn
+                            if self.source_config.useFqnForFiltering
+                            else table_name
+                        ),
+                    ):
+                        self.status.filter(
+                            table_fqn,
+                            "Table Filtered Out",
+                        )
+                        continue
+                    yield table_name, table_type
+        except Exception as err:
+            logger.warning(
+                f"Fetching tables names failed for schema {schema_name} due to - {err}"
+            )
+            logger.debug(traceback.format_exc())
+
     def get_table_columns_dict(
         self, schema_name: str, table_name: str
     ) -> Union[List[Dict], Dict]:
@@ -224,6 +266,18 @@ class CommonNoSQLSource(DatabaseServiceSource, ABC):
         # pylint: disable=unused-argument
         return None
 
+    def get_table_columns(self, schema_name: str, table_name: str) -> List[Column]:
+        """
+        Method to return all columns of a table
+        """
+        import pandas as pd  # pylint: disable=import-outside-toplevel
+
+        df = pd.DataFrame.from_records(
+            list(self.get_table_columns_dict(schema_name, table_name))
+        )
+        column_parser = DataFrameColumnParser.create(df)
+        return column_parser.get_columns()
+
     def yield_table(
         self, table_name_and_type: Tuple[str, TableType]
     ) -> Iterable[Either[CreateTableRequest]]:
@@ -231,19 +285,14 @@ class CommonNoSQLSource(DatabaseServiceSource, ABC):
         From topology.
         Prepare a table request and pass it to the sink
         """
-        import pandas as pd  # pylint: disable=import-outside-toplevel
 
         table_name, table_type = table_name_and_type
         schema_name = self.context.get().database_schema
         try:
-            data = self.get_table_columns_dict(schema_name, table_name)
-            df = pd.DataFrame.from_records(list(data))
-            column_parser = DataFrameColumnParser.create(df)
-            columns = column_parser.get_columns()
             table_request = CreateTableRequest(
                 name=EntityName(table_name),
                 tableType=table_type,
-                columns=columns,
+                columns=self.get_table_columns(schema_name, table_name),
                 tableConstraints=self.get_table_constraints(
                     schema_name=schema_name,
                     table_name=table_name,

metadata/ingestion/source/database/common_pg_mappings.py

@@ -0,0 +1,49 @@
+"""
+Common Postgresql mappings
+"""
+
+from sqlalchemy import String as SqlAlchemyString
+from sqlalchemy.dialects.postgresql.base import ischema_names
+
+from metadata.generated.schema.entity.data.table import (
+    PartitionIntervalTypes,
+    TableType,
+)
+from metadata.ingestion.source.database.column_type_parser import create_sqlalchemy_type
+
+INTERVAL_TYPE_MAP = {
+    "list": PartitionIntervalTypes.COLUMN_VALUE,
+    "hash": PartitionIntervalTypes.COLUMN_VALUE,
+    "range": PartitionIntervalTypes.TIME_UNIT,
+}
+
+RELKIND_MAP = {
+    "r": TableType.Regular,
+    "p": TableType.Partitioned,
+    "f": TableType.Foreign,
+    "v": TableType.View,
+}
+
+GEOMETRY = create_sqlalchemy_type("GEOMETRY")
+POINT = create_sqlalchemy_type("POINT")
+POLYGON = create_sqlalchemy_type("POLYGON")
+
+ischema_names.update(
+    {
+        "geometry": GEOMETRY,
+        "point": POINT,
+        "polygon": POLYGON,
+        "box": create_sqlalchemy_type("BOX"),
+        "bpchar": SqlAlchemyString,
+        "circle": create_sqlalchemy_type("CIRCLE"),
+        "line": create_sqlalchemy_type("LINE"),
+        "lseg": create_sqlalchemy_type("LSEG"),
+        "path": create_sqlalchemy_type("PATH"),
+        "pg_lsn": create_sqlalchemy_type("PG_LSN"),
+        "pg_snapshot": create_sqlalchemy_type("PG_SNAPSHOT"),
+        "tsquery": create_sqlalchemy_type("TSQUERY"),
+        "txid_snapshot": create_sqlalchemy_type("TXID_SNAPSHOT"),
+        "xid": SqlAlchemyString,
+        "xml": create_sqlalchemy_type("XML"),
+    }
+)

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

@@ -26,6 +26,7 @@ from metadata.ingestion.ometa.ometa_api import OpenMetadata
 from metadata.ingestion.source.database.common_nosql_source import (
     SAMPLE_SIZE,
     CommonNoSQLSource,
+    TableNameAndType,
 )
 from metadata.ingestion.source.database.couchbase.models import IndexObject as Index
 from metadata.ingestion.source.database.couchbase.queries import (
@@ -96,13 +97,18 @@ class CouchbaseSource(CommonNoSQLSource):
             logger.debug(traceback.format_exc())
         return []
 
-    def get_table_name_list(self, schema_name: str) -> List[str]:
+    def query_table_names_and_types(
+        self, schema_name: str
+    ) -> Iterable[TableNameAndType]:
         """
         Method to get list of table names available within schema db
         """
         try:
             scope_object = self.context.get().scope_dict.get(schema_name)
-            return [collection.name for collection in scope_object.collections]
+            return [
+                TableNameAndType(name=collection.name)
+                for collection in scope_object.collections
+            ]
         except Exception as exp:
             logger.debug(
                 f"Failed to list collection names for scope [{schema_name}]: {exp}"