openmetadata-ingestion 1.3.0.1__py3-none-any.whl → 1.3.1.0__py3-none-any.whl

metadata/ingestion/source/database/bigtable/metadata.py

@@ -0,0 +1,224 @@
+#  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.
+"""
+Bigtable source methods.
+"""
+import traceback
+from typing import Dict, Iterable, List, Optional, Union
+
+from google.cloud.bigtable import row_filters
+from google.cloud.bigtable.instance import Instance
+from google.cloud.bigtable.table import Table
+
+from metadata.generated.schema.entity.data.table import (
+    ConstraintType,
+    TableConstraint,
+    TableType,
+)
+from metadata.generated.schema.entity.services.connections.database.bigTableConnection import (
+    BigTableConnection,
+)
+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.bigtable.client import MultiProjectClient
+from metadata.ingestion.source.database.bigtable.models import Row
+from metadata.ingestion.source.database.common_nosql_source import (
+    SAMPLE_SIZE as GLOBAL_SAMPLE_SIZE,
+)
+from metadata.ingestion.source.database.common_nosql_source import CommonNoSQLSource
+from metadata.ingestion.source.database.multi_db_source import MultiDBSource
+from metadata.utils.logger import ingestion_logger
+
+logger = ingestion_logger()
+
+# BigTable group's its columns in column families. We make an assumption that if the table has a big number of
+# columns, we at least get a sample of the first 100 column families.
+MAX_COLUMN_FAMILIES = 100
+SAMPLES_PER_COLUMN_FAMILY = 100
+
+ProjectId = str
+InstanceId = str
+TableId = str
+
+
+class BigtableSource(CommonNoSQLSource, MultiDBSource):
+    """
+    Implements the necessary methods to extract database metadata from Google BigTable Source.
+    BigTable is a NoSQL database service for handling large amounts of data. Tha mapping is as follows:
+      project -> instance -> table -> column_family.column
+      (database) (schema)
+    For more infor about BigTable: https://cloud.google.com/bigtable/?hl=en
+    All data types are registered as bytes.
+    """
+
+    def __init__(self, config: WorkflowSource, metadata: OpenMetadata):
+        super().__init__(config, metadata)
+        self.client: MultiProjectClient = self.connection_obj
+
+        # ths instances and tables are cached to avoid making redundant requests to the API.
+        self.instances: Dict[ProjectId, Dict[InstanceId, Instance]] = {}
+        self.tables: Dict[ProjectId, Dict[InstanceId, Dict[TableId, Table]]] = {}
+
+    @classmethod
+    def create(cls, config_dict, metadata: OpenMetadata):
+        config: WorkflowSource = WorkflowSource.parse_obj(config_dict)
+        connection: BigTableConnection = config.serviceConnection.__root__.config
+        if not isinstance(connection, BigTableConnection):
+            raise InvalidSourceException(
+                f"Expected BigTableConnection, but got {connection}"
+            )
+        return cls(config, metadata)
+
+    def get_configured_database(self) -> Optional[str]:
+        """
+        This connector uses "virtual databases" in the form of GCP projects.
+        The concept of a default project for the GCP client is not useful here because the project ID
+        is always an explicit part of the connection. Therefore, this method returns None and the databases
+        are resolved using `self.get_database_names`.
+        """
+        return None
+
+    def get_database_names(self) -> Iterable[str]:
+        return self.get_database_names_raw()
+
+    def get_database_names_raw(self) -> Iterable[str]:
+        yield from self.client.project_ids()
+
+    def get_schema_name_list(self) -> List[str]:
+        project_id = self.context.database
+        try:
+            # the first element is a list of instances
+            # the second element is another collection (seems empty) and I do not know what is its purpose
+            instances, _ = self.client.list_instances(project_id=project_id)
+            self.instances[project_id] = {
+                instance.instance_id: instance for instance in instances
+            }
+            return list(self.instances[project_id].keys())
+        except Exception as err:
+            logger.debug(traceback.format_exc())
+            logger.error(
+                f"Failed to list BigTable instances in project {project_id}: {err}"
+            )
+            raise
+
+    def get_table_name_list(self, schema_name: str) -> List[str]:
+        project_id = self.context.database
+        try:
+            instance = self._get_instance(project_id, schema_name)
+            if instance is None:
+                raise RuntimeError(f"Instance {project_id}/{schema_name} not found.")
+            tables = instance.list_tables()
+            for table in tables:
+                self._set_nested(
+                    self.tables,
+                    [project_id, instance.instance_id, table.table_id],
+                    table,
+                )
+            return list(self.tables[project_id][schema_name].keys())
+        except Exception as err:
+            logger.debug(traceback.format_exc())
+            # add context to the error message
+            logger.error(
+                f"Failed to list BigTable table names in {project_id}.{schema_name}: {err}"
+            )
+        return []
+
+    def get_table_constraints(
+        self, db_name: str, schema_name: str, table_name: str
+    ) -> List[TableConstraint]:
+        return [
+            TableConstraint(
+                constraintType=ConstraintType.PRIMARY_KEY, columns=["row_key"]
+            )
+        ]
+
+    def get_table_columns_dict(
+        self, schema_name: str, table_name: str
+    ) -> Union[List[Dict], Dict]:
+        project_id = self.context.database
+        try:
+            table = self._get_table(project_id, schema_name, table_name)
+            if table is None:
+                raise RuntimeError(
+                    f"Table {project_id}/{schema_name}/{table_name} not found."
+                )
+            column_families = table.list_column_families()
+            # all BigTable tables have a "row_key" column. Even if there are no records in the table.
+            records = [{"row_key": b"row_key"}]
+            # In order to get a "good" sample of data, we try to distribute the sampling
+            # across multiple column families.
+            for column_family in list(column_families.keys())[:MAX_COLUMN_FAMILIES]:
+                records.extend(
+                    self._get_records_for_column_family(
+                        table, column_family, SAMPLES_PER_COLUMN_FAMILY
+                    )
+                )
+                if len(records) >= GLOBAL_SAMPLE_SIZE:
+                    break
+            return records
+        except Exception as err:
+            logger.debug(traceback.format_exc())
+            logger.warning(
+                f"Failed to read BigTable rows for [{project_id}.{schema_name}.{table_name}]: {err}"
+            )
+        return []
+
+    def get_source_url(
+        self,
+        database_name: Optional[str] = None,
+        schema_name: Optional[str] = None,
+        table_name: Optional[str] = None,
+        table_type: Optional[TableType] = None,
+    ) -> Optional[str]:
+        """
+        Method to get the source url for a BigTable table
+        """
+        try:
+            if schema_name and table_name:
+                return (
+                    "https://console.cloud.google.com/bigtable/instances/"
+                    f"{schema_name}/tables/{table_name}/overview?project={database_name}"
+                )
+        except Exception as exc:
+            logger.debug(traceback.format_exc())
+            logger.error(f"Unable to get source url: {exc}")
+        return None
+
+    @staticmethod
+    def _set_nested(dct: dict, keys: List[str], value: any) -> None:
+        for key in keys[:-1]:
+            dct = dct.setdefault(key, {})
+        dct[keys[-1]] = value
+
+    @staticmethod
+    def _get_records_for_column_family(
+        table: Table, column_family: str, limit: int
+    ) -> List[Dict]:
+        filter_ = row_filters.ColumnRangeFilter(column_family_id=column_family)
+        rows = table.read_rows(limit=limit, filter_=filter_)
+        return [Row.from_partial_row(row).to_record() for row in rows]
+
+    def _get_table(
+        self, project_id: str, schema_name: str, table_name: str
+    ) -> Optional[Table]:
+        try:
+            return self.tables[project_id][schema_name][table_name]
+        except KeyError:
+            return None
+
+    def _get_instance(self, project_id: str, schema_name: str) -> Optional[Instance]:
+        try:
+            return self.instances[project_id][schema_name]
+        except KeyError:
+            return None

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

