openmetadata-ingestion 1.1.5.0__py3-none-any.whl → 1.1.6.0__py3-none-any.whl

metadata/generated/schema/settings/__init__.py

@@ -1,3 +1,3 @@
 # generated by datamodel-codegen:
 #   filename:  schema
-#   timestamp: 2023-09-15T12:45:24+00:00
+#   timestamp: 2023-09-25T11:57:29+00:00

metadata/generated/schema/settings/settings.py

@@ -1,6 +1,6 @@
 # generated by datamodel-codegen:
 #   filename:  settings/settings.json
-#   timestamp: 2023-09-15T12:45:24+00:00
+#   timestamp: 2023-09-25T11:57:29+00:00
 
 from __future__ import annotations
 

metadata/generated/schema/system/__init__.py

@@ -1,3 +1,3 @@
 # generated by datamodel-codegen:
 #   filename:  schema
-#   timestamp: 2023-09-15T12:45:24+00:00
+#   timestamp: 2023-09-25T11:57:29+00:00

metadata/generated/schema/system/eventPublisherJob.py

@@ -1,6 +1,6 @@
 # generated by datamodel-codegen:
 #   filename:  system/eventPublisherJob.json
-#   timestamp: 2023-09-15T12:45:24+00:00
+#   timestamp: 2023-09-25T11:57:29+00:00
 
 from __future__ import annotations
 

metadata/generated/schema/system/sqlMigrationScript.py

@@ -1,6 +1,6 @@
 # generated by datamodel-codegen:
 #   filename:  system/sqlMigrationScript.json
-#   timestamp: 2023-09-15T12:45:24+00:00
+#   timestamp: 2023-09-25T11:57:29+00:00
 
 from __future__ import annotations
 

metadata/generated/schema/tests/__init__.py

@@ -1,3 +1,3 @@
 # generated by datamodel-codegen:
 #   filename:  schema
-#   timestamp: 2023-09-15T12:45:24+00:00
+#   timestamp: 2023-09-25T11:57:29+00:00

metadata/generated/schema/tests/basic.py

@@ -1,6 +1,6 @@
 # generated by datamodel-codegen:
 #   filename:  tests/basic.json
-#   timestamp: 2023-09-15T12:45:24+00:00
+#   timestamp: 2023-09-25T11:57:29+00:00
 
 from __future__ import annotations
 

metadata/generated/schema/tests/customMetric.py

@@ -1,6 +1,6 @@
 # generated by datamodel-codegen:
 #   filename:  tests/customMetric.json
-#   timestamp: 2023-09-15T12:45:24+00:00
+#   timestamp: 2023-09-25T11:57:29+00:00
 
 from __future__ import annotations
 

metadata/generated/schema/tests/testCase.py

@@ -1,6 +1,6 @@
 # generated by datamodel-codegen:
 #   filename:  tests/testCase.json
-#   timestamp: 2023-09-15T12:45:24+00:00
+#   timestamp: 2023-09-25T11:57:29+00:00
 
 from __future__ import annotations
 

metadata/generated/schema/tests/testDefinition.py

@@ -1,6 +1,6 @@
 # generated by datamodel-codegen:
 #   filename:  tests/testDefinition.json
-#   timestamp: 2023-09-15T12:45:24+00:00
+#   timestamp: 2023-09-25T11:57:29+00:00
 
 from __future__ import annotations
 

metadata/generated/schema/tests/testSuite.py

@@ -1,6 +1,6 @@
 # generated by datamodel-codegen:
 #   filename:  tests/testSuite.json
-#   timestamp: 2023-09-15T12:45:24+00:00
+#   timestamp: 2023-09-25T11:57:29+00:00
 
 from __future__ import annotations
 

metadata/generated/schema/type/__init__.py

@@ -1,3 +1,3 @@
 # generated by datamodel-codegen:
 #   filename:  schema
-#   timestamp: 2023-09-15T12:45:24+00:00
+#   timestamp: 2023-09-25T11:57:29+00:00

