fixes #12591: add BigTable (#15122)

* feat(connector): add BigTable

* bigtable work

1. docstrings
2. tests
3. created a Row BaseModel
4. implemented a ClassConverter

* docs moved to separate PR

* format files

* minor cosmetic

- removed TODO
- changed headers' year to 2024 for new files
- fixed typos

* format

* formatting and comments

1. added missing docstrings.
2. abstracted the _find_instance method.
3. aliased the IDs used in the BigTable connection

* added comment regarding private key

* added comments regarding column families

* enclose get_schema_name_list in `try/except/else`

* format

* streamlined get_schema_name_list to include all logic in the try block

ingestion/setup.py

@@ -143,6 +143,7 @@
         VERSIONS["pyarrow"],
         "sqlalchemy-bigquery>=1.2.2",
     },
+    "bigtable": {"google-cloud-bigtable>=2.0.0", VERSIONS["pandas"]},
     "clickhouse": {"clickhouse-driver~=0.2", "clickhouse-sqlalchemy~=0.2"},
     "dagster": {
         VERSIONS["pymysql"],

ingestion/src/metadata/examples/workflows/bigtable.yaml

@@ -0,0 +1,32 @@
+source:
+  type: bigtable
+  serviceName: local_bigtable
+  serviceConnection:
+    config:
+      type: BigTable
+      credentials:
+        gcpConfig:
+          type: service_account
+          projectId: project_id
+          privateKeyId: private_key_id
+          privateKey: private_key
+          clientEmail: gcpuser@project_id.iam.gserviceaccount.com
+          clientId: client_id
+          authUri: https://accounts.google.com/o/oauth2/auth
+          tokenUri: https://oauth2.googleapis.com/token
+          authProviderX509CertUrl: https://www.googleapis.com/oauth2/v1/certs
+          clientX509CertUrl: https://www.googleapis.com/oauth2/v1/certs
+
+  sourceConfig:
+    config:
+      type: DatabaseMetadata
+sink:
+  type: metadata-rest
+  config: {}
+workflowConfig:
+  loggerLevel: DEBUG
+  openMetadataServerConfig:
+    hostPort: http://localhost:8585/api
+    authProvider: openmetadata
+    securityConfig:
+      jwtToken: "eyJraWQiOiJHYjM4OWEtOWY3Ni1nZGpzLWE5MmotMDI0MmJrOTQzNTYiLCJ0eXAiOiJKV1QiLCJhbGciOiJSUzI1NiJ9.eyJzdWIiOiJhZG1pbiIsImlzQm90IjpmYWxzZSwiaXNzIjoib3Blbi1tZXRhZGF0YS5vcmciLCJpYXQiOjE2NjM5Mzg0NjIsImVtYWlsIjoiYWRtaW5Ab3Blbm1ldGFkYXRhLm9yZyJ9.tS8um_5DKu7HgzGBzS1VTA5uUjKWOCU0B_j08WXBiEC0mr0zNREkqVfwFDD-d24HlNEbrqioLsBuFRiwIWKc1m_ZlVQbG7P36RUxhuv2vbSp80FKyNM-Tj93FDzq91jsyNmsQhyNv_fNr3TXfzzSPjHt8Go0FMMP66weoKMgW2PbXlhVKwEuXUHyakLLzewm9UMeQaEiRzhiTMU3UkLXcKbYEJJvfNFcLwSl9W8JCO_l0Yj3ud-qt_nQYEZwqW6u5nfdQllN133iikV4fM5QZsMCnm8Rq1mvLR0y9bmJiD7fwM1tmJ791TUWqmKaTnP49U493VanKpUAfzIiOiIbhg"

ingestion/src/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)

ingestion/src/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,
+    )