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

metadata/ingestion/models/patch_request.py

@@ -80,6 +80,7 @@ ALLOWED_COMMON_PATCH_FIELDS = {
     "description": True,
     "owner": True,
     "tags": True,
+    "sourceHash": True,
     # Table Entity Fields
     "tableType": True,
     "columns": {"__all__": ALLOWED_COLUMN_FIELDS},
@@ -210,9 +211,9 @@ def _determine_restricted_operation(
     Only retain add operation for restrict_update_fields fields
     """
     path = patch_ops.get("path")
-    op = patch_ops.get("op")
+    ops = patch_ops.get("op")
     for field in restrict_update_fields or []:
-        if field in path and op != PatchOperation.ADD.value:
+        if field in path and ops != PatchOperation.ADD.value:
             return False
     return True
 

metadata/ingestion/ometa/client.py

@@ -21,6 +21,7 @@ from requests.exceptions import HTTPError
 
 from metadata.config.common import ConfigModel
 from metadata.ingestion.ometa.credentials import URL, get_api_version
+from metadata.utils.execution_time_tracker import calculate_execution_time
 from metadata.utils.logger import ometa_logger
 
 logger = ometa_logger()
@@ -251,6 +252,7 @@ class REST:
 
         return None
 
+    @calculate_execution_time(context="GET")
     def get(self, path, data=None):
         """
         GET method
@@ -264,6 +266,7 @@ class REST:
         """
         return self._request("GET", path, data)
 
+    @calculate_execution_time(context="POST")
     def post(self, path, data=None):
         """
         POST method
@@ -277,6 +280,7 @@ class REST:
         """
         return self._request("POST", path, data)
 
+    @calculate_execution_time(context="PUT")
     def put(self, path, data=None):
         """
         PUT method
@@ -290,6 +294,7 @@ class REST:
         """
         return self._request("PUT", path, data)
 
+    @calculate_execution_time(context="PATCH")
     def patch(self, path, data=None):
         """
         PATCH method
@@ -308,6 +313,7 @@ class REST:
             headers={"Content-type": "application/json-patch+json"},
         )
 
+    @calculate_execution_time(context="DELETE")
     def delete(self, path, data=None):
         """
         DELETE method

metadata/ingestion/ometa/mixins/custom_property_mixin.py

@@ -15,7 +15,8 @@ To be used by OpenMetadata class
 """
 from typing import Dict
 
-from metadata.generated.schema.api.data.createCustomProperty import PropertyType
+from metadata.generated.schema.type.customProperty import PropertyType
+from metadata.generated.schema.type.entityReference import EntityReference
 from metadata.ingestion.models.custom_properties import (
     CustomPropertyDataTypes,
     CustomPropertyType,
@@ -54,16 +55,6 @@ class OMetaCustomPropertyMixin:
             f"/metadata/types/name/{entity_type}?category=field"
         )
 
-        # Get the data type of the custom property
-        if not ometa_custom_property.createCustomPropertyRequest.propertyType:
-            custom_property_type = self.get_custom_property_type(
-                data_type=ometa_custom_property.custom_property_type
-            )
-            property_type = PropertyType(id=custom_property_type.id, type="type")
-            ometa_custom_property.createCustomPropertyRequest.propertyType = (
-                property_type
-            )
-
         resp = self.client.put(
             f"/metadata/types/{entity_schema.get('id')}",
             data=ometa_custom_property.createCustomPropertyRequest.json(),
@@ -78,3 +69,12 @@ class OMetaCustomPropertyMixin:
         """
         resp = self.client.get(f"/metadata/types/name/{data_type.value}?category=field")
         return CustomPropertyType(**resp)
+
+    def get_property_type_ref(self, data_type: CustomPropertyDataTypes) -> PropertyType:
+        """
+        Get the PropertyType for custom properties
+        """
+        custom_property_type = self.get_custom_property_type(data_type=data_type)
+        return PropertyType(
+            __root__=EntityReference(id=custom_property_type.id, type="type")
+        )

metadata/ingestion/ometa/mixins/patch_mixin.py

@@ -135,7 +135,6 @@ class OMetaPatchMixin(OMetaPatchMixinBase):
             Updated Entity
         """
         try:
-
             patch = build_patch(
                 source=source,
                 destination=destination,

metadata/ingestion/ometa/mixins/suggestions_mixin.py

@@ -0,0 +1,41 @@
+#  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.
+"""
+Mixin class containing Suggestions specific methods
+
+To be used by OpenMetadata class
+"""
+from metadata.generated.schema.entity.feed.suggestion import Suggestion
+from metadata.ingestion.ometa.client import REST
+from metadata.utils.logger import ometa_logger
+
+logger = ometa_logger()
+
+
+class OMetaSuggestionsMixin:
+    """
+    OpenMetadata API methods related to the Suggestion Entity
+
+    To be inherited by OpenMetadata
+    """
+
+    client: REST
+
+    def update_suggestion(self, suggestion: Suggestion) -> Suggestion:
+        """
+        Update an existing Suggestion with new fields
+        """
+        resp = self.client.put(
+            f"{self.get_suffix(Suggestion)}/{str(suggestion.id.__root__)}",
+            data=suggestion.json(),
+        )
+
+        return Suggestion(**resp)

metadata/ingestion/ometa/ometa_api.py

@@ -50,6 +50,7 @@ from metadata.ingestion.ometa.mixins.role_policy_mixin import OMetaRolePolicyMix
 from metadata.ingestion.ometa.mixins.search_index_mixin import OMetaSearchIndexMixin
 from metadata.ingestion.ometa.mixins.server_mixin import OMetaServerMixin
 from metadata.ingestion.ometa.mixins.service_mixin import OMetaServiceMixin
+from metadata.ingestion.ometa.mixins.suggestions_mixin import OMetaSuggestionsMixin
 from metadata.ingestion.ometa.mixins.table_mixin import OMetaTableMixin
 from metadata.ingestion.ometa.mixins.tests_mixin import OMetaTestsMixin
 from metadata.ingestion.ometa.mixins.topic_mixin import OMetaTopicMixin
@@ -108,6 +109,7 @@ class OpenMetadata(
     OMetaRolePolicyMixin,
     OMetaSearchIndexMixin,
     OMetaCustomPropertyMixin,
+    OMetaSuggestionsMixin,
     Generic[T, C],
 ):
     """
@@ -246,11 +248,9 @@ class OpenMetadata(
         )
         return entity_class
 
-    def create_or_update(self, data: C) -> T:
+    def _create(self, data: C, method: str) -> T:
         """
-        We allow CreateEntity for PUT, so we expect a type C.
-
-        We PUT to the endpoint and return the Entity generated result
+        Internal logic to run POST vs. PUT
         """
         entity = data.__class__
         is_create = "create" in data.__class__.__name__.lower()
@@ -262,15 +262,23 @@ class OpenMetadata(
             raise InvalidEntityException(
                 f"PUT operations need a CreateEntity, not {entity}"
             )
-        resp = self.client.put(
-            self.get_suffix(entity), data=data.json(encoder=show_secrets_encoder)
-        )
+
+        fn = getattr(self.client, method)
+        resp = fn(self.get_suffix(entity), data=data.json(encoder=show_secrets_encoder))
         if not resp:
             raise EmptyPayloadException(
                 f"Got an empty response when trying to PUT to {self.get_suffix(entity)}, {data.json()}"
             )
         return entity_class(**resp)
 
+    def create_or_update(self, data: C) -> T:
+        """Run a PUT requesting via create request C"""
+        return self._create(data=data, method="put")
+
+    def create(self, data: C) -> T:
+        """Run a POST requesting via create request C"""
+        return self._create(data=data, method="post")
+
     def get_by_name(
         self,
         entity: Type[T],
@@ -410,7 +418,7 @@ class OpenMetadata(
         self,
         entity: Type[T],
         fields: Optional[List[str]] = None,
-        limit: int = 1000,
+        limit: int = 100,
         params: Optional[Dict[str, str]] = None,
         skip_on_failure: bool = False,
     ) -> Iterable[T]:

metadata/ingestion/ometa/routes.py

@@ -50,6 +50,7 @@ from metadata.generated.schema.api.domains.createDataProduct import (
     CreateDataProductRequest,
 )
 from metadata.generated.schema.api.domains.createDomain import CreateDomainRequest
+from metadata.generated.schema.api.feed.createSuggestion import CreateSuggestionRequest
 from metadata.generated.schema.api.lineage.addLineage import AddLineageRequest
 from metadata.generated.schema.api.policies.createPolicy import CreatePolicyRequest
 from metadata.generated.schema.api.services.createDashboardService import (
@@ -90,6 +91,7 @@ from metadata.generated.schema.api.tests.createTestSuite import CreateTestSuiteR
 from metadata.generated.schema.dataInsight.dataInsightChart import DataInsightChart
 from metadata.generated.schema.dataInsight.kpi.kpi import Kpi
 from metadata.generated.schema.entity.automations.workflow import Workflow
+from metadata.generated.schema.entity.bot import Bot
 from metadata.generated.schema.entity.classification.classification import (
     Classification,
 )
@@ -113,6 +115,7 @@ from metadata.generated.schema.entity.data.table import Table
 from metadata.generated.schema.entity.data.topic import Topic
 from metadata.generated.schema.entity.domains.dataProduct import DataProduct
 from metadata.generated.schema.entity.domains.domain import Domain
+from metadata.generated.schema.entity.feed.suggestion import Suggestion
 from metadata.generated.schema.entity.policies.policy import Policy
 from metadata.generated.schema.entity.services.connections.testConnectionDefinition import (
     TestConnectionDefinition,
@@ -181,6 +184,7 @@ ROUTES = {
     User.__name__: "/users",
     CreateUserRequest.__name__: "/users",
     AuthenticationMechanism.__name__: "/users/auth-mechanism",
+    Bot.__name__: "/bots",  # We won't allow bot creation from the client
     # Roles
     Role.__name__: "/roles",
     CreateRoleRequest.__name__: "/roles",
@@ -225,4 +229,7 @@ ROUTES = {
     CreateDomainRequest.__name__: "/domains",
     DataProduct.__name__: "/dataProducts",
     CreateDataProductRequest.__name__: "/dataProducts",
+    # Suggestions
+    Suggestion.__name__: "/suggestions",
+    CreateSuggestionRequest.__name__: "/suggestions",
 }

metadata/ingestion/sink/metadata_rest.py

@@ -82,7 +82,7 @@ from metadata.ingestion.ometa.ometa_api import OpenMetadata
 from metadata.ingestion.source.dashboard.dashboard_service import DashboardUsage
 from metadata.ingestion.source.database.database_service import DataModelLink
 from metadata.profiler.api.models import ProfilerResponse
-from metadata.utils.helpers import calculate_execution_time
+from metadata.utils.execution_time_tracker import calculate_execution_time
 from metadata.utils.logger import get_log_name, ingestion_logger
 
 logger = ingestion_logger()
@@ -129,7 +129,7 @@ class MetadataRestSink(Sink):  # pylint: disable=too-many-public-methods
         logger.debug(f"Processing Create request {type(record)}")
         return self.write_create_request(record)
 
-    @calculate_execution_time
+    @calculate_execution_time(store=False)
     def _run(self, record: Entity, *_, **__) -> Either[Any]:
         """
         Default implementation for the single dispatch

metadata/ingestion/source/dashboard/dashboard_service.py

@@ -531,7 +531,6 @@ class DashboardServiceSource(TopologyRunnerMixin, Source, ABC):
         )
 
     def check_database_schema_name(self, database_schema_name: str):
-
         """
         Check if the input database schema name is equal to "<default>" and return the input name if it is not.
 

metadata/ingestion/source/dashboard/looker/metadata.py

@@ -661,7 +661,7 @@ class LookerSource(DashboardServiceSource):
             ],
             # Dashboards are created from the UI directly. They are not linked to a project
             # like LookML assets, but rather just organised in folders.
-            project=self._get_dashboard_project(dashboard_details),
+            project=self.get_project_name(dashboard_details),
             sourceUrl=f"{clean_uri(self.service_connection.hostPort)}/dashboards/{dashboard_details.id}",
             service=self.context.dashboard_service,
             owner=self.get_owner_ref(dashboard_details=dashboard_details),
@@ -669,8 +669,7 @@ class LookerSource(DashboardServiceSource):
         yield Either(right=dashboard_request)
         self.register_record(dashboard_request=dashboard_request)
 
-    @staticmethod
-    def _get_dashboard_project(dashboard_details: LookerDashboard) -> Optional[str]:
+    def get_project_name(self, dashboard_details: LookerDashboard) -> Optional[str]:
         """
         Get dashboard project if the folder is informed
         """
@@ -680,7 +679,7 @@ class LookerSource(DashboardServiceSource):
             logger.debug(
                 f"Cannot get folder name from dashboard [{dashboard_details.title}] - [{exc}]"
             )
-            return None
+        return None
 
     @staticmethod
     def _clean_table_name(table_name: str) -> str:
@@ -833,7 +832,6 @@ class LookerSource(DashboardServiceSource):
                     to_entity.id.__root__
                     not in self._added_lineage[from_entity.id.__root__]
                 ):
-
                     self._added_lineage[from_entity.id.__root__].append(
                         to_entity.id.__root__
                     )
@@ -943,7 +941,6 @@ class LookerSource(DashboardServiceSource):
         dashboard_name = self.context.dashboard
 
         try:
-
             dashboard_fqn = fqn.build(
                 metadata=self.metadata,
                 entity_type=Dashboard,

metadata/ingestion/source/dashboard/metabase/client.py

@@ -124,6 +124,10 @@ class MetabaseClient:
         try:
             resp_dashboard = self.client.get(f"/dashboard/{dashboard_id}")
             if resp_dashboard:
+                # Small hack needed to support Metabase versions older than 0.48
+                # https://www.metabase.com/releases/metabase-48#fyi--breaking-changes
+                if "ordered_cards" in resp_dashboard:
+                    resp_dashboard["dashcards"] = resp_dashboard["ordered_cards"]
                 return MetabaseDashboardDetails(**resp_dashboard)
         except Exception:
             logger.debug(traceback.format_exc())

metadata/ingestion/source/dashboard/metabase/metadata.py

@@ -178,7 +178,7 @@ class MetabaseSource(DashboardServiceSource):
         Returns:
             Iterable[CreateChartRequest]
         """
-        charts = dashboard_details.ordered_cards
+        charts = dashboard_details.dashcards
         for chart in charts:
             try:
                 chart_details = chart.card
@@ -225,7 +225,7 @@ class MetabaseSource(DashboardServiceSource):
         if not db_service_name:
             return
         chart_list, dashboard_name = (
-            dashboard_details.ordered_cards,
+            dashboard_details.dashcards,
             str(dashboard_details.id),
         )
         for chart in chart_list:
@@ -324,8 +324,9 @@ class MetabaseSource(DashboardServiceSource):
         self, chart_details: MetabaseChart, db_service_name: str, dashboard_name: str
     ) -> Iterable[Either[AddLineageRequest]]:
         table = self.client.get_table(chart_details.table_id)
+        table_name = table.name or table.display_name
 
-        if table is None or table.display_name is None:
+        if table is None or table_name is None:
             return
 
         database_name = table.db.details.db if table.db and table.db.details else None
@@ -334,7 +335,7 @@ class MetabaseSource(DashboardServiceSource):
             database=database_name,
             service_name=db_service_name,
             database_schema=table.table_schema,
-            table=table.display_name,
+            table=table_name,
         )
 
         to_fqn = fqn.build(

metadata/ingestion/source/dashboard/metabase/models.py

@@ -67,7 +67,7 @@ class MetabaseChart(BaseModel):
     display: Optional[str]
 
 
-class OrderedCard(BaseModel):
+class DashCard(BaseModel):
     card: MetabaseChart
 
 
@@ -77,7 +77,7 @@ class MetabaseDashboardDetails(BaseModel):
     """
 
     description: Optional[str]
-    ordered_cards: List[OrderedCard]
+    dashcards: List[DashCard]
     name: Optional[str]
     id: int
     collection_id: Optional[str]

metadata/ingestion/source/dashboard/superset/api_source.py

@@ -192,7 +192,6 @@ class SupersetAPISource(SupersetSourceMixin):
     def yield_datamodel(
         self, dashboard_details: DashboardResult
     ) -> Iterable[Either[CreateDashboardDataModelRequest]]:
-
         if self.source_config.includeDataModels:
             for chart_id in self._get_charts_of_dashboard(dashboard_details):
                 try:

metadata/ingestion/source/dashboard/superset/db_source.py

@@ -15,7 +15,6 @@ Superset source module
 import traceback
 from typing import Iterable, Optional
 
-from sqlalchemy import sql
 from sqlalchemy.engine import Engine
 from sqlalchemy.engine.url import make_url
 
@@ -87,7 +86,7 @@ class SupersetDBSource(SupersetSourceMixin):
         try:
             if table_name:
                 col_list = self.engine.execute(
-                    sql.text(FETCH_COLUMN), table_name=table_name.lower()
+                    FETCH_COLUMN, table_name=table_name.lower()
                 )
                 return [FetchColumn(**col) for col in col_list]
         except Exception as err:
@@ -216,7 +215,6 @@ class SupersetDBSource(SupersetSourceMixin):
     def yield_datamodel(
         self, dashboard_details: FetchDashboard
     ) -> Iterable[Either[CreateDashboardDataModelRequest]]:
-
         if self.source_config.includeDataModels:
             for chart_id in self._get_charts_of_dashboard(dashboard_details):
                 chart_json = self.all_charts.get(chart_id)

metadata/ingestion/source/dashboard/superset/queries.py

@@ -85,5 +85,5 @@ inner join
 on 
 	t.id=tc.table_id  
 where 
-	table_name= %(table_name)s
+	table_name=%(table_name)s
 """

metadata/ingestion/source/dashboard/tableau/metadata.py

@@ -152,6 +152,23 @@ class TableauSource(DashboardServiceSource):
                 include_tags=self.source_config.includeTags,
             )
 
