feat(SIP-95): catalogs in SQL Lab and datasets (#28376)

apache · May 8, 2024 · ce668d4 · ce668d4
1 parent 07cd1d8
commit ce668d4
Show file tree

Hide file tree

Showing 71 changed files with 841 additions and 99 deletions.
diff --git a/superset-frontend/packages/superset-ui-core/src/query/types/Query.ts b/superset-frontend/packages/superset-ui-core/src/query/types/Query.ts
@@ -316,6 +316,7 @@ export type Query = {
   link?: string;
   progress: number;
   resultsKey: string | null;
+  catalog?: string | null;
   schema?: string;
   sql: string;
   sqlEditorId: string;

diff --git a/superset-frontend/packages/superset-ui-core/src/ui-overrides/types.ts b/superset-frontend/packages/superset-ui-core/src/ui-overrides/types.ts
@@ -168,6 +168,7 @@ export interface SubMenuProps {
 export interface CustomAutoCompleteArgs {
   queryEditorId: string;
   dbId?: string | number;
+  catalog?: string | null;
   schema?: string;
 }
 

diff --git a/superset-frontend/src/SqlLab/actions/sqlLab.js b/superset-frontend/src/SqlLab/actions/sqlLab.js
@@ -55,6 +55,7 @@ export const REMOVE_QUERY = 'REMOVE_QUERY';
 export const EXPAND_TABLE = 'EXPAND_TABLE';
 export const COLLAPSE_TABLE = 'COLLAPSE_TABLE';
 export const QUERY_EDITOR_SETDB = 'QUERY_EDITOR_SETDB';
+export const QUERY_EDITOR_SET_CATALOG = 'QUERY_EDITOR_SET_CATALOG';
 export const QUERY_EDITOR_SET_SCHEMA = 'QUERY_EDITOR_SET_SCHEMA';
 export const QUERY_EDITOR_SET_TITLE = 'QUERY_EDITOR_SET_TITLE';
 export const QUERY_EDITOR_SET_AUTORUN = 'QUERY_EDITOR_SET_AUTORUN';
@@ -326,6 +327,7 @@ export function runQuery(query) {
       database_id: query.dbId,
       json: true,
       runAsync: query.runAsync,
+      catalog: query.catalog,
       schema: query.schema,
       sql: query.sql,
       sql_editor_id: query.sqlEditorId,
@@ -381,6 +383,7 @@ export function runQueryFromSqlEditor(
       sql: qe.selectedText || qe.sql,
       sqlEditorId: qe.id,
       tab: qe.name,
+      catalog: qe.catalog,
       schema: qe.schema,
       tempTable,
       templateParams: qe.templateParams,
@@ -556,7 +559,7 @@ export function addNewQueryEditor() {
     );
     const dbIds = Object.values(databases).map(database => database.id);
     const firstDbId = dbIds.length > 0 ? Math.min(...dbIds) : undefined;
-    const { dbId, schema, queryLimit, autorun } = {
+    const { dbId, catalog, schema, queryLimit, autorun } = {
       ...queryEditors[0],
       ...activeQueryEditor,
       ...(unsavedQueryEditor.id === activeQueryEditor?.id &&
@@ -578,6 +581,7 @@ export function addNewQueryEditor() {
     return dispatch(
       addQueryEditor({
         dbId: dbId || defaultDbId || firstDbId,
+        catalog: catalog ?? null,
         schema: schema ?? null,
         autorun: autorun ?? false,
         sql: `${warning}SELECT ...`,
@@ -600,6 +604,7 @@ export function cloneQueryToNewTab(query, autorun) {
     const queryEditor = {
       name: t('Copy of %s', sourceQueryEditor.name),
       dbId: query.dbId ? query.dbId : null,
+      catalog: query.catalog ? query.catalog : null,
       schema: query.schema ? query.schema : null,
       autorun,
       sql: query.sql,
@@ -656,6 +661,7 @@ export function setTables(tableSchemas) {
       return {
         dbId: tableSchema.database_id,
         queryEditorId: tableSchema.tab_state_id.toString(),
+        catalog: tableSchema.catalog,
         schema: tableSchema.schema,
         name: tableSchema.table,
         expanded: tableSchema.expanded,
@@ -694,6 +700,7 @@ export function switchQueryEditor(queryEditor, displayLimit) {
             autorun: json.autorun,
             dbId: json.database_id,
             templateParams: json.template_params,
+            catalog: json.catalog,
             schema: json.schema,
             queryLimit: json.query_limit,
             remoteId: json.saved_query?.id,
@@ -797,6 +804,14 @@ export function queryEditorSetDb(queryEditor, dbId) {
   return { type: QUERY_EDITOR_SETDB, queryEditor, dbId };
 }
 
+export function queryEditorSetCatalog(queryEditor, catalog) {
+  return {
+    type: QUERY_EDITOR_SET_CATALOG,
+    queryEditor: queryEditor || {},
+    catalog,
+  };
+}
+
 export function queryEditorSetSchema(queryEditor, schema) {
   return {
     type: QUERY_EDITOR_SET_SCHEMA,
@@ -954,12 +969,13 @@ export function mergeTable(table, query, prepend) {
   return { type: MERGE_TABLE, table, query, prepend };
 }
 
-export function addTable(queryEditor, tableName, schemaName) {
+export function addTable(queryEditor, tableName, catalogName, schemaName) {
   return function (dispatch, getState) {
     const query = getUpToDateQuery(getState(), queryEditor, queryEditor.id);
     const table = {
       dbId: query.dbId,
       queryEditorId: query.id,
+      catalog: catalogName,
       schema: schemaName,
       name: tableName,
     };
@@ -983,12 +999,14 @@ export function runTablePreviewQuery(newTable) {
       sqlLab: { databases },
     } = getState();
     const database = databases[newTable.dbId];
-    const { dbId } = newTable;
+    const { dbId, catalog, schema } = newTable;
 
     if (database && !database.disable_data_preview) {
       const dataPreviewQuery = {
         id: shortid.generate(),
         dbId,
+        catalog,
+        schema,
         sql: newTable.selectStar,
         tableName: newTable.name,
         sqlEditorId: null,
@@ -1003,6 +1021,7 @@ export function runTablePreviewQuery(newTable) {
             {
               id: newTable.id,
               dbId: newTable.dbId,
+              catalog: newTable.catalog,
               schema: newTable.schema,
               name: newTable.name,
               queryEditorId: newTable.queryEditorId,
@@ -1180,6 +1199,7 @@ export function popStoredQuery(urlId) {
           addQueryEditor({
             name: json.name ? json.name : t('Shared query'),
             dbId: json.dbId ? parseInt(json.dbId, 10) : null,
+            catalog: json.catalog ? json.catalog : null,
             schema: json.schema ? json.schema : null,
             autorun: json.autorun ? json.autorun : false,
             sql: json.sql ? json.sql : 'SELECT ...',
@@ -1215,6 +1235,7 @@ export function popQuery(queryId) {
         const queryData = json.result;
         const queryEditorProps = {
           dbId: queryData.database.id,
+          catalog: queryData.catalog,
           schema: queryData.schema,
           sql: queryData.sql,
           name: t('Copy of %s', queryData.tab_name),
@@ -1268,12 +1289,13 @@ export function createDatasourceFailed(err) {
 export function createDatasource(vizOptions) {
   return dispatch => {
     dispatch(createDatasourceStarted());
-    const { dbId, schema, datasourceName, sql } = vizOptions;
+    const { dbId, catalog, schema, datasourceName, sql } = vizOptions;
     return SupersetClient.post({
       endpoint: '/api/v1/dataset/',
       headers: { 'Content-Type': 'application/json' },
       body: JSON.stringify({
         database: dbId,
+        catalog,
         schema,
         sql,
         table_name: datasourceName,

diff --git a/superset-frontend/src/SqlLab/actions/sqlLab.test.js b/superset-frontend/src/SqlLab/actions/sqlLab.test.js
@@ -419,6 +419,7 @@ describe('async actions', () => {
           queryEditor: {
             name: 'Copy of Dummy query editor',
             dbId: 1,
+            catalog: query.catalog,
             schema: query.schema,
             autorun: true,
             sql: 'SELECT * FROM something',
@@ -481,6 +482,7 @@ describe('async actions', () => {
               sql: expect.stringContaining('SELECT ...'),
               name: `Untitled Query 7`,
               dbId: defaultQueryEditor.dbId,
+              catalog: defaultQueryEditor.catalog,
               schema: defaultQueryEditor.schema,
               autorun: false,
               queryLimit:
@@ -607,6 +609,24 @@ describe('async actions', () => {
       });
     });
 
+    describe('queryEditorSetCatalog', () => {
+      it('updates the tab state in the backend', () => {
+        expect.assertions(1);
+
+        const catalog = 'public';
+        const store = mockStore({});
+        const expectedActions = [
+          {
+            type: actions.QUERY_EDITOR_SET_CATALOG,
+            queryEditor,
+            catalog,
+          },
+        ];
+        store.dispatch(actions.queryEditorSetCatalog(queryEditor, catalog));
+        expect(store.getActions()).toEqual(expectedActions);
+      });
+    });
+
     describe('queryEditorSetSchema', () => {
       it('updates the tab state in the backend', () => {
         expect.assertions(1);
@@ -747,6 +767,7 @@ describe('async actions', () => {
     describe('addTable', () => {
       it('dispatches table state from unsaved change', () => {
         const tableName = 'table';
+        const catalogName = null;
         const schemaName = 'schema';
         const expectedDbId = 473892;
         const store = mockStore({
@@ -759,12 +780,18 @@ describe('async actions', () => {
             },
           },
         });
-        const request = actions.addTable(query, tableName, schemaName);
+        const request = actions.addTable(
+          query,
+          tableName,
+          catalogName,
+          schemaName,
+        );
         request(store.dispatch, store.getState);
         expect(store.getActions()[0]).toEqual(
           expect.objectContaining({
             table: expect.objectContaining({
               name: tableName,
+              catalog: catalogName,
               schema: schemaName,
               dbId: expectedDbId,
             }),
@@ -811,6 +838,7 @@ describe('async actions', () => {
         });
 
         const tableName = 'table';
+        const catalogName = null;
         const schemaName = 'schema';
         const store = mockStore({
           ...initialState,
@@ -829,6 +857,7 @@ describe('async actions', () => {
         const request = actions.runTablePreviewQuery({
           dbId: 1,
           name: tableName,
+          catalog: catalogName,
           schema: schemaName,
         });
         return request(store.dispatch, store.getState).then(() => {

diff --git a/superset-frontend/src/SqlLab/components/AceEditorWrapper/index.tsx b/superset-frontend/src/SqlLab/components/AceEditorWrapper/index.tsx
@@ -74,6 +74,7 @@ const AceEditorWrapper = ({
     'id',
     'dbId',
     'sql',
+    'catalog',
     'schema',
     'templateParams',
     'cursorPosition',
@@ -161,6 +162,7 @@ const AceEditorWrapper = ({
 
   const { data: annotations } = useAnnotations({
     dbId: queryEditor.dbId,
+    catalog: queryEditor.catalog,
     schema: queryEditor.schema,
     sql: currentSql,
     templateParams: queryEditor.templateParams,
@@ -170,6 +172,7 @@ const AceEditorWrapper = ({
     {
       queryEditorId,
       dbId: queryEditor.dbId,
+      catalog: queryEditor.catalog,
       schema: queryEditor.schema,
     },
     !autocomplete,

diff --git a/superset-frontend/src/SqlLab/components/AceEditorWrapper/useKeywords.test.ts b/superset-frontend/src/SqlLab/components/AceEditorWrapper/useKeywords.test.ts
@@ -189,7 +189,12 @@ test('returns column keywords among selected tables', async () => {
     storeWithSqlLab.dispatch(
       tableApiUtil.upsertQueryData(
         'tableMetadata',
-        { dbId: expectDbId, schema: expectSchema, table: expectTable },
+        {
+          dbId: expectDbId,
+          catalog: null,
+          schema: expectSchema,
+          table: expectTable,
+        },
         {
           name: expectTable,
           columns: [
@@ -205,7 +210,12 @@ test('returns column keywords among selected tables', async () => {
     storeWithSqlLab.dispatch(
       tableApiUtil.upsertQueryData(
         'tableMetadata',
-        { dbId: expectDbId, schema: expectSchema, table: unexpectedTable },
+        {
+          dbId: expectDbId,
+          catalog: null,
+          schema: expectSchema,
+          table: unexpectedTable,
+        },
         {
           name: unexpectedTable,
           columns: [
@@ -227,6 +237,7 @@ test('returns column keywords among selected tables', async () => {
       useKeywords({
         queryEditorId: expectQueryEditorId,
         dbId: expectDbId,
+        catalog: null,
         schema: expectSchema,
       }),
     {

diff --git a/superset-frontend/src/SqlLab/components/AceEditorWrapper/useKeywords.ts b/superset-frontend/src/SqlLab/components/AceEditorWrapper/useKeywords.ts
@@ -42,6 +42,7 @@ import { SqlLabRootState } from 'src/SqlLab/types';
 type Params = {
   queryEditorId: string | number;
   dbId?: string | number;
+  catalog?: string | null;
   schema?: string;
 };
 
@@ -58,7 +59,7 @@ const getHelperText = (value: string) =>
 const extensionsRegistry = getExtensionsRegistry();
 
 export function useKeywords(
-  { queryEditorId, dbId, schema }: Params,
+  { queryEditorId, dbId, catalog, schema }: Params,
   skip = false,
 ) {
   const useCustomKeywords = extensionsRegistry.get(
@@ -68,6 +69,7 @@ export function useKeywords(
   const customKeywords = useCustomKeywords?.({
     queryEditorId: String(queryEditorId),
     dbId,
+    catalog,
     schema,
   });
   const dispatch = useDispatch();
@@ -78,13 +80,15 @@ export function useKeywords(
   const { data: schemaOptions } = useSchemasQueryState(
     {
       dbId,
+      catalog: catalog || undefined,
       forceRefresh: false,
     },
     { skip: skipFetch || !dbId },
   );
   const { data: tableData } = useTablesQueryState(
     {
       dbId,
+      catalog,
       schema,
       forceRefresh: false,
     },
@@ -125,6 +129,7 @@ export function useKeywords(
           dbId && schema
             ? {
                 dbId,
+                catalog,
                 schema,
                 table,
               }
@@ -137,7 +142,7 @@ export function useKeywords(
         });
     });
     return [...columns];
-  }, [dbId, schema, apiState, tablesForColumnMetadata]);
+  }, [dbId, catalog, schema, apiState, tablesForColumnMetadata]);
 
   const insertMatch = useEffectEvent((editor: Editor, data: any) => {
     if (data.meta === 'table') {