From ddf4ae261b71274e733f628ac44f216d8763eec1 Mon Sep 17 00:00:00 2001 From: codingwithabhi <63392662+codingwithabhi@users.noreply.github.com> Date: Fri, 21 Jan 2022 17:23:02 +0530 Subject: [PATCH] Bigquery credential configuration updated (#2041) * bigquery-credentials-in-bigquery-json * ui-changes-reverted * credentials-moved-to-options * Update bigquery.json * Update bigquery.py * delete-temp-code-added * delete-temp-code-added * removed-init-from-config * Update datatypes_test.py * Update column_helpers.py * code-formatted Co-authored-by: Abhishek Co-authored-by: = <=> Co-authored-by: Ayush Shah --- ingestion/examples/workflows/bigquery.json | 15 ++++++++++-- .../src/metadata/ingestion/source/bigquery.py | 23 +++++++++++++++---- .../src/metadata/utils/column_helpers.py | 1 + ingestion/tests/unit/datatypes_test.py | 3 +++ 4 files changed, 36 insertions(+), 6 deletions(-) diff --git a/ingestion/examples/workflows/bigquery.json b/ingestion/examples/workflows/bigquery.json index 5da6b1cd1e0..fab27380b2a 100644 --- a/ingestion/examples/workflows/bigquery.json +++ b/ingestion/examples/workflows/bigquery.json @@ -6,8 +6,19 @@ "host_port": "bigquery.googleapis.com", "username": "gcpuser@project_id.iam.gserviceaccount.com", "service_name": "gcp_bigquery", - "options": { - "credentials_path": "examples/creds/bigquery-cred.json" + "options":{ + "credentials":{ + "type": "service_account", + "project_id": "project_id", + "private_key_id": "private_key_id", + "private_key": "", + "client_email": "gcpuser@project_id.iam.gserviceaccount.com", + "client_id": "", + "auth_uri": "https://accounts.google.com/o/oauth2/auth", + "token_uri": "https://oauth2.googleapis.com/token", + "auth_provider_x509_cert_url": "https://www.googleapis.com/oauth2/v1/certs", + "client_x509_cert_url": "" + } }, "table_filter_pattern": { "excludes": [ diff --git a/ingestion/src/metadata/ingestion/source/bigquery.py b/ingestion/src/metadata/ingestion/source/bigquery.py index a7be6f7e5ac..64aa3fdec17 100644 --- a/ingestion/src/metadata/ingestion/source/bigquery.py +++ b/ingestion/src/metadata/ingestion/source/bigquery.py @@ -10,7 +10,8 @@ # limitations under the License. import os -from typing import Optional, Tuple +from typing import Optional, Tuple, Any +import json, tempfile, logging from sqlalchemy_bigquery import _types from sqlalchemy_bigquery._struct import STRUCT @@ -72,11 +73,18 @@ class BigquerySource(SQLSource): def create(cls, config_dict, metadata_config_dict, ctx): config: SQLConnectionConfig = BigQueryConfig.parse_obj(config_dict) metadata_config = MetadataServerConfig.parse_obj(metadata_config_dict) - os.environ["GOOGLE_APPLICATION_CREDENTIALS"] = config.options[ - "credentials_path" - ] + if config.options.get("credentials", None): + cred_path = create_credential_temp_file(config.options.get("credentials")) + os.environ["GOOGLE_APPLICATION_CREDENTIALS"] = cred_path + del config.options["credentials"] + config.options["credentials_path"] = cred_path return cls(config, metadata_config, ctx) + def close(self): + super().close() + if self.config.options["credentials_path"]: + os.unlink(self.config.options["credentials_path"]) + def standardize_schema_table_names( self, schema: str, table: str ) -> Tuple[str, str]: @@ -89,3 +97,10 @@ class BigquerySource(SQLSource): def parse_raw_data_type(self, raw_data_type): return raw_data_type.replace(", ", ",").replace(" ", ":").lower() + + +def create_credential_temp_file(credentials: dict) -> str: + with tempfile.NamedTemporaryFile(delete=False) as fp: + cred_json = json.dumps(credentials, indent=4, separators=(",", ": ")) + fp.write(cred_json.encode()) + return fp.name diff --git a/ingestion/src/metadata/utils/column_helpers.py b/ingestion/src/metadata/utils/column_helpers.py index e4e99414a23..9d6df790433 100644 --- a/ingestion/src/metadata/utils/column_helpers.py +++ b/ingestion/src/metadata/utils/column_helpers.py @@ -127,6 +127,7 @@ _column_string_mapping = { "XML": "BINARY", "XMLTYPE": "BINARY", "CURSOR": "BINARY", + "TIMESTAMP_NTZ": "TIMESTAMP", "TIMESTAMP_LTZ": "TIMESTAMP", "TIMESTAMP_TZ": "TIMESTAMP", } diff --git a/ingestion/tests/unit/datatypes_test.py b/ingestion/tests/unit/datatypes_test.py index 533aa378d50..d033ebdb7fa 100644 --- a/ingestion/tests/unit/datatypes_test.py +++ b/ingestion/tests/unit/datatypes_test.py @@ -90,6 +90,9 @@ SQLTYPES = [ "XML", "XMLTYPE", "YEAR", + "TIMESTAMP_NTZ", + "TIMESTAMP_LTZ", + "TIMESTAMP_TZ", ]