+    def _get_datamodel_sql_query(self, data_model: DataSource) -> Optional[str]:
+        """
+        Method to fetch the custom sql query from the tableau datamodels
+        """
+        try:
+            sql_queries = []
+            for table in data_model.upstreamTables or []:
+                for referenced_query in table.referencedByQueries or []:
+                    sql_queries.append(referenced_query.query)
+            return "\n\n".join(sql_queries) or None
+        except Exception as exc:
+            logger.debug(traceback.format_exc())
+            logger.warning(
+                f"Error processing queries for datamodel [{data_model.id}]: {exc}"
+            )
+        return None
+
     def yield_datamodel(
         self, dashboard_details: TableauDashboard
     ) -> Iterable[Either[CreateDashboardDataModelRequest]]:
@@ -171,6 +188,7 @@ class TableauSource(DashboardServiceSource):
                         dataModelType=DataModelType.TableauDataModel.value,
                         serviceType=DashboardServiceType.Tableau.value,
                         columns=self.get_column_info(data_model),
+                        sql=self._get_datamodel_sql_query(data_model=data_model),
                     )
                     yield Either(right=data_model_request)
                     self.register_record_datamodel(datamodel_request=data_model_request)
@@ -199,7 +217,7 @@ class TableauSource(DashboardServiceSource):
         try:
             dashboard_url = (
                 f"{clean_uri(str(self.config.serviceConnection.__root__.config.hostPort))}"
-                f"/#{urlparse(dashboard_details.webpageUrl).fragment}"
+                f"/#{urlparse(dashboard_details.webpageUrl).fragment}/views"
             )
             dashboard_request = CreateDashboardRequest(
                 name=dashboard_details.id,

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

@@ -54,9 +54,7 @@ class AthenaLakeFormationClient:
             )
         return None
 
