fix: add get_schema_location to metastore loader (#1214)

pinterest · Apr 6, 2023 · 6500610 · 6500610
1 parent 5994a9d
commit 6500610
Show file tree

Hide file tree

Showing 3 changed files with 15 additions and 9 deletions.
diff --git 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
@@ -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
@@ -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)