From 3111f91a0697a88a7f7731559634ba3f8239c607 Mon Sep 17 00:00:00 2001 From: vanshika18 <52087403+vanshika18@users.noreply.github.com> Date: Thu, 27 Jul 2023 23:24:55 +0530 Subject: [PATCH] Support delete temporary credentials files of Google Cloud (#12590) Co-authored-by: Vanshika Kabra --- .../ingestion/source/database/bigquery/metadata.py | 9 +++++++-- ingestion/src/metadata/utils/credentials.py | 10 +++++++++- 2 files changed, 16 insertions(+), 3 deletions(-) diff --git a/ingestion/src/metadata/ingestion/source/database/bigquery/metadata.py b/ingestion/src/metadata/ingestion/source/database/bigquery/metadata.py index 6f367d88ed2..18978f56640 100644 --- a/ingestion/src/metadata/ingestion/source/database/bigquery/metadata.py +++ b/ingestion/src/metadata/ingestion/source/database/bigquery/metadata.py @@ -54,13 +54,13 @@ from metadata.generated.schema.security.credentials.gcpValues import ( 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.database.bigquery.queries import ( BIGQUERY_SCHEMA_DESCRIPTION, ) from metadata.ingestion.source.database.column_type_parser import create_sqlalchemy_type from metadata.ingestion.source.database.common_db_source import CommonDbSourceService from metadata.utils import fqn +from metadata.utils.credentials import GOOGLE_CREDENTIALS from metadata.utils.filters import filter_by_database from metadata.utils.logger import ingestion_logger from metadata.utils.sqlalchemy_utils import is_complex_type @@ -317,7 +317,6 @@ class BigquerySource(CommonDbSourceService): self.service_connection.credentials.gcpConfig.projectId = SingleProjectId( __root__=database_name ) - self.engine = get_connection(self.service_connection) self.inspector = inspect(self.engine) def get_database_names(self) -> Iterable[str]: @@ -430,6 +429,12 @@ class BigquerySource(CommonDbSourceService): 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] def get_source_url( self, diff --git a/ingestion/src/metadata/utils/credentials.py b/ingestion/src/metadata/utils/credentials.py index e29d17620fd..525797e9686 100644 --- a/ingestion/src/metadata/utils/credentials.py +++ b/ingestion/src/metadata/utils/credentials.py @@ -67,8 +67,15 @@ def create_credential_tmp_file(credentials: dict) -> str: with tempfile.NamedTemporaryFile(delete=False) as temp_file: cred_json = json.dumps(credentials, indent=4, separators=(",", ": ")) temp_file.write(cred_json.encode()) + # Get the path of the temporary file + temp_file_path = temp_file.name - return temp_file.name + # The temporary file will be automatically closed when exiting the "with" block, + # but we can explicitly close it here to free up resources immediately. + temp_file.close() + + # Return the path of the temporary file + return temp_file_path def build_google_credentials_dict(gcp_values: GcpCredentialsValues) -> Dict[str, str]: @@ -121,6 +128,7 @@ def set_google_credentials(gcp_credentials: GCPCredentials) -> None: "Overriding default projectid, using the current environment permissions authenticated via gcloud SDK." ) return + credentials_dict = build_google_credentials_dict(gcp_credentials.gcpConfig) tmp_credentials_file = create_credential_tmp_file(credentials=credentials_dict) os.environ[GOOGLE_CREDENTIALS] = tmp_credentials_file