openmetadata-ingestion 1.2.4.0__py3-none-any.whl → 1.3.2.0rc1__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/column_helpers.py

@@ -14,16 +14,6 @@ their raw information from the source
 """
 
 
-def remove_table_from_column_name(table_name: str, raw_column_name: str) -> str:
-    """
-    Given a column `table.column`, return only `column`.
-
-    Note that we might have columns which have real dots
-    "." in the name, so we cannot just split.
-    """
-    return raw_column_name.replace(table_name + ".", "")
-
-
 def truncate_column_name(col_name: str):
     """
     OpenMetadata table column specification limits column name to 128 characters.

metadata/ingestion/source/database/column_type_parser.py

@@ -369,7 +369,9 @@ class ColumnTypeParser:
 
     @staticmethod
     def _parse_struct_fields_string(stuct_type: str) -> Dict[str, object]:
-        parts = ColumnTypeParser._ignore_brackets_split(stuct_type, ",")
+        parts = ColumnTypeParser._ignore_brackets_split(
+            stuct_type, ",", skip_no_child_validation=True
+        )
         columns = []
         for part in parts:
             name_and_type = ColumnTypeParser._ignore_brackets_split(part, ":")
@@ -387,9 +389,10 @@ class ColumnTypeParser:
             field_type = ColumnTypeParser._parse_datatype_string(name_and_type[1])
             field_type["name"] = field_name
             columns.append(field_type)
+
         return {
             "children": columns,
-            "dataTypeDisplay": f"struct<{stuct_type}>",
+            "dataTypeDisplay": f"struct<{stuct_type or 'unknown'}>",
             "dataType": "STRUCT",
         }
 
@@ -431,7 +434,9 @@ class ColumnTypeParser:
         }
 
     @staticmethod
-    def _ignore_brackets_split(string: str, separator: str) -> List[str]:
+    def _ignore_brackets_split(
+        string: str, separator: str, skip_no_child_validation: bool = False
+    ) -> List[str]:
         parts = []
         buf = ""
         level = 0
@@ -452,9 +457,10 @@ class ColumnTypeParser:
             else:
                 buf += char
 
-        if len(buf) == 0:
+        if len(buf) == 0 and not skip_no_child_validation:
             raise ValueError(f"The {separator} cannot be the last char: {string}")
-        parts.append(buf)
+        if buf:
+            parts.append(buf)
         return parts
 
     @staticmethod

metadata/ingestion/source/database/common_db_source.py

@@ -41,13 +41,16 @@ from metadata.generated.schema.entity.data.table import (
     TablePartition,
     TableType,
 )
+from metadata.generated.schema.entity.services.ingestionPipelines.status import (
+    StackTraceError,
+)
 from metadata.generated.schema.metadataIngestion.databaseServiceMetadataPipeline import (
     DatabaseServiceMetadataPipeline,
 )
 from metadata.generated.schema.metadataIngestion.workflow import (
     Source as WorkflowSource,
 )
-from metadata.ingestion.api.models import Either, StackTraceError
+from metadata.ingestion.api.models import Either
 from metadata.ingestion.lineage.sql_lineage import get_column_fqn
 from metadata.ingestion.models.ometa_classification import OMetaTagAndClassification
 from metadata.ingestion.ometa.ometa_api import OpenMetadata
@@ -59,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()
@@ -171,6 +174,7 @@ class CommonDbSourceService(
                 service=self.context.database_service,
                 description=self.get_database_description(database_name),
                 sourceUrl=self.get_source_url(database_name=database_name),
+                tags=self.get_database_tag_labels(database_name=database_name),
             )
         )
 
@@ -209,6 +213,7 @@ class CommonDbSourceService(
                     database_name=self.context.database,
                     schema_name=schema_name,
                 ),
+                tags=self.get_schema_tag_labels(schema_name=schema_name),
             )
         )
 
@@ -246,6 +251,23 @@ class CommonDbSourceService(
             for table_name in self.inspector.get_table_names(schema_name) or []
         ]
 