@@ -0,0 +1,60 @@
+#  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.
+"""
+Bigtable source models.
+"""
+from typing import Dict, List
+
+from google.cloud.bigtable.row import PartialRowData
+from pydantic import BaseModel
+
+
+class Value(BaseModel):
+    """A Bigtable cell value."""
+
+    timestamp: int
+    value: bytes
+
+
+class Cell(BaseModel):
+    """A Bigtable cell."""
+
+    values: List[Value]
+
+
+class Row(BaseModel):
+    """A Bigtable row."""
+
+    cells: Dict[str, Dict[bytes, Cell]]
+    row_key: bytes
+
+    @classmethod
+    def from_partial_row(cls, row: PartialRowData):
+        cells = {}
+        for column_family, cf_cells in row.cells.items():
+            cells.setdefault(column_family, {})
+            for column, cell in cf_cells.items():
+                cells[column_family][column] = Cell(
+                    values=[Value(timestamp=c.timestamp, value=c.value) for c in cell]
+                )
+        return cls(cells=cells, row_key=row.row_key)
+
+    def to_record(self) -> Dict[str, bytes]:
+        record = {}
+        for column_family, cells in self.cells.items():
+            for column, cell in cells.items():
+                # Since each cell can have multiple values and the API returns them in descending order
+                # from latest to oldest, we only take the latest value. This probably does not matter since
+                # all we care about is data types and all data stored in BigTable is of type `bytes`.
+                record[f"{column_family}.{column.decode()}"] = cell.values[0].value
+        record["row_key"] = self.row_key
+
+        return record