metadata/generated/schema/type/auditLog.py

@@ -1,6 +1,6 @@
 # generated by datamodel-codegen:
 #   filename:  type/auditLog.json
-#   timestamp: 2023-09-15T12:45:24+00:00
+#   timestamp: 2023-09-25T11:57:29+00:00
 
 from __future__ import annotations
 

metadata/generated/schema/type/basic.py

@@ -1,6 +1,6 @@
 # generated by datamodel-codegen:
 #   filename:  type/basic.json
-#   timestamp: 2023-09-15T12:45:24+00:00
+#   timestamp: 2023-09-25T11:57:29+00:00
 
 from __future__ import annotations
 

metadata/generated/schema/type/changeEvent.py

@@ -1,6 +1,6 @@
 # generated by datamodel-codegen:
 #   filename:  type/changeEvent.json
-#   timestamp: 2023-09-15T12:45:24+00:00
+#   timestamp: 2023-09-25T11:57:29+00:00
 
 from __future__ import annotations
 

metadata/generated/schema/type/collectionDescriptor.py

@@ -1,6 +1,6 @@
 # generated by datamodel-codegen:
 #   filename:  type/collectionDescriptor.json
-#   timestamp: 2023-09-15T12:45:24+00:00
+#   timestamp: 2023-09-25T11:57:29+00:00
 
 from __future__ import annotations
 

metadata/generated/schema/type/csvDocumentation.py

@@ -1,6 +1,6 @@
 # generated by datamodel-codegen:
 #   filename:  type/csvDocumentation.json
-#   timestamp: 2023-09-15T12:45:24+00:00
+#   timestamp: 2023-09-25T11:57:29+00:00
 
 from __future__ import annotations
 

metadata/generated/schema/type/csvErrorType.py

@@ -1,6 +1,6 @@
 # generated by datamodel-codegen:
 #   filename:  type/csvErrorType.json
-#   timestamp: 2023-09-15T12:45:24+00:00
+#   timestamp: 2023-09-25T11:57:29+00:00
 
 from __future__ import annotations
 

metadata/generated/schema/type/csvFile.py

@@ -1,6 +1,6 @@
 # generated by datamodel-codegen:
 #   filename:  type/csvFile.json
-#   timestamp: 2023-09-15T12:45:24+00:00
+#   timestamp: 2023-09-25T11:57:29+00:00
 
 from __future__ import annotations
 

metadata/generated/schema/type/csvImportResult.py

@@ -1,6 +1,6 @@
 # generated by datamodel-codegen:
 #   filename:  type/csvImportResult.json
-#   timestamp: 2023-09-15T12:45:24+00:00
+#   timestamp: 2023-09-25T11:57:29+00:00
 
 from __future__ import annotations
 

metadata/generated/schema/type/dailyCount.py

@@ -1,6 +1,6 @@
 # generated by datamodel-codegen:
 #   filename:  type/dailyCount.json
-#   timestamp: 2023-09-15T12:45:24+00:00
+#   timestamp: 2023-09-25T11:57:29+00:00
 
 from __future__ import annotations
 

metadata/generated/schema/type/databaseConnectionConfig.py

@@ -1,6 +1,6 @@
 # generated by datamodel-codegen:
 #   filename:  type/databaseConnectionConfig.json
-#   timestamp: 2023-09-15T12:45:24+00:00
+#   timestamp: 2023-09-25T11:57:29+00:00
 
 from __future__ import annotations
 

metadata/generated/schema/type/entityHistory.py

@@ -1,6 +1,6 @@
 # generated by datamodel-codegen:
 #   filename:  type/entityHistory.json
-#   timestamp: 2023-09-15T12:45:24+00:00
+#   timestamp: 2023-09-25T11:57:29+00:00
 
 from __future__ import annotations
 

metadata/generated/schema/type/entityLineage.py

@@ -1,6 +1,6 @@
 # generated by datamodel-codegen:
 #   filename:  type/entityLineage.json