+    def query_view_names_and_types(
+        self, schema_name: str
+    ) -> Iterable[TableNameAndType]:
+        """
+        Connect to the source database to get the view
+        name and type. By default, use the inspector method
+        to get the names and pass the View type.
+
+        This is useful for sources where we need fine-grained
+        logic on how to handle table types, e.g., material views,...
+        """
+
+        return [
+            TableNameAndType(name=table_name, type_=TableType.View)
+            for table_name in self.inspector.get_view_names(schema_name) or []
+        ]
+
     def get_tables_name_and_type(self) -> Optional[Iterable[Tuple[str, str]]]:
         """
         Handle table and views.
@@ -285,8 +307,10 @@ class CommonDbSourceService(
                     yield table_name, table_and_type.type_
 
             if self.source_config.includeViews:
-                for view_name in self.inspector.get_view_names(schema_name):
-                    view_name = self.standardize_table_name(schema_name, view_name)
+                for view_and_type in self.query_view_names_and_types(schema_name):
+                    view_name = self.standardize_table_name(
+                        schema_name, view_and_type.name
+                    )
                     view_fqn = fqn.build(
                         self.metadata,
                         entity_type=Table,
@@ -307,7 +331,7 @@ class CommonDbSourceService(
                             "Table Filtered Out",
                         )
                         continue
-                    yield view_name, TableType.View
+                    yield view_name, view_and_type.type_
         except Exception as err:
             logger.warning(
                 f"Fetching tables names failed for schema {schema_name} due to - {err}"
@@ -317,7 +341,7 @@ class CommonDbSourceService(
     def get_view_definition(
         self, table_type: str, table_name: str, schema_name: str, inspector: Inspector
     ) -> Optional[str]:
-        if table_type == TableType.View:
+        if table_type in (TableType.View, TableType.MaterializedView):
             try:
                 view_definition = inspector.get_view_definition(table_name, schema_name)
                 view_definition = (
@@ -381,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]]:
@@ -439,6 +463,7 @@ class CommonDbSourceService(
                     database_name=self.context.database,
                     table_type=table_type,
                 ),
+                owner=self.get_owner_ref(table_name=table_name),
             )
 
             is_partitioned, partition_details = self.get_table_partition_details(
@@ -469,7 +494,7 @@ class CommonDbSourceService(
             error = f"Unexpected exception to yield table [{table_name}]: {exc}"
             yield Either(
                 left=StackTraceError(
-                    name=table_name, error=error, stack_trace=traceback.format_exc()
+                    name=table_name, error=error, stackTrace=traceback.format_exc()
                 )
             )
 

metadata/ingestion/source/database/common_nosql_source.py

@@ -28,14 +28,21 @@ 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,
+)
 from metadata.generated.schema.metadataIngestion.databaseServiceMetadataPipeline import (
     DatabaseServiceMetadataPipeline,
 )
 from metadata.generated.schema.metadataIngestion.workflow import (
     Source as WorkflowSource,
 )
-from metadata.ingestion.api.models import Either, StackTraceError
+from metadata.ingestion.api.models import Either
 from metadata.ingestion.models.ometa_classification import OMetaTagAndClassification
 from metadata.ingestion.ometa.ometa_api import OpenMetadata
 from metadata.ingestion.source.connections import get_connection
@@ -43,7 +50,7 @@ from metadata.ingestion.source.database.database_service import DatabaseServiceS
 from metadata.ingestion.source.database.stored_procedures_mixin import QueryByProcedure
 from metadata.utils import fqn
 from metadata.utils.constants import DEFAULT_DATABASE
-from metadata.utils.datalake.datalake_utils import get_columns
+from metadata.utils.datalake.datalake_utils import DataFrameColumnParser
 from metadata.utils.filters import filter_by_schema, filter_by_table
 from metadata.utils.logger import ingestion_logger
 
@@ -200,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]]:
@@ -214,12 +230,17 @@ class CommonNoSQLSource(DatabaseServiceSource, ABC):
         try:
             data = self.get_table_columns_dict(schema_name, table_name)
             df = pd.DataFrame.from_records(list(data))
-            columns = get_columns(df)
+            column_parser = DataFrameColumnParser.create(df)
+            columns = column_parser.get_columns()
             table_request = CreateTableRequest(
                 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,
@@ -242,7 +263,7 @@ class CommonNoSQLSource(DatabaseServiceSource, ABC):
                 left=StackTraceError(
                     name=table_name,
                     error=f"Unexpected exception to yield table [{table_name}]: {exc}",
-                    stack_trace=traceback.format_exc(),
+                    stackTrace=traceback.format_exc(),
                 )
             )