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

apache · May 7, 2024 · b8fc10d · b8fc10d
1 parent e90246f
commit b8fc10d
Show file tree

Hide file tree

Showing 51 changed files with 573 additions and 64 deletions.
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/components/SaveDatasetModal/index.tsx b/superset-frontend/src/SqlLab/components/SaveDatasetModal/index.tsx
@@ -77,6 +77,7 @@ export interface ISaveableDatasource {
   dbId: number;
   sql: string;
   templateParams?: string | object | null;
+  catalog?: string | null;
   schema?: string | null;
   database?: Database;
 }
@@ -292,6 +293,7 @@ export const SaveDatasetModal = ({
       createDatasource({
         sql: datasource.sql,
         dbId: datasource.dbId || datasource?.database?.id,
+        catalog: datasource?.catalog,
         schema: datasource?.schema,
         templateParams,
         datasourceName: datasetName,

diff --git a/superset-frontend/src/SqlLab/components/SaveQuery/index.tsx b/superset-frontend/src/SqlLab/components/SaveQuery/index.tsx
@@ -48,7 +48,7 @@ export type QueryPayload = {
   description?: string;
   id?: string;
   remoteId?: number;
-} & Pick<QueryEditor, 'dbId' | 'schema' | 'sql'>;
+} & Pick<QueryEditor, 'dbId' | 'catalog' | 'schema' | 'sql'>;
 
 const Styles = styled.span`
   span[role='img'] {
@@ -78,6 +78,7 @@ const SaveQuery = ({
     'dbId',
     'latestQueryId',
     'queryLimit',
+    'catalog',
     'schema',
     'selectedText',
     'sql',
@@ -115,6 +116,7 @@ const SaveQuery = ({
     description,
     dbId: query.dbId ?? 0,
     sql: query.sql,
+    catalog: query.catalog,
     schema: query.schema,
     templateParams: query.templateParams,
     remoteId: query?.remoteId || undefined,

diff --git a/superset-frontend/src/SqlLab/components/SqlEditorLeftBar/index.tsx b/superset-frontend/src/SqlLab/components/SqlEditorLeftBar/index.tsx
@@ -34,6 +34,7 @@ import {
   removeTables,
   collapseTable,
   expandTable,
+  queryEditorSetCatalog,
   queryEditorSetSchema,
   setDatabases,
   addDangerToast,
@@ -115,13 +116,17 @@ const SqlEditorLeftBar = ({
     shallowEqual,
   );
   const dispatch = useDispatch();
-  const queryEditor = useQueryEditor(queryEditorId, ['dbId', 'schema']);
+  const queryEditor = useQueryEditor(queryEditorId, [
+    'dbId',
+    'catalog',
+    'schema',
+  ]);
 
   const [emptyResultsWithSearch, setEmptyResultsWithSearch] = useState(false);
   const [userSelectedDb, setUserSelected] = useState<DatabaseObject | null>(
     null,
   );
-  const { schema } = queryEditor;
+  const { catalog, schema } = queryEditor;
 
   useEffect(() => {
     const bool = querystring.parse(window.location.search).db;
@@ -138,9 +143,9 @@ const SqlEditorLeftBar = ({
     }
   }, [database]);
 
-  const onEmptyResults = (searchText?: string) => {
+  const onEmptyResults = useCallback((searchText?: string) => {
     setEmptyResultsWithSearch(!!searchText);
-  };
+  }, []);
 
   const onDbChange = ({ id: dbId }: { id: number }) => {
     setEmptyState?.(false);
@@ -152,7 +157,11 @@ const SqlEditorLeftBar = ({
     [tables],
   );
 
-  const onTablesChange = (tableNames: string[], schemaName: string) => {
+  const onTablesChange = (
+    tableNames: string[],
+    catalogName: string | null,
+    schemaName: string,
+  ) => {
     if (!schemaName) {
       return;
     }
@@ -169,7 +178,7 @@ const SqlEditorLeftBar = ({
     });
 
     tablesToAdd.forEach(tableName => {
-      dispatch(addTable(queryEditor, tableName, schemaName));
+      dispatch(addTable(queryEditor, tableName, catalogName, schemaName));
     });
 
     dispatch(removeTables(currentTables));
@@ -210,6 +219,15 @@ const SqlEditorLeftBar = ({
   const shouldShowReset = window.location.search === '?reset=1';
   const tableMetaDataHeight = height - 130; // 130 is the height of the selects above
 
+  const handleCatalogChange = useCallback(
+    (catalog: string | null) => {
+      if (queryEditor) {
+        dispatch(queryEditorSetCatalog(queryEditor, catalog));
+      }
+    },
+    [dispatch, queryEditor],
+  );
+
   const handleSchemaChange = useCallback(
     (schema: string) => {
       if (queryEditor) {
@@ -246,9 +264,11 @@ const SqlEditorLeftBar = ({
         getDbList={handleDbList}
         handleError={handleError}
         onDbChange={onDbChange}
+        onCatalogChange={handleCatalogChange}
+        catalog={catalog}
         onSchemaChange={handleSchemaChange}
-        onTableSelectChange={onTablesChange}
         schema={schema}
+        onTableSelectChange={onTablesChange}
         tableValue={selectedTableNames}
         sqlLabMode
       />

diff --git a/superset-frontend/src/SqlLab/components/TabbedSqlEditors/index.tsx b/superset-frontend/src/SqlLab/components/TabbedSqlEditors/index.tsx
@@ -111,6 +111,7 @@ class TabbedSqlEditors extends React.PureComponent<TabbedSqlEditorsProps> {
       queryId,
       dbid,
       dbname,
+      catalog,
       schema,
       autorun,
       new: isNewQuery,
@@ -149,6 +150,7 @@ class TabbedSqlEditors extends React.PureComponent<TabbedSqlEditorsProps> {
         const newQueryEditor = {
           name,
           dbId: databaseId,
+          catalog,
           schema,
           autorun,
           sql,

diff --git a/superset-frontend/src/SqlLab/components/TableElement/index.tsx b/superset-frontend/src/SqlLab/components/TableElement/index.tsx
@@ -101,7 +101,7 @@ const StyledCollapsePanel = styled(Collapse.Panel)`
 `;
 
 const TableElement = ({ table, ...props }: TableElementProps) => {
-  const { dbId, schema, name, expanded } = table;
+  const { dbId, catalog, schema, name, expanded } = table;
   const theme = useTheme();
   const dispatch = useDispatch();
   const {
@@ -112,6 +112,7 @@ const TableElement = ({ table, ...props }: TableElementProps) => {
   } = useTableMetadataQuery(
     {
       dbId,
+      catalog,
       schema,
       table: name,
     },
@@ -125,6 +126,7 @@ const TableElement = ({ table, ...props }: TableElementProps) => {
   } = useTableExtendedMetadataQuery(
     {
       dbId,
+      catalog,
       schema,
       table: name,
     },

diff --git a/superset-frontend/src/SqlLab/reducers/getInitialState.ts b/superset-frontend/src/SqlLab/reducers/getInitialState.ts
@@ -89,6 +89,7 @@ export default function getInitialState({
         autorun: Boolean(activeTab.autorun),
         templateParams: activeTab.template_params || undefined,
         dbId: activeTab.database_id,
+        catalog: activeTab.catalog,
         schema: activeTab.schema,
         queryLimit: activeTab.query_limit,
         hideLeftBar: activeTab.hide_left_bar,
@@ -121,6 +122,7 @@ export default function getInitialState({
         const table = {
           dbId: tableSchema.database_id,
           queryEditorId: tableSchema.tab_state_id.toString(),
+          catalog: tableSchema.catalog,
           schema: tableSchema.schema,
           name: tableSchema.table,
           expanded: tableSchema.expanded,

diff --git a/superset-frontend/src/SqlLab/reducers/sqlLab.js b/superset-frontend/src/SqlLab/reducers/sqlLab.js
@@ -109,6 +109,7 @@ export default function sqlLabReducer(state = {}, action) {
         remoteId: progenitor.remoteId,
         name: t('Copy of %s', progenitor.name),
         dbId: action.query.dbId ? action.query.dbId : null,
+        catalog: action.query.catalog ? action.query.catalog : null,
         schema: action.query.schema ? action.query.schema : null,
         autorun: true,
         sql: action.query.sql,
@@ -180,6 +181,7 @@ export default function sqlLabReducer(state = {}, action) {
         if (
           xt.dbId === at.dbId &&
           xt.queryEditorId === at.queryEditorId &&
+          xt.catalog === at.catalog &&
           xt.schema === at.schema &&
           xt.name === at.name
         ) {
@@ -410,6 +412,8 @@ export default function sqlLabReducer(state = {}, action) {
       return state;
     },
     [actions.LOAD_QUERY_EDITOR]() {
+      console.log(state);
+      console.log(state.unsavedQueryEditor);
       const mergeUnsavedState = alterInArr(
         state,
         'queryEditors',
@@ -503,6 +507,18 @@ export default function sqlLabReducer(state = {}, action) {
         ),
       };
     },
+    [actions.QUERY_EDITOR_SET_CATALOG]() {
+      return {
+        ...state,
+        ...alterUnsavedQueryEditorState(
+          state,
+          {
+            catalog: action.catalog,
+          },
+          action.queryEditor.id,
+        ),
+      };
+    },
     [actions.QUERY_EDITOR_SET_SCHEMA]() {
       return {
         ...state,

diff --git a/superset-frontend/src/SqlLab/types.ts b/superset-frontend/src/SqlLab/types.ts
@@ -50,6 +50,7 @@ export interface QueryEditor {
   dbId?: number;
   name: string;
   title?: string; // keep it optional for backward compatibility
+  catalog?: string | null;
   schema?: string;
   autorun: boolean;
   sql: string;
@@ -81,6 +82,7 @@ export type UnsavedQueryEditor = Partial<QueryEditor>;
 export interface Table {
   id: string;
   dbId: number;
+  catalog: string | null;
   schema: string;
   name: string;
   queryEditorId: QueryEditor['id'];