-#   timestamp: 2023-09-15T12:45:24+00:00
+#   timestamp: 2023-09-25T11:57:29+00:00
 
 from __future__ import annotations
 

metadata/generated/schema/type/entityReference.py

@@ -1,6 +1,6 @@
 # generated by datamodel-codegen:
 #   filename:  type/entityReference.json
-#   timestamp: 2023-09-15T12:45:24+00:00
+#   timestamp: 2023-09-25T11:57:29+00:00
 
 from __future__ import annotations
 

metadata/generated/schema/type/entityReferenceList.py

@@ -1,6 +1,6 @@
 # generated by datamodel-codegen:
 #   filename:  type/entityReferenceList.json
-#   timestamp: 2023-09-15T12:45:24+00:00
+#   timestamp: 2023-09-25T11:57:29+00:00
 
 from __future__ import annotations
 

metadata/generated/schema/type/entityRelationship.py

@@ -1,6 +1,6 @@
 # generated by datamodel-codegen:
 #   filename:  type/entityRelationship.json
-#   timestamp: 2023-09-15T12:45:24+00:00
+#   timestamp: 2023-09-25T11:57:29+00:00
 
 from __future__ import annotations
 

metadata/generated/schema/type/entityUsage.py

@@ -1,6 +1,6 @@
 # generated by datamodel-codegen:
 #   filename:  type/entityUsage.json
-#   timestamp: 2023-09-15T12:45:24+00:00
+#   timestamp: 2023-09-25T11:57:29+00:00
 
 from __future__ import annotations
 

metadata/generated/schema/type/filterPattern.py

@@ -1,6 +1,6 @@
 # generated by datamodel-codegen:
 #   filename:  type/filterPattern.json
-#   timestamp: 2023-09-15T12:45:24+00:00
+#   timestamp: 2023-09-25T11:57:29+00:00
 
 from __future__ import annotations
 

metadata/generated/schema/type/function.py

@@ -1,6 +1,6 @@
 # generated by datamodel-codegen:
 #   filename:  type/function.json
-#   timestamp: 2023-09-15T12:45:24+00:00
+#   timestamp: 2023-09-25T11:57:29+00:00
 
 from __future__ import annotations
 

metadata/generated/schema/type/include.py

@@ -1,6 +1,6 @@
 # generated by datamodel-codegen:
 #   filename:  type/include.json
-#   timestamp: 2023-09-15T12:45:24+00:00
+#   timestamp: 2023-09-25T11:57:29+00:00
 
 from __future__ import annotations
 

metadata/generated/schema/type/jdbcConnection.py

@@ -1,6 +1,6 @@
 # generated by datamodel-codegen:
 #   filename:  type/jdbcConnection.json
-#   timestamp: 2023-09-15T12:45:24+00:00
+#   timestamp: 2023-09-25T11:57:29+00:00
 
 from __future__ import annotations
 

metadata/generated/schema/type/paging.py

@@ -1,6 +1,6 @@
 # generated by datamodel-codegen:
 #   filename:  type/paging.json
-#   timestamp: 2023-09-15T12:45:24+00:00
+#   timestamp: 2023-09-25T11:57:29+00:00
 
 from __future__ import annotations
 

metadata/generated/schema/type/profile.py

@@ -1,6 +1,6 @@
 # generated by datamodel-codegen:
 #   filename:  type/profile.json
-#   timestamp: 2023-09-15T12:45:24+00:00
+#   timestamp: 2023-09-25T11:57:29+00:00
 
 from __future__ import annotations
 

metadata/generated/schema/type/queryParserData.py

@@ -1,6 +1,6 @@
 # generated by datamodel-codegen:
 #   filename:  type/queryParserData.json
-#   timestamp: 2023-09-15T12:45:24+00:00
+#   timestamp: 2023-09-25T11:57:29+00:00
 
 from __future__ import annotations
 

metadata/generated/schema/type/reaction.py