-    def get_table_and_column_tags(
-        self, schema_name: str, table_name: str
-    ) -> Optional[LFTags]:
+    def get_table_and_column_tags(self, schema_name: str, table_name: str) -> LFTags:
         """
         Method to call the API and get the table and column tags
         """
@@ -79,4 +77,4 @@ class AthenaLakeFormationClient:
             logger.warning(
                 f"Unable to get lf tags for table resource [{table_name}] due to: {exc}"
             )
-        return None
+        return LFTags()

metadata/ingestion/source/database/azuresql/metadata.py

@@ -77,7 +77,6 @@ class AzuresqlSource(CommonDbSourceService, MultiDBSource):
         yield from self._execute_database_query(AZURE_SQL_GET_DATABASES)
 
     def get_database_names(self) -> Iterable[str]:
-
         if not self.config.serviceConnection.__root__.config.ingestAllDatabases:
             configured_db = self.config.serviceConnection.__root__.config.database
             self.set_inspector(database_name=configured_db)

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

@@ -12,7 +12,7 @@
 """
 Source connection helper
 """
-
+import traceback
 from typing import Any
 
 from pydantic import BaseModel
@@ -26,7 +26,14 @@ from metadata.generated.schema.security.credentials.gcpValues import (
     SingleProjectId,
 )
 from metadata.ingestion.source.connections import get_connection
