openmetadata-ingestion 1.3.0.0.dev0__py3-none-any.whl → 1.3.1.0__py3-none-any.whl

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

@@ -58,6 +58,28 @@ BIGQUERY_TABLE_AND_TYPE = textwrap.dedent(
     """
 )
 
+BIGQUERY_TABLE_CONSTRAINTS = textwrap.dedent(
+    """
+    SELECT * 
+    FROM `{project_id}`.{schema_name}.INFORMATION_SCHEMA.CONSTRAINT_COLUMN_USAGE 
+    WHERE table_name = '{table_name}' AND constraint_name LIKE '%pk$';
+    """
+)
+
+BIGQUERY_FOREIGN_CONSTRAINTS = textwrap.dedent(
+    """
+    SELECT 
+      c.table_name AS referred_table, 
+      r.table_schema as referred_schema, 
+      r.constraint_name as name, 
+      c.column_name as referred_columns,
+      c.column_name as constrained_columns 
+    FROM `{project_id}`.{schema_name}.INFORMATION_SCHEMA.CONSTRAINT_COLUMN_USAGE c 
+    JOIN `{project_id}`.{schema_name}.INFORMATION_SCHEMA.TABLE_CONSTRAINTS r ON c.constraint_name = r.constraint_name 
+    WHERE r.constraint_type = 'FOREIGN KEY' AND r.table_name='{table_name}';
+    """
+)
+
 BIGQUERY_GET_STORED_PROCEDURES = textwrap.dedent(
     """
 SELECT 

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

@@ -0,0 +1,62 @@
+#  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.
+"""A client for Google Cloud Bigtable that supports multiple projects."""
+from functools import partial
+from typing import List, Optional, Type
+
+from google import auth
+from google.cloud.bigtable import Client
+
+NoProject = object()
+
+
+class MultiProjectClient:
+    """Google Cloud Client does not support ad-hoc project switching. This class wraps the client and allows
+    switching between projects. If no project is specified, the client will not have a project set and will try
+    to resolve it from ADC.
+    Example usage:
+    ```
+    from google.cloud.bigtable import Client
+    client = MultiProjectClient(Client, project_ids=["project1", "project2"])
+    instances_project1 = client.list_instances("project1")
+    instances_project2 = client.list_instances("project2")
+    """
+
+    def __init__(
+        self,
+        client_class: Type[Client],
+        project_ids: Optional[List[str]] = None,
+        **client_kwargs,
+    ):
+        if project_ids:
+            self.clients = {
+                project_id: client_class(project=project_id, **client_kwargs)
+                for project_id in project_ids
+            }
+        else:
+            self.clients = {NoProject: client_class(**client_kwargs)}
+
+    def project_ids(self):
+        if NoProject in self.clients:
+            _, project_id = auth.default()
+            return [project_id]
+        return list(self.clients.keys())
+
+    def __getattr__(self, client_method):
+        """Return the underlying client method as a partial function so we can inject the project_id."""
+        return partial(self._call, client_method)
+
+    def _call(self, method, project_id, *args, **kwargs):
+        """Call the method on the client for the given project_id. The args and kwargs are passed through."""
+        client = self.clients.get(project_id, self.clients.get(NoProject))
+        if not client:
+            raise ValueError(f"Project {project_id} not found")
+        return getattr(client, method)(*args, **kwargs)

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

@@ -0,0 +1,116 @@
+#  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 connection"""
+from typing import List, Optional
+
+from google.cloud.bigtable import Client
+
+from metadata.generated.schema.entity.automations.workflow import (
+    Workflow as AutomationWorkflow,
+)
+from metadata.generated.schema.entity.services.connections.database.bigTableConnection import (
+    BigTableConnection,
+)
+from metadata.generated.schema.security.credentials.gcpValues import (
+    GcpCredentialsValues,
+    SingleProjectId,
+)
+from metadata.ingestion.connections.test_connections import (
+    SourceConnectionException,
+    test_connection_steps,
+)
+from metadata.ingestion.ometa.ometa_api import OpenMetadata
+from metadata.ingestion.source.database.bigtable.client import MultiProjectClient
+from metadata.utils.credentials import set_google_credentials
+from metadata.utils.logger import ingestion_logger
+
+logger = ingestion_logger()
+
+
+def get_connection(connection: BigTableConnection):
+    set_google_credentials(connection.credentials)
+    project_ids = None
+    if isinstance(connection.credentials.gcpConfig, GcpCredentialsValues):
+        project_ids = (
+            [connection.credentials.gcpConfig.projectId.__root__]
+            if isinstance(connection.credentials.gcpConfig.projectId, SingleProjectId)
+            else connection.credentials.gcpConfig.projectId.__root__
+        )
+    # admin=True is required to list instances and tables
+    return MultiProjectClient(client_class=Client, project_ids=project_ids, admin=True)
+
+
+def get_nested_index(lst: list, index: List[int], default=None):
+    try:
+        for i in index:
+            lst = lst[i]
+        return lst
+    except IndexError:
+        return default
+
+
+class Tester:
+    """
+    A wrapper class that holds state. We need it because the different testing stages
+    are not independent of each other. For example, we need to list instances before we can list
+    """
+
+    def __init__(self, client: MultiProjectClient):
+        self.client = client
+        self.project_id = None
+        self.instance = None
+        self.table = None
+
+    def list_instances(self):
+        self.project_id = list(self.client.clients.keys())[0]
+        instances = list(self.client.list_instances(project_id=self.project_id))
+        self.instance = get_nested_index(instances, [0, 0])
+
+    def list_tables(self):
+        if not self.instance:
+            raise SourceConnectionException(
+                f"No instances found in project {self.project_id}"
+            )
+        tables = list(self.instance.list_tables())
+        self.table = tables[0]
+
+    def get_row(self):
+        if not self.table:
+            raise SourceConnectionException(
+                f"No tables found in project {self.instance.project_id} and instance {self.instance.instance_id}"
+            )
+        self.table.read_rows(limit=1)
+
+
+def test_connection(
+    metadata: OpenMetadata,
+    client: MultiProjectClient,
+    service_connection: BigTableConnection,
+    automation_workflow: Optional[AutomationWorkflow] = None,
+) -> None:
+    """
+    Test connection. This can be executed either as part
+    of a metadata workflow or during an Automation Workflow
+    """
+    tester = Tester(client)
+
+    test_fn = {
+        "GetInstances": tester.list_instances,
+        "GetTables": tester.list_tables,
+        "GetRows": tester.get_row,
+    }
+
+    test_connection_steps(
+        metadata=metadata,
+        test_fn=test_fn,
+        service_type=service_connection.type.value,
+        automation_workflow=automation_workflow,
+    )

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,
 )
@@ -46,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
 
@@ -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]]:
@@ -217,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,

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

@@ -94,7 +94,6 @@ class DatabricksClient:
         Method returns List the history of queries through SQL warehouses
         """
         try:
-
             data = {}
             daydiff = end_date - start_date