@@ -1,6 +1,6 @@
 # generated by datamodel-codegen:
 #   filename:  type/reaction.json
-#   timestamp: 2023-09-15T12:45:24+00:00
+#   timestamp: 2023-09-25T11:57:29+00:00
 
 from __future__ import annotations
 

metadata/generated/schema/type/schedule.py

@@ -1,6 +1,6 @@
 # generated by datamodel-codegen:
 #   filename:  type/schedule.json
-#   timestamp: 2023-09-15T12:45:24+00:00
+#   timestamp: 2023-09-25T11:57:29+00:00
 
 from __future__ import annotations
 

metadata/generated/schema/type/schema.py

@@ -1,6 +1,6 @@
 # generated by datamodel-codegen:
 #   filename:  type/schema.json
-#   timestamp: 2023-09-15T12:45:24+00:00
+#   timestamp: 2023-09-25T11:57:29+00:00
 
 from __future__ import annotations
 

metadata/generated/schema/type/tableQuery.py

@@ -1,6 +1,6 @@
 # generated by datamodel-codegen:
 #   filename:  type/tableQuery.json
-#   timestamp: 2023-09-15T12:45:24+00:00
+#   timestamp: 2023-09-25T11:57:29+00:00
 
 from __future__ import annotations
 

metadata/generated/schema/type/tableUsageCount.py

@@ -1,6 +1,6 @@
 # generated by datamodel-codegen:
 #   filename:  type/tableUsageCount.json
-#   timestamp: 2023-09-15T12:45:24+00:00
+#   timestamp: 2023-09-25T11:57:29+00:00
 
 from __future__ import annotations
 

metadata/generated/schema/type/tagLabel.py

@@ -1,6 +1,6 @@
 # generated by datamodel-codegen:
 #   filename:  type/tagLabel.json
-#   timestamp: 2023-09-15T12:45:24+00:00
+#   timestamp: 2023-09-25T11:57:29+00:00
 
 from __future__ import annotations
 

metadata/generated/schema/type/usageDetails.py

@@ -1,6 +1,6 @@
 # generated by datamodel-codegen:
 #   filename:  type/usageDetails.json
-#   timestamp: 2023-09-15T12:45:24+00:00
+#   timestamp: 2023-09-25T11:57:29+00:00
 
 from __future__ import annotations
 

metadata/generated/schema/type/usageRequest.py

@@ -1,6 +1,6 @@
 # generated by datamodel-codegen:
 #   filename:  type/usageRequest.json
-#   timestamp: 2023-09-15T12:45:24+00:00
+#   timestamp: 2023-09-25T11:57:29+00:00
 
 from __future__ import annotations
 

metadata/generated/schema/type/votes.py

@@ -1,6 +1,6 @@
 # generated by datamodel-codegen:
 #   filename:  type/votes.json
-#   timestamp: 2023-09-15T12:45:24+00:00
+#   timestamp: 2023-09-25T11:57:29+00:00
 
 from __future__ import annotations
 

metadata/ingestion/lineage/sql_lineage.py