+from metadata.ingestion.source.database.bigquery.queries import (
+    BIGQUERY_FOREIGN_CONSTRAINTS,
+    BIGQUERY_TABLE_CONSTRAINTS,
+)
 from metadata.utils.bigquery_utils import get_bigquery_client
+from metadata.utils.logger import ingestion_logger
+
+logger = ingestion_logger()
 
 
 class InspectorWrapper(BaseModel):
@@ -64,3 +71,63 @@ def get_inspector_details(
     inspector = inspect(engine)
 
     return InspectorWrapper(client=client, engine=engine, inspector=inspector)
+
+
+def get_pk_constraint(
+    self, connection, table_name, schema=None, **kw
+):  # pylint: disable=unused-argument
+    """
+    This function overrides to get primary key constraint
+    """
+    try:
+        table_constraints = connection.engine.execute(
+            BIGQUERY_TABLE_CONSTRAINTS.format(
+                project_id=connection.engine.url.host,
+                schema_name=schema,
+                table_name=table_name,
+            )
+        )
+        col_name = []
+        for table_constraint in table_constraints:
+            col_name.append(table_constraint.column_name)
+        return {"constrained_columns": tuple(col_name)}
+    except Exception as exc:
+        logger.debug(traceback.format_exc())
+        logger.warning(
+            f"Error while fetching primary key constraint error for table [{schema}.{table_name}]: {exc}"
+        )
+        return {"constrained_columns": []}
+
+
+def get_foreign_keys(
+    self, connection, table_name, schema=None, **kw
+):  # pylint: disable=unused-argument
+    """
+    This function overrides to get foreign key constraint
+    """
+    try:
+        table_constraints = connection.engine.execute(
+            BIGQUERY_FOREIGN_CONSTRAINTS.format(
+                project_id=connection.engine.url.host,
+                schema_name=schema,
+                table_name=table_name,
+            )
+        )
+        col_name = []
+        for table_constraint in table_constraints:
+            col_name.append(
+                {
+                    "name": table_constraint.name,
+                    "referred_schema": table_constraint.referred_schema,
+                    "referred_table": table_constraint.referred_table,
+                    "constrained_columns": [table_constraint.constrained_columns],
+                    "referred_columns": [table_constraint.referred_columns],
+                }
+            )
+        return col_name
+    except Exception as exc:
+        logger.debug(traceback.format_exc())
+        logger.warning(
+            f"Error while fetching foreign key constraint error for table [{schema}.{table_name}]: {exc}"
+        )
+        return []

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

@@ -56,7 +56,11 @@ 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,
@@ -194,6 +198,8 @@ 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(
@@ -245,7 +251,9 @@ class BigquerySource(
             test_connection_fn(
                 self.metadata, inspector_details.engine, self.service_connection
             )
-            if os.environ[GOOGLE_CREDENTIALS]:
+            # 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(
@@ -436,7 +444,8 @@ class BigquerySource(
         inspector_details = get_inspector_details(
             database_name=database_name, service_connection=self.service_connection
         )
-        self.temp_credentials_file_path.append(os.environ[GOOGLE_CREDENTIALS])
+        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