From 26eae8d5ce443f8e49cf85507aea6e9159695485 Mon Sep 17 00:00:00 2001 From: "J.C. Zhong" Date: Thu, 6 Apr 2023 17:45:30 +0000 Subject: [PATCH] fix: add get_schema_location to metastore loader --- .../lib/metastore/base_metastore_loader.py | 4 ++++ .../metastore/loaders/hive_metastore_loader.py | 3 +++ .../lib/table_upload/exporter/s3_exporter.py | 17 ++++++++--------- 3 files changed, 15 insertions(+), 9 deletions(-) diff --git a/querybook/server/lib/metastore/base_metastore_loader.py b/querybook/server/lib/metastore/base_metastore_loader.py index 02b88bc90..7cb6053f5 100644 --- a/querybook/server/lib/metastore/base_metastore_loader.py +++ b/querybook/server/lib/metastore/base_metastore_loader.py @@ -538,6 +538,10 @@ def get_data_element(self, data_element_name: str) -> Optional[DataElementTuple] """Override this to get data element by name""" pass + def get_schema_location(self, schema_name: str) -> str: + """Get schema location, used by table uploader""" + pass + @abstractclassmethod def get_metastore_params_template(self) -> AllFormField: """Override this to get the form field required for the metastore diff --git a/querybook/server/lib/metastore/loaders/hive_metastore_loader.py b/querybook/server/lib/metastore/loaders/hive_metastore_loader.py index c9e78eecd..7d0376e29 100644 --- a/querybook/server/lib/metastore/loaders/hive_metastore_loader.py +++ b/querybook/server/lib/metastore/loaders/hive_metastore_loader.py @@ -98,6 +98,9 @@ def get_partitions( self.hmc, schema_name, table_name, conditions ) + def get_schema_location(self, schema_name: str) -> str: + return self.hmc.get_database(schema_name).locationUri + def _get_hmc(self, metastore_dict): return HiveMetastoreClient( hmss_ro_addrs=metastore_dict["metastore_params"]["hms_connection"] diff --git a/querybook/server/lib/table_upload/exporter/s3_exporter.py b/querybook/server/lib/table_upload/exporter/s3_exporter.py index eca1efccc..886305bce 100644 --- a/querybook/server/lib/table_upload/exporter/s3_exporter.py +++ b/querybook/server/lib/table_upload/exporter/s3_exporter.py @@ -75,18 +75,17 @@ def destination_s3_root(self, session=None) -> str: return sanitize_s3_url_with_trailing_slash(s3_path) + schema_name + "/" if self._exporter_config.get("use_schema_location", False): - # Defer import since this is only needed for this option - from lib.metastore.loaders.hive_metastore_loader import HMSMetastoreLoader - query_engine = get_query_engine_by_id(self._engine_id, session=session) - metastore: HMSMetastoreLoader = get_metastore_loader( - query_engine.metastore_id, session=session + metastore = get_metastore_loader(query_engine.metastore_id, session=session) + + if metastore is None: + raise Exception("Invalid metastore") + + schema_location_uri = metastore.get_schema_location( + self._table_config["schema_name"] ) - if metastore is None or not isinstance(metastore, HMSMetastoreLoader): + if not schema_location_uri: raise Exception("Invalid metastore to use use_schema_location option") - schema_location_uri = metastore.hmc.get_database( - self._table_config["schema_name"] - ).locationUri return sanitize_s3_url_with_trailing_slash(schema_location_uri)