feat(sqllab): Format sql (#25344)

This commit is contained in:
JUST.in DO IT 2023-11-03 09:38:11 -07:00 committed by GitHub
parent d619078d25
commit 24a2213445
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
8 changed files with 119 additions and 4 deletions

View File

@ -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);

View File

@ -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);

View File

@ -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({

View File

@ -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,

View File

@ -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<Props> = ({
[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<Props> = ({
}),
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<Props> = ({
? t('Schedule the query periodically')
: t('You must run the query successfully first');
return (
<Menu css={{ width: theme.gridUnit * 44 }}>
<Menu css={{ width: theme.gridUnit * 50 }}>
<Menu.Item css={{ display: 'flex', justifyContent: 'space-between' }}>
{' '}
<span>{t('Autocomplete')}</span>{' '}
@ -622,6 +635,7 @@ const SqlEditor: React.FC<Props> = ({
/>
</Menu.Item>
)}
<Menu.Item onClick={formatCurrentQuery}>{t('Format SQL')}</Menu.Item>
{!isEmpty(scheduledQueriesConf) && (
<Menu.Item>
<ScheduleQueryButton

View File

@ -19,7 +19,9 @@ from typing import Any, cast, Optional
from urllib import parse
import simplejson as json
import sqlparse
from flask import request, Response
from flask_appbuilder import permission_name
from flask_appbuilder.api import expose, protect, rison, safe
from flask_appbuilder.models.sqla.interface import SQLAInterface
from marshmallow import ValidationError
@ -46,6 +48,7 @@ from superset.sqllab.query_render import SqlQueryRenderImpl
from superset.sqllab.schemas import (
EstimateQueryCostSchema,
ExecutePayloadSchema,
FormatQueryPayloadSchema,
QueryExecutionResponseSchema,
sql_lab_get_results_schema,
SQLLabBootstrapSchema,
@ -78,6 +81,7 @@ class SqlLabRestApi(BaseSupersetApi):
estimate_model_schema = EstimateQueryCostSchema()
execute_model_schema = ExecutePayloadSchema()
format_model_schema = FormatQueryPayloadSchema()
apispec_parameter_schemas = {
"sql_lab_get_results_schema": sql_lab_get_results_schema,
@ -185,6 +189,52 @@ class SqlLabRestApi(BaseSupersetApi):
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/<string:client_id>/")
@protect()
@statsd_metrics

View File

@ -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)

View File

@ -223,6 +223,19 @@ class TestSqlLabApi(SupersetTestCase):
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()