openmetadata-ingestion 1.2.4.0__py3-none-any.whl → 1.3.2.0rc1__py3-none-any.whl

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

@@ -29,7 +29,7 @@ from metadata.generated.schema.api.data.createDatabaseSchema import (
 from metadata.generated.schema.api.data.createStoredProcedure import (
     CreateStoredProcedureRequest,
 )
-from metadata.generated.schema.entity.data.database import Database, EntityName
+from metadata.generated.schema.entity.data.database import Database
 from metadata.generated.schema.entity.data.databaseSchema import DatabaseSchema
 from metadata.generated.schema.entity.data.storedProcedure import StoredProcedureCode
 from metadata.generated.schema.entity.data.table import (
@@ -40,20 +40,27 @@ from metadata.generated.schema.entity.data.table import (
 from metadata.generated.schema.entity.services.connections.database.bigQueryConnection import (
     BigQueryConnection,
 )
+from metadata.generated.schema.entity.services.ingestionPipelines.status import (
+    StackTraceError,
+)
 from metadata.generated.schema.metadataIngestion.workflow import (
     Source as WorkflowSource,
 )
 from metadata.generated.schema.security.credentials.gcpValues import (
     GcpCredentialsValues,
 )
-from metadata.generated.schema.type.basic import SourceUrl
+from metadata.generated.schema.type.basic import EntityName, SourceUrl
 from metadata.generated.schema.type.tagLabel import TagLabel
-from metadata.ingestion.api.models import Either, StackTraceError
+from metadata.ingestion.api.models import Either
 from metadata.ingestion.api.steps import InvalidSourceException
 from metadata.ingestion.models.ometa_classification import OMetaTagAndClassification
 from metadata.ingestion.ometa.ometa_api import OpenMetadata
 from metadata.ingestion.source.connections import get_test_connection_fn
-from metadata.ingestion.source.database.bigquery.helper import get_inspector_details
+from metadata.ingestion.source.database.bigquery.helper import (
+    get_foreign_keys,
+    get_inspector_details,
+    get_pk_constraint,
+)
 from metadata.ingestion.source.database.bigquery.models import (
     STORED_PROC_LANGUAGE_MAP,
     BigQueryStoredProcedure,
@@ -61,6 +68,7 @@ from metadata.ingestion.source.database.bigquery.models import (
 from metadata.ingestion.source.database.bigquery.queries import (
     BIGQUERY_GET_STORED_PROCEDURE_QUERIES,
     BIGQUERY_GET_STORED_PROCEDURES,
+    BIGQUERY_LIFE_CYCLE_QUERY,
     BIGQUERY_SCHEMA_DESCRIPTION,
     BIGQUERY_TABLE_AND_TYPE,
 )
@@ -69,6 +77,9 @@ from metadata.ingestion.source.database.common_db_source import (
     CommonDbSourceService,
     TableNameAndType,
 )
+from metadata.ingestion.source.database.life_cycle_query_mixin import (
+    LifeCycleQueryMixin,
+)
 from metadata.ingestion.source.database.multi_db_source import MultiDBSource
 from metadata.ingestion.source.database.stored_procedures_mixin import (
     QueryByProcedure,
@@ -80,11 +91,8 @@ from metadata.utils.filters import filter_by_database
 from metadata.utils.helpers import get_start_and_end
 from metadata.utils.logger import ingestion_logger
 from metadata.utils.sqlalchemy_utils import is_complex_type
-from metadata.utils.tag_utils import (
-    get_ometa_tag_and_classification,
-    get_tag_label,
-    get_tag_labels,
-)
+from metadata.utils.tag_utils import get_ometa_tag_and_classification, get_tag_label
+from metadata.utils.tag_utils import get_tag_labels as fetch_tag_labels_om
 
 _bigquery_table_types = {
     "BASE TABLE": TableType.Regular,
@@ -190,9 +198,13 @@ def _build_formatted_table_id(table):
 BigQueryDialect._build_formatted_table_id = (  # pylint: disable=protected-access
     _build_formatted_table_id
 )
+BigQueryDialect.get_pk_constraint = get_pk_constraint
+BigQueryDialect.get_foreign_keys = get_foreign_keys
 
 
-class BigquerySource(StoredProcedureMixin, CommonDbSourceService, MultiDBSource):
+class BigquerySource(
+    LifeCycleQueryMixin, StoredProcedureMixin, CommonDbSourceService, MultiDBSource
+):
     """
     Implements the necessary methods to extract
     Database metadata from Bigquery Source
@@ -204,12 +216,14 @@ class BigquerySource(StoredProcedureMixin, CommonDbSourceService, MultiDBSource)
         # as per service connection config, which would result in an error.
         self.test_connection = lambda: None
         super().__init__(config, metadata)
-        self.temp_credentials = None
         self.client = None
+        # Used to delete temp json file created while initializing bigquery client
+        self.temp_credentials_file_path = []
         # Upon invoking the set_project_id method, we retrieve a comprehensive
         # list of all project IDs. Subsequently, after the invokation,
         # we proceed to test the connections for each of these project IDs
         self.project_ids = self.set_project_id()
+        self.life_cycle_query = BIGQUERY_LIFE_CYCLE_QUERY
         self.test_connection = self._test_connection
         self.test_connection()
 
@@ -237,6 +251,10 @@ class BigquerySource(StoredProcedureMixin, CommonDbSourceService, MultiDBSource)
             test_connection_fn(
                 self.metadata, inspector_details.engine, self.service_connection
             )
+            # GOOGLE_CREDENTIALS may not have been set,
+            # to avoid key error, we use `get` for dict
+            if os.environ.get(GOOGLE_CREDENTIALS):
+                self.temp_credentials_file_path.append(os.environ[GOOGLE_CREDENTIALS])
 
     def query_table_names_and_types(
         self, schema_name: str
@@ -281,6 +299,7 @@ class BigquerySource(StoredProcedureMixin, CommonDbSourceService, MultiDBSource)
                         classification_name=key,
                         tag_description="Bigquery Dataset Label",
                         classification_description="",
+                        include_tags=self.source_config.includeTags,
                     )
             # Fetching policy tags on the column level
             list_project_ids = [self.context.database]
@@ -300,13 +319,14 @@ class BigquerySource(StoredProcedureMixin, CommonDbSourceService, MultiDBSource)
                         classification_name=taxonomy.display_name,
                         tag_description="Bigquery Policy Tag",
                         classification_description="",
+                        include_tags=self.source_config.includeTags,
                     )
         except Exception as exc:
             yield Either(
                 left=StackTraceError(
                     name="Tags and Classifications",
                     error=f"Skipping Policy Tag ingestion due to: {exc}",
-                    stack_trace=traceback.format_exc(),
+                    stackTrace=traceback.format_exc(),
                 )
             )
 
@@ -353,18 +373,18 @@ class BigquerySource(StoredProcedureMixin, CommonDbSourceService, MultiDBSource)
                 schema_name=schema_name,
             ),
         )
-
-        dataset_obj = self.client.get_dataset(schema_name)
-        if dataset_obj.labels:
-            database_schema_request_obj.tags = []
-            for label_classification, label_tag_name in dataset_obj.labels.items():
-                database_schema_request_obj.tags.append(
-                    get_tag_label(
+        if self.source_config.includeTags:
+            dataset_obj = self.client.get_dataset(schema_name)
+            if dataset_obj.labels:
+                database_schema_request_obj.tags = []
+                for label_classification, label_tag_name in dataset_obj.labels.items():
+                    tag_label = get_tag_label(
                         metadata=self.metadata,
                         tag_name=label_tag_name,
                         classification_name=label_classification,
                     )
-                )
+                    if tag_label:
+                        database_schema_request_obj.tags.append(tag_label)
         yield Either(right=database_schema_request_obj)
 
     def get_table_obj(self, table_name: str):
@@ -373,7 +393,7 @@ class BigquerySource(StoredProcedureMixin, CommonDbSourceService, MultiDBSource)
         bq_table_fqn = fqn._build(database, schema_name, table_name)
         return self.client.get_table(bq_table_fqn)
 
-    def yield_table_tag_details(self, table_name_and_type: Tuple[str, str]):
+    def yield_table_tags(self, table_name_and_type: Tuple[str, str]):
         table_name, _ = table_name_and_type
         table_obj = self.get_table_obj(table_name=table_name)
         if table_obj.labels:
@@ -383,6 +403,7 @@ class BigquerySource(StoredProcedureMixin, CommonDbSourceService, MultiDBSource)
                     classification_name=key,
                     tag_description="Bigquery Table Label",
                     classification_description="",
+                    include_tags=self.source_config.includeTags,
                 )
 
     def get_tag_labels(self, table_name: str) -> Optional[List[TagLabel]]:
@@ -411,7 +432,7 @@ class BigquerySource(StoredProcedureMixin, CommonDbSourceService, MultiDBSource)
         is properly informed
         """
         if column.get("policy_tags"):
-            return get_tag_labels(
+            return fetch_tag_labels_om(
                 metadata=self.metadata,
                 tags=[column["policy_tags"]],
                 classification_name=column["taxonomy"],
@@ -423,7 +444,8 @@ class BigquerySource(StoredProcedureMixin, CommonDbSourceService, MultiDBSource)
         inspector_details = get_inspector_details(
             database_name=database_name, service_connection=self.service_connection
         )
-
+        if os.environ.get(GOOGLE_CREDENTIALS):
+            self.temp_credentials_file_path.append(os.environ[GOOGLE_CREDENTIALS])
         self.client = inspector_details.client
         self.engine = inspector_details.engine
         self.inspector = inspector_details.inspector
@@ -516,15 +538,14 @@ class BigquerySource(StoredProcedureMixin, CommonDbSourceService, MultiDBSource)
 
     def close(self):
         super().close()
-        if self.temp_credentials:
-            os.unlink(self.temp_credentials)
         os.environ.pop("GOOGLE_CLOUD_PROJECT", "")
         if isinstance(
             self.service_connection.credentials.gcpConfig, GcpCredentialsValues
         ) and (GOOGLE_CREDENTIALS in os.environ):
-            tmp_credentials_file = os.environ[GOOGLE_CREDENTIALS]
-            os.remove(tmp_credentials_file)
             del os.environ[GOOGLE_CREDENTIALS]
+            for temp_file_path in self.temp_credentials_file_path:
+                if os.path.exists(temp_file_path):
+                    os.remove(temp_file_path)
 
     def _get_source_url(
         self,
@@ -635,7 +656,7 @@ class BigquerySource(StoredProcedureMixin, CommonDbSourceService, MultiDBSource)
                 left=StackTraceError(
                     name=stored_procedure.name,
                     error=f"Error yielding Stored Procedure [{stored_procedure.name}] due to [{exc}]",
-                    stack_trace=traceback.format_exc(),
+                    stackTrace=traceback.format_exc(),
                 )
             )
 

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 
@@ -75,7 +97,6 @@ BIGQUERY_GET_STORED_PROCEDURE_QUERIES = textwrap.dedent(
     """
 WITH SP_HISTORY AS (
   SELECT
-    job_id,
     query AS query_text,
     start_time,
     end_time,
@@ -90,7 +111,6 @@ WITH SP_HISTORY AS (
 ),
 Q_HISTORY AS (
   SELECT
-    job_id,
     project_id as database_name,
     user_email as user_name,
     statement_type as query_type,
@@ -109,8 +129,6 @@ Q_HISTORY AS (
     AND error_result is NULL
 )
 SELECT
-  SP.job_id as procedure_id,
-  Q.job_id as query_id,
   Q.query_type as query_type,
   SP.query_text as procedure_text,
   Q.query_text as query_text,
@@ -130,3 +148,14 @@ JOIN Q_HISTORY Q
 ORDER BY procedure_start_time DESC
 """
 )
+
+BIGQUERY_LIFE_CYCLE_QUERY = textwrap.dedent(
+    """
+select 
+table_name as table_name,
+creation_time as created_at
+from `{schema_name}`.INFORMATION_SCHEMA.TABLES
+where table_schema = '{schema_name}'
+and table_catalog = '{database_name}'
+"""
+)

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

@@ -14,6 +14,7 @@ Handle big query usage extraction
 from abc import ABC
 from copy import deepcopy
 from datetime import datetime
+from typing import Optional
 
 from google import auth
 
@@ -86,3 +87,15 @@ class BigqueryQueryParserSource(QueryParserSource, ABC):
                 yield inspector_details.engine
         else:
             yield self.engine
+
+    def check_life_cycle_query(
+        self, query_type: Optional[str], query_text: Optional[str]
+    ) -> bool:
+        """
+        returns true if query is to be used for life cycle processing.
+
+        Override if we have specific parameters
+        """
+        if query_type != "SELECT":
+            return True
+        return False

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

@@ -26,6 +26,4 @@ class BigqueryUsageSource(BigqueryQueryParserSource, UsageSource):
 
     sql_stmt = BIGQUERY_STATEMENT
 
-    filters = """
-        AND statement_type = "SELECT"
-    """
+    filters = ""

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,
+    )