feat(SIP-95): new endpoint for table metadata (#28122)

apache · Apr 25, 2024 · 6cf681d · john-bodley · May 4, 2024 · 6cf681d
1 parent 52f8734
commit 6cf681d
Show file tree

Hide file tree

Showing 71 changed files with 1,051 additions and 516 deletions.
diff --git a/superset-frontend/src/SqlLab/actions/sqlLab.test.js b/superset-frontend/src/SqlLab/actions/sqlLab.test.js
@@ -508,10 +508,11 @@ describe('async actions', () => {
     fetchMock.delete(updateTableSchemaEndpoint, {});
     fetchMock.post(updateTableSchemaEndpoint, JSON.stringify({ id: 1 }));
 
-    const getTableMetadataEndpoint = 'glob:**/api/v1/database/*/table/*/*/';
+    const getTableMetadataEndpoint =
+      'glob:**/api/v1/database/*/table_metadata/*';
     fetchMock.get(getTableMetadataEndpoint, {});
     const getExtraTableMetadataEndpoint =
-      'glob:**/api/v1/database/*/table_metadata/extra/';
+      'glob:**/api/v1/database/*/table_metadata/extra/*';
     fetchMock.get(getExtraTableMetadataEndpoint, {});
 
     let isFeatureEnabledMock;

diff --git a/superset-frontend/src/SqlLab/components/SqlEditorLeftBar/SqlEditorLeftBar.test.tsx b/superset-frontend/src/SqlLab/components/SqlEditorLeftBar/SqlEditorLeftBar.test.tsx
@@ -61,13 +61,13 @@ beforeEach(() => {
       },
     ],
   });
-  fetchMock.get('glob:*/api/v1/database/*/table/*/*', {
+  fetchMock.get('glob:*/api/v1/database/*/table_metadata/*', {
     status: 200,
     body: {
       columns: table.columns,
     },
   });
-  fetchMock.get('glob:*/api/v1/database/*/table_metadata/extra/', {
+  fetchMock.get('glob:*/api/v1/database/*/table_metadata/extra/*', {
     status: 200,
     body: {},
   });

diff --git a/superset-frontend/src/SqlLab/components/TableElement/TableElement.test.tsx b/superset-frontend/src/SqlLab/components/TableElement/TableElement.test.tsx
@@ -47,9 +47,10 @@ jest.mock(
       <div data-test="mock-column-element">{column.name}</div>
     ),
 );
-const getTableMetadataEndpoint = 'glob:**/api/v1/database/*/table/*/*/';
+const getTableMetadataEndpoint =
+  /\/api\/v1\/database\/\d+\/table_metadata\/(?:\?.*)?$/;
 const getExtraTableMetadataEndpoint =
-  'glob:**/api/v1/database/*/table_metadata/extra/*';
+  /\/api\/v1\/database\/\d+\/table_metadata\/extra\/(?:\?.*)?$/;
 const updateTableSchemaEndpoint = 'glob:*/tableschemaview/*/expanded';
 
 beforeEach(() => {

diff --git a/superset-frontend/src/features/datasets/AddDataset/DatasetPanel/index.tsx b/superset-frontend/src/features/datasets/AddDataset/DatasetPanel/index.tsx
@@ -74,7 +74,9 @@ const DatasetPanelWrapper = ({
     const { dbId, tableName, schema } = props;
     setLoading(true);
     setHasColumns?.(false);
-    const path = `/api/v1/database/${dbId}/table/${tableName}/${schema}/`;
+    const path = schema
+      ? `/api/v1/database/${dbId}/table_metadata/?name=${tableName}&schema=${schema}`
+      : `/api/v1/database/${dbId}/table_metadata/?name=${tableName}`;
     try {
       const response = await SupersetClient.get({
         endpoint: path,

diff --git a/superset-frontend/src/hooks/apiResources/tables.ts b/superset-frontend/src/hooks/apiResources/tables.ts
@@ -114,9 +114,9 @@ const tableApi = api.injectEndpoints({
     }),
     tableMetadata: builder.query<TableMetaData, FetchTableMetadataQueryParams>({
       query: ({ dbId, schema, table }) => ({
-        endpoint: `/api/v1/database/${dbId}/table/${encodeURIComponent(
-          table,
-        )}/${encodeURIComponent(schema)}/`,
+        endpoint: schema
+          ? `/api/v1/database/${dbId}/table_metadata/?name=${table}&schema=${schema}`
+          : `/api/v1/database/${dbId}/table_metadata/?name=${table}`,
         transformResponse: ({ json }: TableMetadataReponse) => json,
       }),
     }),

diff --git a/superset/commands/database/tables.py b/superset/commands/database/tables.py
@@ -51,6 +51,7 @@ def run(self) -> dict[str, Any]:
                 datasource_names=sorted(
                     DatasourceName(*datasource_name)
                     for datasource_name in self._model.get_all_table_names_in_schema(
+                        catalog=None,
                         schema=self._schema_name,
                         force=self._force,
                         cache=self._model.table_cache_enabled,
@@ -65,6 +66,7 @@ def run(self) -> dict[str, Any]:
                 datasource_names=sorted(
                     DatasourceName(*datasource_name)
                     for datasource_name in self._model.get_all_view_names_in_schema(
+                        catalog=None,
                         schema=self._schema_name,
                         force=self._force,
                         cache=self._model.table_cache_enabled,

diff --git a/superset/commands/database/validate_sql.py b/superset/commands/database/validate_sql.py
@@ -61,11 +61,12 @@ def run(self) -> list[dict[str, Any]]:
             raise ValidatorSQLUnexpectedError()
         sql = self._properties["sql"]
         schema = self._properties.get("schema")
+        catalog = self._properties.get("catalog")
         try:
             timeout = current_app.config["SQLLAB_VALIDATION_TIMEOUT"]
             timeout_msg = f"The query exceeded the {timeout} seconds timeout."
             with utils.timeout(seconds=timeout, error_message=timeout_msg):
-                errors = self._validator.validate(sql, schema, self._model)
+                errors = self._validator.validate(sql, catalog, schema, self._model)
             return [err.to_dict() for err in errors]
         except Exception as ex:
             logger.exception(ex)

diff --git a/superset/commands/dataset/create.py b/superset/commands/dataset/create.py
@@ -34,6 +34,7 @@
 from superset.daos.exceptions import DAOCreateFailedError
 from superset.exceptions import SupersetSecurityException
 from superset.extensions import db, security_manager
+from superset.sql_parse import Table
 
 logger = logging.getLogger(__name__)
 
@@ -61,12 +62,15 @@ def validate(self) -> None:
         exceptions: list[ValidationError] = []
         database_id = self._properties["database"]
         table_name = self._properties["table_name"]
-        schema = self._properties.get("schema", None)
-        sql = self._properties.get("sql", None)
+        schema = self._properties.get("schema")
+        catalog = self._properties.get("catalog")
+        sql = self._properties.get("sql")
         owner_ids: Optional[list[int]] = self._properties.get("owners")
 
+        table = Table(table_name, schema, catalog)
+
         # Validate uniqueness
-        if not DatasetDAO.validate_uniqueness(database_id, schema, table_name):
+        if not DatasetDAO.validate_uniqueness(database_id, table):
             exceptions.append(DatasetExistsValidationError(table_name))
 
         # Validate/Populate database
@@ -80,7 +84,7 @@ def validate(self) -> None:
         if (
             database
             and not sql
-            and not DatasetDAO.validate_table_exists(database, table_name, schema)
+            and not DatasetDAO.validate_table_exists(database, table)
         ):
             exceptions.append(TableNotFoundValidationError(table_name))
 
@@ -89,6 +93,7 @@ def validate(self) -> None:
                 security_manager.raise_for_access(
                     database=database,
                     sql=sql,
+                    catalog=catalog,
                     schema=schema,
                 )
             except SupersetSecurityException as ex:

diff --git a/superset/commands/dataset/importers/v1/utils.py b/superset/commands/dataset/importers/v1/utils.py
@@ -32,6 +32,7 @@
 from superset.commands.exceptions import ImportFailedError
 from superset.connectors.sqla.models import SqlaTable
 from superset.models.core import Database
+from superset.sql_parse import Table
 from superset.utils.core import get_user
 
 logger = logging.getLogger(__name__)
@@ -164,7 +165,9 @@ def import_dataset(
         db.session.flush()
 
     try:
-        table_exists = dataset.database.has_table_by_name(dataset.table_name)
+        table_exists = dataset.database.has_table(
+            Table(dataset.table_name, dataset.schema),
+        )
     except Exception:  # pylint: disable=broad-except
         # MySQL doesn't play nice with GSheets table names
         logger.warning(
@@ -217,7 +220,10 @@ def load_data(data_uri: str, dataset: SqlaTable, database: Database) -> None:
         )
     else:
         logger.warning("Loading data outside the import transaction")
-        with database.get_sqla_engine() as engine:
+        with database.get_sqla_engine(
+            catalog=dataset.catalog,
+            schema=dataset.schema,
+        ) as engine:
             df.to_sql(
                 dataset.table_name,
                 con=engine,

diff --git a/superset/commands/dataset/update.py b/superset/commands/dataset/update.py
@@ -41,6 +41,7 @@
 from superset.daos.dataset import DatasetDAO
 from superset.daos.exceptions import DAOUpdateFailedError
 from superset.exceptions import SupersetSecurityException
+from superset.sql_parse import Table
 
 logger = logging.getLogger(__name__)
 
@@ -90,9 +91,8 @@ def validate(self) -> None:
         # Validate uniqueness
         if not DatasetDAO.validate_update_uniqueness(
             self._model.database_id,
-            self._model.schema,
+            Table(table_name, self._model.schema, self._model.catalog),
             self._model_id,
-            table_name,
         ):
             exceptions.append(DatasetExistsValidationError(table_name))
         # Validate/Populate database not allowed to change

diff --git a/superset/commands/sql_lab/estimate.py b/superset/commands/sql_lab/estimate.py
@@ -17,7 +17,7 @@
 from __future__ import annotations
 
 import logging
-from typing import Any
+from typing import Any, TypedDict
 
 from flask_babel import gettext as __
 
@@ -27,7 +27,6 @@
 from superset.exceptions import SupersetErrorException, SupersetTimeoutException
 from superset.jinja_context import get_template_processor
 from superset.models.core import Database
-from superset.sqllab.schemas import EstimateQueryCostSchema
 from superset.utils import core as utils
 
 config = app.config
@@ -37,18 +36,28 @@
 logger = logging.getLogger(__name__)
 
 
+class EstimateQueryCostType(TypedDict):
+    database_id: int
+    sql: str
+    template_params: dict[str, Any]
+    catalog: str | None
+    schema: str | None
+
+
 class QueryEstimationCommand(BaseCommand):
     _database_id: int
     _sql: str
     _template_params: dict[str, Any]
     _schema: str
     _database: Database
+    _catalog: str | None
 
-    def __init__(self, params: EstimateQueryCostSchema) -> None:
-        self._database_id = params.get("database_id")
+    def __init__(self, params: EstimateQueryCostType) -> None:
+        self._database_id = params["database_id"]
         self._sql = params.get("sql", "")
         self._template_params = params.get("template_params", {})
-        self._schema = params.get("schema", "")
+        self._schema = params.get("schema") or ""
+        self._catalog = params.get("catalog")
 
     def validate(self) -> None:
         self._database = db.session.query(Database).get(self._database_id)
@@ -77,7 +86,11 @@ def run(
         try:
             with utils.timeout(seconds=timeout, error_message=timeout_msg):
                 cost = self._database.db_engine_spec.estimate_query_cost(
-                    self._database, self._schema, sql, utils.QuerySource.SQL_LAB
+                    self._database,
+                    self._catalog,
+                    self._schema,
+                    sql,
+                    utils.QuerySource.SQL_LAB,
                 )
         except SupersetTimeoutException as ex:
             logger.exception(ex)