@@ -12,7 +12,7 @@
 Helper functions to handle SQL lineage operations
 """
 import traceback
-from typing import Any, Iterable, Iterator, List, Optional
+from typing import Any, Iterable, Iterator, List, Optional, Tuple
 
 from metadata.generated.schema.api.lineage.addLineage import AddLineageRequest
 from metadata.generated.schema.entity.data.table import Table
@@ -118,6 +118,34 @@ def search_table_entities(
         return None
 
 
+def get_table_fqn_from_query_name(
+    table_name: str,
+) -> Tuple[Optional[str], Optional[str], Optional[str]]:
+    """
+    Method to extract database, schema and table name
+    from raw table name used in query
+    """
+
+    split_table = table_name.split(".")
+    empty_list: List[Any] = [None]  # Otherwise, there's a typing error in the concat
+
+    if len(split_table) > 3:
+        # In case of bigquery, it is possible that tables within information schema when
+        # referred with their fully qualified name may look like this
+        # `project-id.dataset-id.information_schema.table-name` in such cases there
+        # will be 4 values to unpack vs the expected 3 values, hence in such case we
+        # just pick the table name and keep the database and schema name as none
+
+        table = split_table[-1]
+        database_query, schema_query = None, None
+    else:
+        database_query, schema_query, table = (
+            empty_list * (3 - len(split_table))
+        ) + split_table
+
+    return database_query, schema_query, table
+
+
 def get_table_entities_from_query(
     metadata: OpenMetadata,
     service_name: str,
@@ -141,12 +169,7 @@ def get_table_entities_from_query(
     # First try to find the data from the given db and schema (with table name as given or uppercase)
     # Otherwise, pick it up from the table_name str (with table name as given or uppercase)
 
-    split_table = table_name.split(".")
-    empty_list: List[Any] = [None]  # Otherwise, there's a typing error in the concat
-
-    database_query, schema_query, table = (
-        empty_list * (3 - len(split_table))
-    ) + split_table
+    database_query, schema_query, table = get_table_fqn_from_query_name(table_name)
 
     table_entities = search_table_entities(
         metadata=metadata,

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

@@ -16,7 +16,6 @@ import os
 from functools import partial
 from typing import Optional
 
-from google import auth
 from google.cloud.datacatalog_v1 import PolicyTagManagerClient
 from sqlalchemy.engine import Engine
 
@@ -98,8 +97,6 @@ def test_connection(
     Test connection. This can be executed either as part
     of a metadata workflow or during an Automation Workflow
     """
-    _, project_ids = auth.default()
-    project_ids = project_ids if isinstance(project_ids, list) else [project_ids]
 
     def get_tags(taxonomies):
         for taxonomy in taxonomies:
