openmetadata-ingestion 1.7.5.0__py3-none-any.whl → 1.8.0.0__py3-none-any.whl

metadata/ingestion/sink/metadata_rest.py

@@ -79,7 +79,7 @@ from metadata.ingestion.models.tests_data import (
     OMetaTestSuiteSample,
 )
 from metadata.ingestion.models.user import OMetaUserProfile
-from metadata.ingestion.ometa.client import APIError
+from metadata.ingestion.ometa.client import APIError, LimitsException
 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
@@ -118,6 +118,7 @@ class MetadataRestSink(Sink):  # pylint: disable=too-many-public-methods
         self.metadata = metadata
         self.role_entities = {}
         self.team_entities = {}
+        self.limit_reached = set()
 
     @classmethod
     def create(
@@ -166,16 +167,30 @@ class MetadataRestSink(Sink):  # pylint: disable=too-many-public-methods
         Send to OM the request creation received as is.
         :param entity_request: Create Entity request
         """
-        created = self.metadata.create_or_update(entity_request)
-        if created:
-            return Either(right=created)
+        if type(entity_request).__name__ in self.limit_reached:
+            # If the limit has been reached, we don't need to try to ingest the entity
+            # Note: We use PatchRequest to update the entity, so updating is not affected by the limit
+            return None
+        try:
+            created = self.metadata.create_or_update(entity_request)
+            if created:
+                return Either(right=created)
 
-        error = f"Failed to ingest {type(entity_request).__name__}"
-        return Either(
-            left=StackTraceError(
-                name=type(entity_request).__name__, error=error, stackTrace=None
+            error = f"Failed to ingest {type(entity_request).__name__}"
+            return Either(
+                left=StackTraceError(
+                    name=type(entity_request).__name__, error=error, stackTrace=None
+                )
+            )
+        except LimitsException as _:
+            self.limit_reached.add(type(entity_request).__name__)
+            return Either(
+                left=StackTraceError(
+                    name=type(entity_request).__name__,
+                    error=f"Limit reached for {type(entity_request).__name__}",
+                    stackTrace=None,
+                )
             )
-        )
 
     @_run_dispatch.register
     def patch_entity(self, record: PatchRequest) -> Either[Entity]:
@@ -329,12 +344,26 @@ class MetadataRestSink(Sink):  # pylint: disable=too-many-public-methods
             team = self.metadata.create_or_update(create_team)
             self.team_entities[team.name.root] = str(team.id.root)
             return team
+        except LimitsException as _:
+            if type(create_team).__name__ in self.limit_reached:
+                # Note: We do not have a way to patch the team,
+                # so we try to put and handle exception
+                return None
+            self.limit_reached.add(type(create_team).__name__)
+            return Either(
+                left=StackTraceError(
+                    name=type(create_team).__name__,
+                    error=f"Limit reached for {type(create_team).__name__}",
+                    stackTrace=None,
+                )
+            )
         except Exception as exc:
             logger.debug(traceback.format_exc())
             logger.error(f"Unexpected error creating team [{create_team}]: {exc}")
 
         return None
 
+    # pylint: disable=too-many-branches
     @_run_dispatch.register
     def write_users(self, record: OMetaUserProfile) -> Either[User]:
         """
@@ -388,8 +417,22 @@ class MetadataRestSink(Sink):  # pylint: disable=too-many-public-methods
         metadata_user = CreateUserRequest(**user_profile)
 
         # Create user
-        user = self.metadata.create_or_update(metadata_user)
-        return Either(right=user)
+        try:
+            user = self.metadata.create_or_update(metadata_user)
+            return Either(right=user)
+        except LimitsException as _:
+            if type(metadata_user).__name__ in self.limit_reached:
+                # Note: We do not have a way to patch the user,
+                # so we try to put and handle exception
+                return None
+            self.limit_reached.add(type(metadata_user).__name__)
+            return Either(
+                left=StackTraceError(
+                    name=type(metadata_user).__name__,
+                    error=f"Limit reached for {type(metadata_user).__name__}",
+                    stackTrace=None,
+                )
+            )
 
     @_run_dispatch.register
     def delete_entity(self, record: DeleteEntity) -> Either[Entity]:

metadata/ingestion/source/api/api_service.py

@@ -39,9 +39,6 @@ from metadata.ingestion.api.delete import delete_entity_from_source
 from metadata.ingestion.api.models import Either
 from metadata.ingestion.api.steps import Source
 from metadata.ingestion.api.topology_runner import TopologyRunnerMixin
-from metadata.ingestion.connections.test_connections import (
-    raise_test_connection_exception,
-)
 from metadata.ingestion.models.delete_entity import DeleteEntity
 from metadata.ingestion.models.topology import (
     NodeStage,
@@ -50,7 +47,7 @@ from metadata.ingestion.models.topology import (
     TopologyNode,
 )
 from metadata.ingestion.ometa.ometa_api import OpenMetadata
-from metadata.ingestion.source.connections import get_connection, get_test_connection_fn
+from metadata.ingestion.source.connections import get_connection, test_connection_common
 from metadata.utils import fqn
 from metadata.utils.logger import ingestion_logger
 
@@ -177,11 +174,9 @@ class ApiServiceSource(TopologyRunnerMixin, Source, ABC):
         """By default, nothing to close"""
 
     def test_connection(self) -> None:
-        test_connection_fn = get_test_connection_fn(self.service_connection)
-        result = test_connection_fn(
+        test_connection_common(
             self.metadata, self.connection_obj, self.service_connection
         )
-        raise_test_connection_exception(result)
 
     def mark_api_collections_as_deleted(self) -> Iterable[Either[DeleteEntity]]:
         """Method to mark the api collection as deleted"""

metadata/ingestion/source/connections.py

@@ -14,11 +14,16 @@ Main entrypoints to create and test connections
 for any source.
 """
 import traceback
-from typing import Any, Callable
+from typing import Any, Callable, Optional, Type
 
 from pydantic import BaseModel
-from sqlalchemy.engine import Engine
 
+from metadata.ingestion.connections.connection import BaseConnection
+from metadata.ingestion.connections.test_connections import (
+    raise_test_connection_exception,
+)
+from metadata.ingestion.ometa.ometa_api import OpenMetadata
+from metadata.utils.class_helper import get_service_type_from_source_type
 from metadata.utils.importer import import_connection_fn
 from metadata.utils.logger import cli_logger
 
@@ -28,10 +33,71 @@ GET_CONNECTION_FN_NAME = "get_connection"
 TEST_CONNECTION_FN_NAME = "test_connection"
 
 
+# TODO: This is a temporary solution to get the connection class from the service spec.
+# Once we migrate all connectors we shouldn't need this.
+def _get_connection_class_from_spec(
+    connection: BaseModel,
+) -> Optional[Type[BaseConnection]]:
+    """
+    Helper method to get the connection class from the connection spec.
+    Returns the connection class if successful, None otherwise.
+    """
+    from metadata.utils.service_spec.service_spec import (  # pylint: disable=import-outside-toplevel
+        BaseSpec,
+        import_connection_class,
+    )
+
+    connection_type = getattr(connection, "type", None)
+    if connection_type:
+        service_type = get_service_type_from_source_type(connection_type.value)
+        try:
+            spec = BaseSpec.get_for_source(service_type, connection_type.value.lower())
+            if getattr(spec, "connection_class", None):
+                connection_class = import_connection_class(
+                    service_type, connection_type.value.lower()
+                )
+                return connection_class
+        except Exception:
+            logger.error(
+                f"Error importing connection class for {connection_type.value}"
+            )
+            logger.debug(traceback.format_exc())
+    return None
+
+
+def _get_connection_fn_from_service_spec(connection: BaseModel) -> Optional[Callable]:
+    """
+    Import the get_connection function from the source, or use ServiceSpec connection_class if defined.
+    """
+    connection_class = _get_connection_class_from_spec(connection)
+    if connection_class:
+
+        def _get_client(conn):
+            return connection_class(conn).get_client()
+
+        return _get_client
+    return None
+
+
+def _get_test_fn_from_service_spec(connection: BaseModel) -> Optional[Callable]:
+    """
+    Import the get_connection function from the source, or use ServiceSpec connection_class if defined.
+    """
+    connection_class = _get_connection_class_from_spec(connection)
+    if connection_class:
+        return connection_class(connection).test_connection
+    return None
+
+
 def get_connection_fn(connection: BaseModel) -> Callable:
     """
-    Import the get_connection function from the source
+    Import the get_connection function from the source, or use ServiceSpec connection_class if defined.
     """
+    # Try ServiceSpec path first
+    connection_fn = _get_connection_fn_from_service_spec(connection)
+    if connection_fn:
+        return connection_fn
+    # Fallback to default
     return import_connection_fn(
         connection=connection, function_name=GET_CONNECTION_FN_NAME
     )
@@ -41,6 +107,10 @@ def get_test_connection_fn(connection: BaseModel) -> Callable:
     """
     Import the test_connection function from the source
     """
+    test_fn = _get_test_fn_from_service_spec(connection)
+    if test_fn:
+        return test_fn
+    # Fallback to default
     return import_connection_fn(
         connection=connection, function_name=TEST_CONNECTION_FN_NAME
     )
@@ -54,15 +124,11 @@ def get_connection(connection: BaseModel) -> Any:
     return get_connection_fn(connection)(connection)
 
 
-def kill_active_connections(engine: Engine):
-    """
-    Method to kill the active connections
-    as well as idle connections in the engine
-    """
+def test_connection_common(metadata: OpenMetadata, connection_obj, service_connection):
+    test_connection_fn = get_test_connection_fn(service_connection)
+    # TODO: Remove this once we migrate all connectors to use the new test connection function
     try:
-        active_conn = engine.pool.checkedout() + engine.pool.checkedin()
-        if active_conn:
-            engine.dispose()
-    except Exception as exc:
-        logger.warning(f"Error Killing the active connections {exc}")
-        logger.debug(traceback.format_exc())
+        result = test_connection_fn(metadata)
+    except TypeError:
+        result = test_connection_fn(metadata, connection_obj, service_connection)
+    raise_test_connection_exception(result)

metadata/ingestion/source/connections_utils.py

@@ -0,0 +1,32 @@
+#  Copyright 2025 Collate
+#  Licensed under the Collate Community License, Version 1.0 (the "License");
+#  you may not use this file except in compliance with the License.
+#  You may obtain a copy of the License at
+#  https://github.com/open-metadata/OpenMetadata/blob/main/ingestion/LICENSE
+#  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.
+
+import traceback
+
+from sqlalchemy.engine import Engine
+
+from metadata.utils.logger import cli_logger
+
+logger = cli_logger()
+
+
+def kill_active_connections(engine: Engine):
+    """
+    Method to kill the active connections
+    as well as idle connections in the engine
+    """
+    try:
+        active_conn = engine.pool.checkedout() + engine.pool.checkedin()
+        if active_conn:
+            engine.dispose()
+    except Exception as exc:
+        logger.warning(f"Error Killing the active connections {exc}")
+        logger.debug(traceback.format_exc())

metadata/ingestion/source/dashboard/dashboard_service.py

@@ -55,9 +55,6 @@ from metadata.ingestion.api.delete import delete_entity_from_source
 from metadata.ingestion.api.models import Either, Entity
 from metadata.ingestion.api.steps import Source
 from metadata.ingestion.api.topology_runner import C, TopologyRunnerMixin
-from metadata.ingestion.connections.test_connections import (
-    raise_test_connection_exception,
-)
 from metadata.ingestion.lineage.sql_lineage import get_column_fqn
 from metadata.ingestion.models.delete_entity import DeleteEntity
 from metadata.ingestion.models.ometa_classification import OMetaTagAndClassification
@@ -70,7 +67,7 @@ from metadata.ingestion.models.topology import (
     TopologyNode,
 )
 from metadata.ingestion.ometa.ometa_api import OpenMetadata
-from metadata.ingestion.source.connections import get_connection, get_test_connection_fn
+from metadata.ingestion.source.connections import get_connection, test_connection_common
 from metadata.utils import fqn
 from metadata.utils.filters import filter_by_dashboard, filter_by_project
 from metadata.utils.logger import ingestion_logger
@@ -595,11 +592,9 @@ class DashboardServiceSource(TopologyRunnerMixin, Source, ABC):
             yield dashboard_details
 
     def test_connection(self) -> None:
-        test_connection_fn = get_test_connection_fn(self.service_connection)
-        result = test_connection_fn(
+        test_connection_common(
             self.metadata, self.connection_obj, self.service_connection
         )
-        raise_test_connection_exception(result)
 
     def prepare(self):
         """By default, nothing to prepare"""

metadata/ingestion/source/dashboard/powerbi/metadata.py

@@ -455,7 +455,7 @@ class PowerbiSource(DashboardServiceSource):
                 measures.append(Column(**parsed_measure.model_dump()))
             except Exception as err:
                 logger.debug(traceback.format_exc())
-                logger.warning(f"Error processing datamodel nested measure: {err}")
+                logger.debug(f"Error processing datamodel nested measure: {err}")
         return measures
 
     def _get_child_columns(self, table: PowerBiTable) -> List[Column]:
@@ -631,20 +631,26 @@ class PowerbiSource(DashboardServiceSource):
                 entity=Dashboard,
                 fqn=report_fqn,
             )
-            datamodel_fqn = fqn.build(
-                self.metadata,
-                entity_type=DashboardDataModel,
-                service_name=self.context.get().dashboard_service,
-                data_model_name=dashboard_details.datasetId,
-            )
-            datamodel_entity = self.metadata.get_by_name(
-                entity=DashboardDataModel,
-                fqn=datamodel_fqn,
-            )
+            if dashboard_details.datasetId:
+                datamodel_fqn = fqn.build(
+                    self.metadata,
+                    entity_type=DashboardDataModel,
+                    service_name=self.context.get().dashboard_service,
+                    data_model_name=dashboard_details.datasetId,
+                )
+                datamodel_entity = self.metadata.get_by_name(
+                    entity=DashboardDataModel,
+                    fqn=datamodel_fqn,
+                )
 
-            if datamodel_entity and report_entity:
-                yield self._get_add_lineage_request(
-                    to_entity=report_entity, from_entity=datamodel_entity
+                if datamodel_entity and report_entity:
+                    yield self._get_add_lineage_request(
+                        to_entity=report_entity, from_entity=datamodel_entity
+                    )
+            else:
+                logger.debug(
+                    f"Skipping datamodel and report lineage for"
+                    f" {dashboard_details.id} as datasetId is not found"
                 )
 
         except Exception as exc:  # pylint: disable=broad-except
@@ -1184,6 +1190,12 @@ class PowerbiSource(DashboardServiceSource):
         Method to process the dashboard owners
         """
         try:
+            if not self.source_config.includeOwners:
+                logger.debug(
+                    f"Skipping owner processing for {dashboard_details.id} "
+                    f"as includeOwners is False"
+                )
+                return None
             owner_ref_list = []  # to assign multiple owners to entity if they exist
             for owner in dashboard_details.users or []:
                 owner_ref = None
@@ -1198,7 +1210,7 @@ class PowerbiSource(DashboardServiceSource):
                     access_right = owner.dashboardUserAccessRight
 
                 if owner.userType != "Member":
-                    logger.warning(
+                    logger.debug(
                         f"User is not a member of {dashboard_details.id}:"
                         f" ({owner.displayName}, {owner.email})"
                     )
@@ -1211,7 +1223,7 @@ class PowerbiSource(DashboardServiceSource):
                         try:
                             owner_email = EmailStr._validate(owner.email)
                         except PydanticCustomError:
-                            logger.warning(f"Invalid email for owner: {owner.email}")
+                            logger.debug(f"Invalid email for owner: {owner.email}")
                             owner_email = None
                         if owner_email:
                             try:
@@ -1219,7 +1231,7 @@ class PowerbiSource(DashboardServiceSource):
                                     owner_email.lower()
                                 )
                             except Exception as err:
-                                logger.warning(
+                                logger.debug(
                                     f"Could not process owner data with email"
                                     f" {owner.email} in {dashboard_details.id}: {err}"
                                 )
@@ -1229,14 +1241,14 @@ class PowerbiSource(DashboardServiceSource):
                                 name=owner.displayName
                             )
                         except Exception as err:
-                            logger.warning(
+                            logger.debug(
                                 f"Could not process owner data with name"
                                 f" {owner.displayName} in {dashboard_details.id}: {err}"
                             )
                     if owner_ref:
                         owner_ref_list.append(owner_ref.root[0])
                 else:
-                    logger.warning(
+                    logger.debug(
                         f"User does not have owner, admin or write access to"
                         f" {dashboard_details.id}: ({owner.displayName}, {owner.email})"
                     )
@@ -1254,7 +1266,7 @@ class PowerbiSource(DashboardServiceSource):
                     if owner_ref and owner_ref.root[0] not in owner_ref_list:
                         owner_ref_list.append(owner_ref.root[0])
                 except Exception as err:
-                    logger.warning(f"Could not fetch owner data due to {err}")
+                    logger.debug(f"Could not fetch current active user due to {err}")
             if len(owner_ref_list) > 0:
                 logger.debug(
                     f"Successfully fetched owners data for {dashboard_details.id}"

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

@@ -24,7 +24,7 @@ from metadata.generated.schema.entity.services.connections.dashboard.supersetCon
     SupersetConnection,
 )
 from metadata.generated.schema.entity.services.connections.database.mysqlConnection import (
-    MysqlConnection,
+    MysqlConnection as MysqlConnectionConfig,
 )
 from metadata.generated.schema.entity.services.connections.database.postgresConnection import (
     PostgresConnection,
@@ -46,16 +46,16 @@ from metadata.ingestion.source.dashboard.superset.queries import (
     FETCH_ALL_CHARTS_TEST,
     FETCH_DASHBOARDS_TEST,
 )
-from metadata.ingestion.source.database.mysql.connection import (
-    get_connection as mysql_get_connection,
-)
+from metadata.ingestion.source.database.mysql.connection import MySQLConnection
 from metadata.ingestion.source.database.postgres.connection import (
     get_connection as pg_get_connection,
 )
 from metadata.utils.constants import THREE_MIN
 
 
-def get_connection(connection: SupersetConnection) -> SupersetAPIClient:
+def get_connection(
+    connection: SupersetConnection,
+) -> Union[SupersetAPIClient, Engine, None]:
     """
     Create connection
     """
@@ -63,8 +63,8 @@ def get_connection(connection: SupersetConnection) -> SupersetAPIClient:
         return SupersetAPIClient(connection)
     if isinstance(connection.connection, PostgresConnection):
         return pg_get_connection(connection=connection.connection)
-    if isinstance(connection.connection, MysqlConnection):
-        return mysql_get_connection(connection=connection.connection)
+    if isinstance(connection.connection, MysqlConnectionConfig):
+        return MySQLConnection(connection.connection).get_client()
     return None
 
 
@@ -89,7 +89,7 @@ def test_connection(
     else:
         test_fn["CheckAccess"] = partial(test_connection_engine_step, client)
         test_fn["GetDashboards"] = partial(test_query, client, FETCH_DASHBOARDS_TEST)
-        if isinstance(service_connection.connection, MysqlConnection):
+        if isinstance(service_connection.connection, MysqlConnectionConfig):
             test_fn["GetCharts"] = partial(
                 test_query, client, FETCH_ALL_CHARTS_TEST.replace('"', "`")
             )

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

@@ -93,7 +93,7 @@ LIMIT 1
 FETCH_COLUMN = """
 select 
 	tc.id, 
-    	tc.table_name ,
+    	t.table_name ,
     	tc.column_name, 
 		tc.table_id, 
     	tc.type,

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

@@ -13,6 +13,7 @@
 Tableau Source Model module
 """
 
+import uuid
 from typing import Dict, List, Optional, Set, Union
 
 from pydantic import BaseModel, ConfigDict, Field, field_validator
@@ -28,9 +29,18 @@ class TableauBaseModel(BaseModel):
 
     model_config = ConfigDict(extra="allow")
 
-    id: str
+    # in case of personal space workbooks, the project id is returned as a UUID
+    id: Union[str, uuid.UUID]
     name: Optional[str] = None
 
+    # pylint: disable=no-self-argument
+    @field_validator("id", mode="before")
+    def coerce_uuid_to_string(cls, value):
+        """Ensure id is always stored as a string internally"""
+        if isinstance(value, uuid.UUID):
+            return str(value)
+        return value
+
     def __hash__(self):
         return hash(self.id)
 
@@ -187,7 +197,7 @@ class TableauDashboard(TableauBaseModel):
     tags: Optional[Set] = []
     webpageUrl: Optional[str] = None
     charts: Optional[List[TableauChart]] = None
-    dataModels: List[DataSource] = []
+    dataModels: Optional[List[DataSource]] = []
     custom_sql_queries: Optional[List[str]] = None
     user_views: Optional[int] = None
 

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

@@ -68,21 +68,25 @@ def get_connection_url(connection: BigQueryConnection) -> str:
             connection.credentials.gcpConfig.projectId, SingleProjectId
         ):
             if not connection.credentials.gcpConfig.projectId.root:
-                return f"{connection.scheme.value}://{connection.credentials.gcpConfig.projectId or ''}"
+                return f"{connection.scheme.value}://{connection.billingProjectId or connection.credentials.gcpConfig.projectId.root or ''}"
             if (
                 not connection.credentials.gcpConfig.privateKey
                 and connection.credentials.gcpConfig.projectId.root
             ):
                 project_id = connection.credentials.gcpConfig.projectId.root
-                os.environ["GOOGLE_CLOUD_PROJECT"] = project_id
-            return f"{connection.scheme.value}://{connection.credentials.gcpConfig.projectId.root}"
+                os.environ["GOOGLE_CLOUD_PROJECT"] = (
+                    connection.billingProjectId or project_id
+                )
+            return f"{connection.scheme.value}://{connection.billingProjectId or connection.credentials.gcpConfig.projectId.root}"
         elif isinstance(connection.credentials.gcpConfig.projectId, MultipleProjectId):
             for project_id in connection.credentials.gcpConfig.projectId.root:
                 if not connection.credentials.gcpConfig.privateKey and project_id:
                     # Setting environment variable based on project id given by user / set in ADC
-                    os.environ["GOOGLE_CLOUD_PROJECT"] = project_id
-                return f"{connection.scheme.value}://{project_id}"
-            return f"{connection.scheme.value}://"
+                    os.environ["GOOGLE_CLOUD_PROJECT"] = (
+                        connection.billingProjectId or project_id
+                    )
+                return f"{connection.scheme.value}://{connection.billingProjectId or project_id}"
+            return f"{connection.scheme.value}://{connection.billingProjectId or ''}"
 
     # If gcpConfig is the JSON key path and projectId is defined, we use it by default
     elif (
@@ -92,13 +96,13 @@ def get_connection_url(connection: BigQueryConnection) -> str:
         if isinstance(  # pylint: disable=no-else-return
             connection.credentials.gcpConfig.projectId, SingleProjectId
         ):
-            return f"{connection.scheme.value}://{connection.credentials.gcpConfig.projectId.root}"
+            return f"{connection.scheme.value}://{connection.billingProjectId or connection.credentials.gcpConfig.projectId.root}"
 
         elif isinstance(connection.credentials.gcpConfig.projectId, MultipleProjectId):
             for project_id in connection.credentials.gcpConfig.projectId.root:
-                return f"{connection.scheme.value}://{project_id}"
+                return f"{connection.scheme.value}://{connection.billingProjectId or project_id}"
 
-    return f"{connection.scheme.value}://"
+    return f"{connection.scheme.value}://{connection.billingProjectId or ''}"
 
 
 def get_connection(connection: BigQueryConnection) -> Engine:

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

@@ -72,7 +72,9 @@ def get_inspector_details(
                 "lifetime"
             ] = new_service_connection.credentials.gcpImpersonateServiceAccount.lifetime
 
-    client = get_bigquery_client(project_id=database_name, **kwargs)
+    client = get_bigquery_client(
+        project_id=new_service_connection.billingProjectId or database_name, **kwargs
+    )
     engine = get_connection(new_service_connection)
     inspector = inspect(engine)
 
@@ -86,15 +88,15 @@ def get_pk_constraint(
     This function overrides to get primary key constraint
     """
     try:
-        constraints = PK_CACHE.get(f"{connection.engine.url.host}.{schema}")
+        project, schema = schema.split(".")
+        constraints = PK_CACHE.get(f"{project}.{schema}")
         if constraints is None:
             constraints = connection.engine.execute(
                 BIGQUERY_TABLE_CONSTRAINTS.format(
-                    project_id=connection.engine.url.host,
-                    schema_name=schema,
+                    project_id=project, schema_name=schema
                 )
             )
-            PK_CACHE[f"{connection.engine.url.host}.{schema}"] = constraints.fetchall()
+            PK_CACHE[f"{project}.{schema}"] = constraints.fetchall()
 
         col_name = []
         table_constraints = [row for row in constraints if row.table_name == table_name]
@@ -116,15 +118,15 @@ def get_foreign_keys(
     This function overrides to get foreign key constraint
     """
     try:
-        constraints = FK_CACHE.get(f"{connection.engine.url.host}.{schema}")
+        project, schema = schema.split(".")
+        constraints = FK_CACHE.get(f"{project}.{schema}")
         if constraints is None:
             constraints = connection.engine.execute(
                 BIGQUERY_FOREIGN_CONSTRAINTS.format(
-                    project_id=connection.engine.url.host,
-                    schema_name=schema,
+                    project_id=project, schema_name=schema
                 )
             )
-            FK_CACHE[f"{connection.engine.url.host}.{schema}"] = constraints.fetchall()
+            FK_CACHE[f"{project}.{schema}"] = constraints.fetchall()
 
         col_name = []
         table_constraints = [row for row in constraints if row.table_name == table_name]