metadata/ingestion/source/database/common_db_source.py

@@ -62,8 +62,8 @@ from metadata.ingestion.source.database.stored_procedures_mixin import QueryByPr
 from metadata.ingestion.source.models import TableView
 from metadata.utils import fqn
 from metadata.utils.db_utils import get_view_lineage
+from metadata.utils.execution_time_tracker import calculate_execution_time_generator
 from metadata.utils.filters import filter_by_table
-from metadata.utils.helpers import calculate_execution_time_generator
 from metadata.utils.logger import ingestion_logger
 
 logger = ingestion_logger()
@@ -405,7 +405,7 @@ class CommonDbSourceService(
         """Not Implemented"""
         yield from []
 
-    @calculate_execution_time_generator
+    @calculate_execution_time_generator(store=False)
     def yield_table(
         self, table_name_and_type: Tuple[str, str]
     ) -> Iterable[Either[CreateTableRequest]]:

metadata/ingestion/source/database/common_nosql_source.py

@@ -28,7 +28,11 @@ 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 Table, TableType
+from metadata.generated.schema.entity.data.table import (
+    Table,
+    TableConstraint,
+    TableType,
+)
 from metadata.generated.schema.entity.services.ingestionPipelines.status import (
     StackTraceError,
 )
@@ -203,6 +207,15 @@ class CommonNoSQLSource(DatabaseServiceSource, ABC):
         need to be overridden by sources
         """
 
+    def get_table_constraints(
+        self,
+        db_name: str,
+        schema_name: str,
+        table_name: str,
+    ) -> Optional[List[TableConstraint]]:
+        # pylint: disable=unused-argument
+        return None
+
     def yield_table(
         self, table_name_and_type: Tuple[str, str]
     ) -> Iterable[Either[CreateTableRequest]]:
@@ -223,7 +236,11 @@ class CommonNoSQLSource(DatabaseServiceSource, ABC):
                 name=table_name,
                 tableType=table_type,
                 columns=columns,
-                tableConstraints=None,
+                tableConstraints=self.get_table_constraints(
+                    schema_name=schema_name,
+                    table_name=table_name,
+                    db_name=self.context.database,
+                ),
                 databaseSchema=fqn.build(
                     metadata=self.metadata,
                     entity_type=DatabaseSchema,

metadata/ingestion/source/database/databricks/metadata.py

@@ -13,7 +13,7 @@
 import re
 import traceback
 from copy import deepcopy
-from typing import Iterable, Optional, Tuple
+from typing import Iterable, Optional, Tuple, Union
 
 from pyhive.sqlalchemy_hive import _type_map
 from sqlalchemy import types, util
@@ -66,6 +66,7 @@ logger = ingestion_logger()
 
 DATABRICKS_TAG = "DATABRICK TAG"
 DATABRICKS_TAG_CLASSIFICATION = "DATABRICK TAG CLASSIFICATION"
+DEFAULT_TAG_VALUE = "NONE"
 
 
 class STRUCT(String):
@@ -268,6 +269,10 @@ class DatabricksSource(CommonDbSourceService, MultiDBSource):
         super().__init__(config, metadata)
         self.is_older_version = False
         self._init_version()
+        self.catalog_tags = {}
+        self.schema_tags = {}
+        self.table_tags = {}
+        self.column_tags = {}
 
     def _init_version(self):
         try:
@@ -314,10 +319,103 @@ class DatabricksSource(CommonDbSourceService, MultiDBSource):
         else:
             yield DEFAULT_DATABASE
 
+    def _clear_tag_cache(self) -> None:
+        """
+        Method to clean any existing tags available in memory
+        """
+        self.catalog_tags.clear()
+        self.table_tags.clear()
+        self.schema_tags.clear()
+        self.column_tags.clear()
+
+    def _add_to_tag_cache(
+        self, tag_dict: dict, key: Union[str, Tuple], value: Tuple[str, str]
+    ):
+        if tag_dict.get(key):
+            tag_dict.get(key).append(value)
+        else:
+            tag_dict[key] = [value]
+
+    def populate_tags_cache(self, database_name: str) -> None:
+        """
+        Method to fetch all the tags and populate the relevant caches
+        """
+        self._clear_tag_cache()
+        if self.source_config.includeTags is False:
+            return
+        try:
+            tags = self.connection.execute(
+                DATABRICKS_GET_CATALOGS_TAGS.format(database_name=database_name)
+            )
+
+            for tag in tags:
+                self._add_to_tag_cache(
+                    self.catalog_tags,
+                    tag.catalog_name,
+                    # tag value is an optional field, if tag value is not available use default tag value
+                    (tag.tag_name, tag.tag_value or DEFAULT_TAG_VALUE),
+                )
+        except Exception as exc:
+            logger.debug(f"Failed to fetch catalog tags due to - {exc}")
+
+        try:
+            tags = self.connection.execute(
+                DATABRICKS_GET_SCHEMA_TAGS.format(database_name=database_name)
+            )
+            for tag in tags:
+                self._add_to_tag_cache(
+                    self.schema_tags,
+                    (tag.catalog_name, tag.schema_name),
+                    # tag value is an optional field, if tag value is not available use default tag value
+                    (tag.tag_name, tag.tag_value or DEFAULT_TAG_VALUE),
+                )
+        except Exception as exc:
+            logger.debug(f"Failed to fetch schema tags due to - {exc}")
+
+        try:
+            tags = self.connection.execute(
+                DATABRICKS_GET_TABLE_TAGS.format(database_name=database_name)
+            )
+            for tag in tags:
+                self._add_to_tag_cache(
+                    self.table_tags,
+                    (tag.catalog_name, tag.schema_name, tag.table_name),
+                    # tag value is an optional field, if tag value is not available use default tag value
+                    (tag.tag_name, tag.tag_value or DEFAULT_TAG_VALUE),
+                )
+        except Exception as exc:
+            logger.debug(f"Failed to fetch table tags due to - {exc}")
+
+        try:
+            tags = self.connection.execute(
+                DATABRICKS_GET_COLUMN_TAGS.format(database_name=database_name)
+            )
+            for tag in tags:
+                tag_table_id = (tag.catalog_name, tag.schema_name, tag.table_name)
+                if self.column_tags.get(tag_table_id):
+                    self._add_to_tag_cache(
+                        self.column_tags.get(tag_table_id),
+                        tag.column_name,
+                        # tag value is an optional field, if tag value is not available use default tag value
+                        (tag.tag_name, tag.tag_value or DEFAULT_TAG_VALUE),
+                    )
+                else:
+                    self.column_tags[tag_table_id] = {
+                        tag.column_name: [
+                            (
+                                tag.tag_name,
+                                tag.tag_value or DEFAULT_TAG_VALUE,
+                            )
+                        ]
+                    }
+        except Exception as exc:
+            logger.debug(f"Failed to fetch column tags due to - {exc}")
+
     def get_database_names(self) -> Iterable[str]:
         configured_catalog = self.service_connection.catalog
         if configured_catalog:
             self.set_inspector(database_name=configured_catalog)
+            self.populate_tags_cache(database_name=configured_catalog)
             yield configured_catalog
         else:
             for new_catalog in self.get_database_names_raw():
@@ -337,6 +435,7 @@ class DatabricksSource(CommonDbSourceService, MultiDBSource):
                     continue
                 try:
                     self.set_inspector(database_name=new_catalog)
+                    self.populate_tags_cache(database_name=new_catalog)
                     yield new_catalog
                 except Exception as exc:
                     logger.error(traceback.format_exc())
@@ -361,10 +460,8 @@ class DatabricksSource(CommonDbSourceService, MultiDBSource):
         Method to yield database tags
         """
         try:
-            tags = self.connection.execute(
-                DATABRICKS_GET_CATALOGS_TAGS.format(database_name=database_name)
-            )
-            for tag in tags:
+            catalog_tags = self.catalog_tags.get(database_name, [])
+            for tag_name, tag_value in catalog_tags:
                 yield from get_ometa_tag_and_classification(
                     tag_fqn=fqn.build(
                         self.metadata,
@@ -372,8 +469,8 @@ class DatabricksSource(CommonDbSourceService, MultiDBSource):
                         service_name=self.context.database_service,
                         database_name=database_name,
                     ),
-                    tags=[tag.tag_value],
-                    classification_name=tag.tag_name,
+                    tags=[tag_value],
+                    classification_name=tag_name,
                     tag_description=DATABRICKS_TAG,
                     classification_description=DATABRICKS_TAG_CLASSIFICATION,
                 )
@@ -394,12 +491,8 @@ class DatabricksSource(CommonDbSourceService, MultiDBSource):
         Method to yield schema tags
         """
         try:
-            tags = self.connection.execute(
-                DATABRICKS_GET_SCHEMA_TAGS.format(
-                    database_name=self.context.database, schema_name=schema_name
-                )
-            )
-            for tag in tags:
+            schema_tags = self.schema_tags.get((self.context.database, schema_name), [])
+            for tag_name, tag_value in schema_tags:
                 yield from get_ometa_tag_and_classification(
                     tag_fqn=fqn.build(
                         self.metadata,
@@ -408,8 +501,8 @@ class DatabricksSource(CommonDbSourceService, MultiDBSource):
                         database_name=self.context.database,
                         schema_name=schema_name,
                     ),
-                    tags=[tag.tag_value],
-                    classification_name=tag.tag_name,
+                    tags=[tag_value],
+                    classification_name=tag_name,
                     tag_description=DATABRICKS_TAG,
                     classification_description=DATABRICKS_TAG_CLASSIFICATION,
                 )
@@ -428,14 +521,10 @@ class DatabricksSource(CommonDbSourceService, MultiDBSource):
     ) -> Iterable[Either[OMetaTagAndClassification]]:
         table_name, _ = table_name_and_type
         try:
-            table_tags = self.connection.execute(
-                DATABRICKS_GET_TABLE_TAGS.format(
-                    database_name=self.context.database,
-                    schema_name=self.context.database_schema,
-                    table_name=table_name,
-                )
+            table_tags = self.table_tags.get(
+                (self.context.database, self.context.database_schema, table_name), []
             )
-            for tag in table_tags:
+            for tag_name, tag_value in table_tags:
                 yield from get_ometa_tag_and_classification(
                     tag_fqn=fqn.build(
                         self.metadata,
@@ -445,35 +534,32 @@ class DatabricksSource(CommonDbSourceService, MultiDBSource):
                         schema_name=self.context.database_schema,
                         table_name=table_name,
                     ),
-                    tags=[tag.tag_value],
-                    classification_name=tag.tag_name,
+                    tags=[tag_value],
+                    classification_name=tag_name,
                     tag_description=DATABRICKS_TAG,
                     classification_description=DATABRICKS_TAG_CLASSIFICATION,
                 )
 
-            column_tags = self.connection.execute(
-                DATABRICKS_GET_COLUMN_TAGS.format(
-                    database_name=self.context.database,
-                    schema_name=self.context.database_schema,
-                    table_name=table_name,
-                )
+            column_tags = self.column_tags.get(
+                (self.context.database, self.context.database_schema, table_name), {}
             )
-            for tag in column_tags:
-                yield from get_ometa_tag_and_classification(
-                    tag_fqn=fqn.build(
-                        self.metadata,
-                        Column,
-                        service_name=self.context.database_service,
-                        database_name=self.context.database,
-                        schema_name=self.context.database_schema,
-                        table_name=table_name,
-                        column_name=tag.column_name,
-                    ),
-                    tags=[tag.tag_value],
-                    classification_name=tag.tag_name,
-                    tag_description=DATABRICKS_TAG,
-                    classification_description=DATABRICKS_TAG_CLASSIFICATION,
-                )
+            for column_name, tags in column_tags.items():
+                for tag_name, tag_value in tags or []:
+                    yield from get_ometa_tag_and_classification(
+                        tag_fqn=fqn.build(
+                            self.metadata,
+                            Column,
+                            service_name=self.context.database_service,
+                            database_name=self.context.database,
+                            schema_name=self.context.database_schema,
+                            table_name=table_name,
+                            column_name=column_name,
+                        ),
+                        tags=[tag_value],
+                        classification_name=tag_name,
+                        tag_description=DATABRICKS_TAG,
+                        classification_description=DATABRICKS_TAG_CLASSIFICATION,
+                    )
 
         except Exception as exc:
             yield Either(

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

@@ -37,8 +37,7 @@ DATABRICKS_GET_SCHEMA_TAGS = textwrap.dedent(
     """
     SELECT 
         * 
-    FROM {database_name}.information_schema.schema_tags 
-    WHERE schema_name='{schema_name}'"""
+    FROM {database_name}.information_schema.schema_tags"""
 )
 
 DATABRICKS_GET_TABLE_TAGS = textwrap.dedent(
@@ -46,7 +45,7 @@ DATABRICKS_GET_TABLE_TAGS = textwrap.dedent(
     SELECT 
         * 
     FROM {database_name}.information_schema.table_tags 
-    WHERE schema_name = '{schema_name}' AND table_name = '{table_name}';"""
+    """
 )
 
 DATABRICKS_GET_COLUMN_TAGS = textwrap.dedent(
@@ -54,5 +53,5 @@ DATABRICKS_GET_COLUMN_TAGS = textwrap.dedent(
     SELECT 
         * 
     FROM {database_name}.information_schema.column_tags 
-    WHERE schema_name='{schema_name}' AND table_name='{table_name}';"""
+    """
 )