@@ -109,12 +106,10 @@ def test_connection(
             return policy_tags
 
     def test_tags():
-        for project in project_ids:
-            taxonomies = PolicyTagManagerClient().list_taxonomies(
-                parent=f"projects/{project}/locations/{service_connection.taxonomyLocation}"
-            )
-            return get_tags(taxonomies)
-        return None
+        taxonomies = PolicyTagManagerClient().list_taxonomies(
+            parent=f"projects/{engine.url.host}/locations/{service_connection.taxonomyLocation}"
+        )
+        return get_tags(taxonomies)
 
     def test_connection_inner(engine):
         test_fn = {
@@ -139,7 +134,4 @@ def test_connection(
             automation_workflow=automation_workflow,
         )
 
-    for project in project_ids:
-        if project in str(engine.url):
-            continue
-        test_connection_inner(engine)
+    test_connection_inner(engine)

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

@@ -0,0 +1,66 @@
+#  Copyright 2021 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.
+
+"""
+Source connection helper
+"""
+
+from typing import Any
+
+from pydantic import BaseModel
+from sqlalchemy import inspect
+
+from metadata.generated.schema.entity.services.connections.database.bigQueryConnection import (
+    BigQueryConnection,
+)
+from metadata.generated.schema.security.credentials.gcpValues import (
+    GcpCredentialsValues,
+    SingleProjectId,
+)
+from metadata.ingestion.source.connections import get_connection
+from metadata.utils.bigquery_utils import get_bigquery_client
+
+
+class InspectorWrapper(BaseModel):
+    client: Any
+    engine: Any
+    inspector: Any
+
+
+def get_inspector_details(
+    database_name: str, service_connection: BigQueryConnection
+) -> InspectorWrapper:
+    """
+    Method to get the bigquery inspector details
+    """
+    # TODO support location property in JSON Schema
+    # TODO support OAuth 2.0 scopes
+    kwargs = {}
+    if isinstance(service_connection.credentials.gcpConfig, GcpCredentialsValues):
+        service_connection.credentials.gcpConfig.projectId = SingleProjectId(
+            __root__=database_name
+        )
+        if service_connection.credentials.gcpImpersonateServiceAccount:
+            kwargs[
+                "impersonate_service_account"
+            ] = (
+                service_connection.credentials.gcpImpersonateServiceAccount.impersonateServiceAccount
+            )
+
+            kwargs[
+                "lifetime"
+            ] = service_connection.credentials.gcpImpersonateServiceAccount.lifetime
+
+    client = get_bigquery_client(project_id=database_name, **kwargs)
+    engine = get_connection(service_connection)
+    inspector = inspect(engine)
+
+    return InspectorWrapper(client=client, engine=engine, inspector=inspector)

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

@@ -17,7 +17,6 @@ from typing import Iterable, List, Optional, Tuple
 
 from google import auth
 from google.cloud.datacatalog_v1 import PolicyTagManagerClient
-from sqlalchemy import inspect
 from sqlalchemy.engine.reflection import Inspector
 from sqlalchemy.sql.sqltypes import Interval
 from sqlalchemy.types import String
@@ -44,12 +43,12 @@ from metadata.generated.schema.metadataIngestion.workflow import (
 )
 from metadata.generated.schema.security.credentials.gcpValues import (
     GcpCredentialsValues,
-    SingleProjectId,
 )
 from metadata.generated.schema.type.tagLabel import TagLabel
 from metadata.ingestion.api.source import InvalidSourceException
 from metadata.ingestion.models.ometa_classification import OMetaTagAndClassification
-from metadata.ingestion.source.connections import get_connection
+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.queries import (
     BIGQUERY_SCHEMA_DESCRIPTION,
     BIGQUERY_TABLE_AND_TYPE,
@@ -60,7 +59,6 @@ from metadata.ingestion.source.database.common_db_source import (
     TableNameAndType,
 )
 from metadata.utils import fqn
-from metadata.utils.bigquery_utils import get_bigquery_client
 from metadata.utils.credentials import GOOGLE_CREDENTIALS
 from metadata.utils.filters import filter_by_database
 from metadata.utils.logger import ingestion_logger
@@ -204,6 +202,12 @@ class BigquerySource(CommonDbSourceService):
         _, project_ids = auth.default()
         return project_ids if isinstance(project_ids, list) else [project_ids]
 
+    def test_connection(self) -> None:
+        for project_id in self.set_project_id():
+            self.set_inspector(project_id)
+            test_connection_fn = get_test_connection_fn(self.service_connection)
+            test_connection_fn(self.metadata, self.engine, self.service_connection)
+
     def query_table_names_and_types(
         self, schema_name: str
     ) -> Iterable[TableNameAndType]:
@@ -374,31 +378,13 @@ class BigquerySource(CommonDbSourceService):
         return None
 
     def set_inspector(self, database_name: str):
-        # TODO support location property in JSON Schema
-        # TODO support OAuth 2.0 scopes
-        kwargs = {}
-        if isinstance(
-            self.service_connection.credentials.gcpConfig, GcpCredentialsValues
-        ):
-            self.service_connection.credentials.gcpConfig.projectId = SingleProjectId(
-                __root__=database_name
-            )
-            if self.service_connection.credentials.gcpImpersonateServiceAccount:
-                kwargs[
-                    "impersonate_service_account"
-                ] = (
-                    self.service_connection.credentials.gcpImpersonateServiceAccount.impersonateServiceAccount
-                )
-
-                kwargs[
-                    "lifetime"
-                ] = (
-                    self.service_connection.credentials.gcpImpersonateServiceAccount.lifetime
-                )
+        inspector_details = get_inspector_details(
+            database_name=database_name, service_connection=self.service_connection
+        )
 
-        self.client = get_bigquery_client(project_id=database_name, **kwargs)
-        self.engine = get_connection(self.service_connection)
-        self.inspector = inspect(self.engine)
+        self.client = inspector_details.client
+        self.engine = inspector_details.engine
+        self.inspector = inspector_details.inspector
 
     def get_database_names(self) -> Iterable[str]:
         for project_id in self.project_ids: