From 24a22134453f61f46fd079821ecb887b72ed4bcc Mon Sep 17 00:00:00 2001 From: "JUST.in DO IT" Date: Fri, 3 Nov 2023 09:38:11 -0700 Subject: [PATCH] feat(sqllab): Format sql (#25344) --- .../src/SqlLab/actions/sqlLab.js | 13 +++++ .../src/SqlLab/actions/sqlLab.test.js | 17 +++++++ .../components/AceEditorWrapper/index.tsx | 6 ++- .../KeyboardShortcutButton/index.tsx | 2 + .../src/SqlLab/components/SqlEditor/index.tsx | 18 ++++++- superset/sqllab/api.py | 50 +++++++++++++++++++ superset/sqllab/schemas.py | 4 ++ tests/integration_tests/sql_lab/api_tests.py | 13 +++++ 8 files changed, 119 insertions(+), 4 deletions(-) diff --git a/superset-frontend/src/SqlLab/actions/sqlLab.js b/superset-frontend/src/SqlLab/actions/sqlLab.js index 57e01088a831d..44b4307a1906e 100644 --- a/superset-frontend/src/SqlLab/actions/sqlLab.js +++ b/superset-frontend/src/SqlLab/actions/sqlLab.js @@ -1033,6 +1033,19 @@ export function queryEditorSetSql(queryEditor, sql) { return { type: QUERY_EDITOR_SET_SQL, queryEditor, sql }; } +export function formatQuery(queryEditor) { + return function (dispatch, getState) { + const { sql } = getUpToDateQuery(getState(), queryEditor); + return SupersetClient.post({ + endpoint: `/api/v1/sqllab/format_sql/`, + body: JSON.stringify({ sql }), + headers: { 'Content-Type': 'application/json' }, + }).then(({ json }) => { + dispatch(queryEditorSetSql(queryEditor, json.result)); + }); + }; +} + export function queryEditorSetAndSaveSql(targetQueryEditor, sql) { return function (dispatch, getState) { const queryEditor = getUpToDateQuery(getState(), targetQueryEditor); diff --git a/superset-frontend/src/SqlLab/actions/sqlLab.test.js b/superset-frontend/src/SqlLab/actions/sqlLab.test.js index 25f80aa1c386a..dbf4e8a5c55f1 100644 --- a/superset-frontend/src/SqlLab/actions/sqlLab.test.js +++ b/superset-frontend/src/SqlLab/actions/sqlLab.test.js @@ -22,6 +22,7 @@ import fetchMock from 'fetch-mock'; import configureMockStore from 'redux-mock-store'; import thunk from 'redux-thunk'; import shortid from 'shortid'; +import { waitFor } from '@testing-library/react'; import * as uiCore from '@superset-ui/core'; import * as actions from 'src/SqlLab/actions/sqlLab'; import { LOG_EVENT } from 'src/logger/actions'; @@ -127,6 +128,22 @@ describe('async actions', () => { }); }); + describe('formatQuery', () => { + const formatQueryEndpoint = 'glob:*/api/v1/sqllab/format_sql/'; + const expectedSql = 'SELECT 1'; + fetchMock.post(formatQueryEndpoint, { result: expectedSql }); + + test('posts to the correct url', async () => { + const store = mockStore(initialState); + store.dispatch(actions.formatQuery(query, queryId)); + await waitFor(() => + expect(fetchMock.calls(formatQueryEndpoint)).toHaveLength(1), + ); + expect(store.getActions()[0].type).toBe(actions.QUERY_EDITOR_SET_SQL); + expect(store.getActions()[0].sql).toBe(expectedSql); + }); + }); + describe('fetchQueryResults', () => { const makeRequest = () => { const request = actions.fetchQueryResults(query); diff --git a/superset-frontend/src/SqlLab/components/AceEditorWrapper/index.tsx b/superset-frontend/src/SqlLab/components/AceEditorWrapper/index.tsx index 3b6b2e67b6249..8529f49d4f8f2 100644 --- a/superset-frontend/src/SqlLab/components/AceEditorWrapper/index.tsx +++ b/superset-frontend/src/SqlLab/components/AceEditorWrapper/index.tsx @@ -146,8 +146,10 @@ const AceEditorWrapper = ({ }; const onChangeText = (text: string) => { - setSql(text); - onChange(text); + if (text !== sql) { + setSql(text); + onChange(text); + } }; const { data: annotations } = useAnnotations({ diff --git a/superset-frontend/src/SqlLab/components/KeyboardShortcutButton/index.tsx b/superset-frontend/src/SqlLab/components/KeyboardShortcutButton/index.tsx index 306e69e7608c1..0599793d862d3 100644 --- a/superset-frontend/src/SqlLab/components/KeyboardShortcutButton/index.tsx +++ b/superset-frontend/src/SqlLab/components/KeyboardShortcutButton/index.tsx @@ -37,6 +37,7 @@ export enum KeyboardShortcut { CMD_OPT_F = 'cmd+opt+f', CTRL_F = 'ctrl+f', CTRL_H = 'ctrl+h', + CTRL_SHIFT_F = 'ctrl+shift+f', } export const KEY_MAP = { @@ -49,6 +50,7 @@ export const KEY_MAP = { [KeyboardShortcut.CTRL_Q]: userOS === 'Windows' ? t('New tab') : undefined, [KeyboardShortcut.CTRL_T]: userOS !== 'Windows' ? t('New tab') : undefined, [KeyboardShortcut.CTRL_P]: t('Previous Line'), + [KeyboardShortcut.CTRL_SHIFT_F]: t('Format SQL'), // default ace editor shortcuts [KeyboardShortcut.CMD_F]: userOS === 'MacOS' ? t('Find') : undefined, [KeyboardShortcut.CTRL_F]: userOS !== 'MacOS' ? t('Find') : undefined, diff --git a/superset-frontend/src/SqlLab/components/SqlEditor/index.tsx b/superset-frontend/src/SqlLab/components/SqlEditor/index.tsx index 83bb80d997a28..609cb917b6f20 100644 --- a/superset-frontend/src/SqlLab/components/SqlEditor/index.tsx +++ b/superset-frontend/src/SqlLab/components/SqlEditor/index.tsx @@ -72,6 +72,7 @@ import { scheduleQuery, setActiveSouthPaneTab, updateSavedQuery, + formatQuery, } from 'src/SqlLab/actions/sqlLab'; import { STATE_TYPE_MAP, @@ -305,6 +306,10 @@ const SqlEditor: React.FC = ({ [ctas, database, defaultQueryLimit, dispatch, queryEditor], ); + const formatCurrentQuery = useCallback(() => { + dispatch(formatQuery(queryEditor)); + }, [dispatch, queryEditor]); + const stopQuery = useCallback(() => { if (latestQuery && ['running', 'pending'].indexOf(latestQuery.state) >= 0) { dispatch(postStopQuery(latestQuery)); @@ -384,8 +389,16 @@ const SqlEditor: React.FC = ({ }), func: stopQuery, }, + { + name: 'formatQuery', + key: KeyboardShortcut.CTRL_SHIFT_F, + descr: KEY_MAP[KeyboardShortcut.CTRL_SHIFT_F], + func: () => { + formatCurrentQuery(); + }, + }, ]; - }, [dispatch, queryEditor.sql, startQuery, stopQuery]); + }, [dispatch, queryEditor.sql, startQuery, stopQuery, formatCurrentQuery]); const hotkeys = useMemo(() => { // Get all hotkeys including ace editor hotkeys @@ -602,7 +615,7 @@ const SqlEditor: React.FC = ({ ? t('Schedule the query periodically') : t('You must run the query successfully first'); return ( - + {' '} {t('Autocomplete')}{' '} @@ -622,6 +635,7 @@ const SqlEditor: React.FC = ({ /> )} + {t('Format SQL')} {!isEmpty(scheduledQueriesConf) && ( Response: result = command.run() return self.response(200, result=result) + @expose("/format_sql/", methods=("POST",)) + @statsd_metrics + @protect() + @permission_name("read") + @event_logger.log_this_with_context( + action=lambda self, *args, **kwargs: f"{self.__class__.__name__}" f".format", + log_to_statsd=False, + ) + def format_sql(self) -> FlaskResponse: + """Format the SQL query. + --- + post: + summary: Format SQL code + requestBody: + description: SQL query + required: true + content: + application/json: + schema: + $ref: '#/components/schemas/FormatQueryPayloadSchema' + responses: + 200: + description: Format SQL result + content: + application/json: + schema: + type: object + properties: + result: + type: string + 400: + $ref: '#/components/responses/400' + 401: + $ref: '#/components/responses/401' + 403: + $ref: '#/components/responses/403' + 500: + $ref: '#/components/responses/500' + """ + try: + model = self.format_model_schema.load(request.json) + result = sqlparse.format(model["sql"], reindent=True, keyword_case="upper") + return self.response(200, result=result) + except ValidationError as error: + return self.response_400(message=error.messages) + @expose("/export//") @protect() @statsd_metrics diff --git a/superset/sqllab/schemas.py b/superset/sqllab/schemas.py index 46ee773222fc2..66f90a6e920a0 100644 --- a/superset/sqllab/schemas.py +++ b/superset/sqllab/schemas.py @@ -42,6 +42,10 @@ class EstimateQueryCostSchema(Schema): ) +class FormatQueryPayloadSchema(Schema): + sql = fields.String(required=True) + + class ExecutePayloadSchema(Schema): database_id = fields.Integer(required=True) sql = fields.String(required=True) diff --git a/tests/integration_tests/sql_lab/api_tests.py b/tests/integration_tests/sql_lab/api_tests.py index 6441033b6ca63..49dd4ea32e7f4 100644 --- a/tests/integration_tests/sql_lab/api_tests.py +++ b/tests/integration_tests/sql_lab/api_tests.py @@ -223,6 +223,19 @@ def test_estimate_valid_request(self): self.assertDictEqual(resp_data, success_resp) self.assertEqual(rv.status_code, 200) + def test_format_sql_request(self): + self.login() + + data = {"sql": "select 1 from my_table"} + rv = self.client.post( + "/api/v1/sqllab/format_sql/", + json=data, + ) + success_resp = {"result": "SELECT 1\nFROM my_table"} + resp_data = json.loads(rv.data.decode("utf-8")) + self.assertDictEqual(resp_data, success_resp) + self.assertEqual(rv.status_code, 200) + @mock.patch("superset.sqllab.commands.results.results_backend_use_msgpack", False) def test_execute_required_params(self): self.login()