From 0ee37623a376cc0be9ac44e3b2c2a84100b32624 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Luis=20S=C3=A1nchez?= Date: Tue, 25 Nov 2025 17:06:46 -0300 Subject: [PATCH 1/4] fix(SqlLab): enhance SQL formatting with Jinja template support. --- .../src/SqlLab/actions/sqlLab.js | 19 +- .../src/SqlLab/actions/sqlLab.test.js | 195 +++++++++++++++++- superset/sqllab/api.py | 29 ++- superset/sqllab/schemas.py | 8 + tests/integration_tests/sql_lab/api_tests.py | 19 ++ 5 files changed, 266 insertions(+), 4 deletions(-) diff --git a/superset-frontend/src/SqlLab/actions/sqlLab.js b/superset-frontend/src/SqlLab/actions/sqlLab.js index 644d3e33463d..bf8e2fe0a63b 100644 --- a/superset-frontend/src/SqlLab/actions/sqlLab.js +++ b/superset-frontend/src/SqlLab/actions/sqlLab.js @@ -914,11 +914,26 @@ export function queryEditorSetAndSaveSql(targetQueryEditor, sql, queryId) { export function formatQuery(queryEditor) { return function (dispatch, getState) { - const { sql } = getUpToDateQuery(getState(), queryEditor); + const qe = getUpToDateQuery(getState(), queryEditor); + const { sql, dbId, templateParams } = qe; + const body = { sql }; + + // Include database_id and template_params if available for Jinja processing + if (dbId) { + body.database_id = dbId; + } + if (templateParams) { + // Send templateParams as a JSON string to match the backend schema + body.template_params = + typeof templateParams === 'string' + ? templateParams + : JSON.stringify(templateParams); + } + return SupersetClient.post({ endpoint: `/api/v1/sqllab/format_sql/`, // TODO (betodealmeida): pass engine as a parameter for better formatting - body: JSON.stringify({ sql }), + body: JSON.stringify(body), headers: { 'Content-Type': 'application/json' }, }).then(({ json }) => { dispatch(queryEditorSetSql(queryEditor, json.result)); diff --git a/superset-frontend/src/SqlLab/actions/sqlLab.test.js b/superset-frontend/src/SqlLab/actions/sqlLab.test.js index 0e9644b69cf3..a90355b87621 100644 --- a/superset-frontend/src/SqlLab/actions/sqlLab.test.js +++ b/superset-frontend/src/SqlLab/actions/sqlLab.test.js @@ -170,7 +170,16 @@ describe('async actions', () => { describe('formatQuery', () => { const formatQueryEndpoint = 'glob:*/api/v1/sqllab/format_sql/'; const expectedSql = 'SELECT 1'; - fetchMock.post(formatQueryEndpoint, { result: expectedSql }); + + beforeEach(() => { + fetchMock.post( + formatQueryEndpoint, + { result: expectedSql }, + { + overwriteRoutes: true, + }, + ); + }); test('posts to the correct url', async () => { const store = mockStore(initialState); @@ -181,6 +190,190 @@ describe('async actions', () => { expect(store.getActions()[0].type).toBe(actions.QUERY_EDITOR_SET_SQL); expect(store.getActions()[0].sql).toBe(expectedSql); }); + + test('sends only sql in request body when no dbId or templateParams', async () => { + const queryEditorWithoutExtras = { + ...defaultQueryEditor, + sql: 'SELECT * FROM table', + dbId: null, + templateParams: null, + }; + const state = { + sqlLab: { + queryEditors: [queryEditorWithoutExtras], + unsavedQueryEditor: {}, + }, + }; + const store = mockStore(state); + + store.dispatch(actions.formatQuery(queryEditorWithoutExtras)); + + await waitFor(() => + expect(fetchMock.calls(formatQueryEndpoint)).toHaveLength(1), + ); + + const call = fetchMock.calls(formatQueryEndpoint)[0]; + const body = JSON.parse(call[1].body); + + expect(body).toEqual({ sql: 'SELECT * FROM table' }); + expect(body.database_id).toBeUndefined(); + expect(body.template_params).toBeUndefined(); + }); + + test('includes database_id in request when dbId is provided', async () => { + const queryEditorWithDb = { + ...defaultQueryEditor, + sql: 'SELECT * FROM table', + dbId: 5, + templateParams: null, + }; + const state = { + sqlLab: { + queryEditors: [queryEditorWithDb], + unsavedQueryEditor: {}, + }, + }; + const store = mockStore(state); + + store.dispatch(actions.formatQuery(queryEditorWithDb)); + + await waitFor(() => + expect(fetchMock.calls(formatQueryEndpoint)).toHaveLength(1), + ); + + const call = fetchMock.calls(formatQueryEndpoint)[0]; + const body = JSON.parse(call[1].body); + + expect(body).toEqual({ + sql: 'SELECT * FROM table', + database_id: 5, + }); + }); + + test('includes template_params as string when provided as string', async () => { + const queryEditorWithTemplateString = { + ...defaultQueryEditor, + sql: 'SELECT * FROM table WHERE id = {{ user_id }}', + dbId: 5, + templateParams: '{"user_id": 123}', + }; + const state = { + sqlLab: { + queryEditors: [queryEditorWithTemplateString], + unsavedQueryEditor: {}, + }, + }; + const store = mockStore(state); + + store.dispatch(actions.formatQuery(queryEditorWithTemplateString)); + + await waitFor(() => + expect(fetchMock.calls(formatQueryEndpoint)).toHaveLength(1), + ); + + const call = fetchMock.calls(formatQueryEndpoint)[0]; + const body = JSON.parse(call[1].body); + + expect(body).toEqual({ + sql: 'SELECT * FROM table WHERE id = {{ user_id }}', + database_id: 5, + template_params: '{"user_id": 123}', + }); + }); + + test('stringifies template_params when provided as object', async () => { + const queryEditorWithTemplateObject = { + ...defaultQueryEditor, + sql: 'SELECT * FROM table WHERE id = {{ user_id }}', + dbId: 5, + templateParams: { user_id: 123, status: 'active' }, + }; + const state = { + sqlLab: { + queryEditors: [queryEditorWithTemplateObject], + unsavedQueryEditor: {}, + }, + }; + const store = mockStore(state); + + store.dispatch(actions.formatQuery(queryEditorWithTemplateObject)); + + await waitFor(() => + expect(fetchMock.calls(formatQueryEndpoint)).toHaveLength(1), + ); + + const call = fetchMock.calls(formatQueryEndpoint)[0]; + const body = JSON.parse(call[1].body); + + expect(body).toEqual({ + sql: 'SELECT * FROM table WHERE id = {{ user_id }}', + database_id: 5, + template_params: '{"user_id":123,"status":"active"}', + }); + }); + + test('dispatches QUERY_EDITOR_SET_SQL with formatted result', async () => { + const formattedSql = 'SELECT\n *\nFROM\n table'; + fetchMock.post( + formatQueryEndpoint, + { result: formattedSql }, + { + overwriteRoutes: true, + }, + ); + + const queryEditorToFormat = { + ...defaultQueryEditor, + sql: 'SELECT * FROM table', + }; + const state = { + sqlLab: { + queryEditors: [queryEditorToFormat], + unsavedQueryEditor: {}, + }, + }; + const store = mockStore(state); + + await store.dispatch(actions.formatQuery(queryEditorToFormat)); + + const dispatchedActions = store.getActions(); + expect(dispatchedActions).toHaveLength(1); + expect(dispatchedActions[0].type).toBe(actions.QUERY_EDITOR_SET_SQL); + expect(dispatchedActions[0].sql).toBe(formattedSql); + }); + + test('uses up-to-date query editor state from store', async () => { + const outdatedQueryEditor = { + ...defaultQueryEditor, + sql: 'OLD SQL', + dbId: 1, + }; + const upToDateQueryEditor = { + ...defaultQueryEditor, + sql: 'SELECT * FROM updated_table', + dbId: 10, + }; + const state = { + sqlLab: { + queryEditors: [upToDateQueryEditor], + unsavedQueryEditor: {}, + }, + }; + const store = mockStore(state); + + // Pass outdated query editor, but expect the function to use the up-to-date one from store + store.dispatch(actions.formatQuery(outdatedQueryEditor)); + + await waitFor(() => + expect(fetchMock.calls(formatQueryEndpoint)).toHaveLength(1), + ); + + const call = fetchMock.calls(formatQueryEndpoint)[0]; + const body = JSON.parse(call[1].body); + + expect(body.sql).toBe('SELECT * FROM updated_table'); + expect(body.database_id).toBe(10); + }); }); // eslint-disable-next-line no-restricted-globals -- TODO: Migrate from describe blocks diff --git a/superset/sqllab/api.py b/superset/sqllab/api.py index d3cd123346bb..82812248a361 100644 --- a/superset/sqllab/api.py +++ b/superset/sqllab/api.py @@ -234,7 +234,34 @@ def format_sql(self) -> FlaskResponse: """ try: model = self.format_model_schema.load(request.json) - result = SQLScript(model["sql"], model.get("engine")).format() + sql = model["sql"] + template_params = model.get("template_params") + database_id = model.get("database_id") + + # Process Jinja templates if template_params and database_id are provided + if template_params and database_id is not None: + database = DatabaseDAO.find_by_id(database_id) + if database: + try: + template_params = ( + json.loads(template_params) + if isinstance(template_params, str) + else template_params + ) + except json.JSONDecodeError: + logger.warning( + "Invalid template parameter %s. Skipping processing", + str(template_params), + ) + template_params = {} + + if template_params: + template_processor = get_template_processor(database=database) + sql = template_processor.process_template( + sql, **template_params + ) + + result = SQLScript(sql, model.get("engine")).format() return self.response(200, result=result) except ValidationError as error: return self.response_400(message=error.messages) diff --git a/superset/sqllab/schemas.py b/superset/sqllab/schemas.py index ce86a990c5b7..3c2413f8d3ba 100644 --- a/superset/sqllab/schemas.py +++ b/superset/sqllab/schemas.py @@ -48,6 +48,14 @@ class EstimateQueryCostSchema(Schema): class FormatQueryPayloadSchema(Schema): sql = fields.String(required=True) engine = fields.String(required=False, allow_none=True) + database_id = fields.Integer( + required=False, allow_none=True, metadata={"description": "The database id"} + ) + template_params = fields.String( + required=False, + allow_none=True, + metadata={"description": "The SQL query template params as JSON string"}, + ) class ExecutePayloadSchema(Schema): diff --git a/tests/integration_tests/sql_lab/api_tests.py b/tests/integration_tests/sql_lab/api_tests.py index 97a7eab9e80d..9454e7c9b104 100644 --- a/tests/integration_tests/sql_lab/api_tests.py +++ b/tests/integration_tests/sql_lab/api_tests.py @@ -288,6 +288,25 @@ def test_format_sql_request(self): self.assertDictEqual(resp_data, success_resp) # noqa: PT009 assert rv.status_code == 200 + def test_format_sql_request_with_jinja(self): + self.login(ADMIN_USERNAME) + example_db = get_example_database() + + data = { + "sql": "select * from {{tbl}}", + "database_id": example_db.id, + "template_params": json.dumps({"tbl": '"Vehicle Sales"'}), + } + rv = self.client.post( + "/api/v1/sqllab/format_sql/", + json=data, + ) + resp_data = json.loads(rv.data.decode("utf-8")) + # Verify that Jinja template was processed before formatting + assert "{{tbl}}" not in resp_data["result"] + assert '"Vehicle Sales"' in resp_data["result"] + assert rv.status_code == 200 + @mock.patch("superset.commands.sql_lab.results.results_backend_use_msgpack", False) def test_execute_required_params(self): self.login(ADMIN_USERNAME) From 528e93a701ce44aadbc09aa64501bdfd69343f71 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Luis=20S=C3=A1nchez?= Date: Wed, 26 Nov 2025 12:59:29 -0300 Subject: [PATCH 2/4] refactor(SqlLab): streamline query formatting by directly destructuring state values, removing solved TODO. --- superset-frontend/src/SqlLab/actions/sqlLab.js | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/superset-frontend/src/SqlLab/actions/sqlLab.js b/superset-frontend/src/SqlLab/actions/sqlLab.js index bf8e2fe0a63b..230bb956aa7e 100644 --- a/superset-frontend/src/SqlLab/actions/sqlLab.js +++ b/superset-frontend/src/SqlLab/actions/sqlLab.js @@ -914,8 +914,7 @@ export function queryEditorSetAndSaveSql(targetQueryEditor, sql, queryId) { export function formatQuery(queryEditor) { return function (dispatch, getState) { - const qe = getUpToDateQuery(getState(), queryEditor); - const { sql, dbId, templateParams } = qe; + const { sql, dbId, templateParams } = getUpToDateQuery(getState(), queryEditor); const body = { sql }; // Include database_id and template_params if available for Jinja processing @@ -932,7 +931,6 @@ export function formatQuery(queryEditor) { return SupersetClient.post({ endpoint: `/api/v1/sqllab/format_sql/`, - // TODO (betodealmeida): pass engine as a parameter for better formatting body: JSON.stringify(body), headers: { 'Content-Type': 'application/json' }, }).then(({ json }) => { From 2d6ded9abef1487b06dece90252d1b6d72eb353a Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Luis=20S=C3=A1nchez?= Date: Fri, 28 Nov 2025 15:59:09 -0300 Subject: [PATCH 3/4] refactor(SqlLab): optimize template parameter processing in SQL formatting --- superset/sqllab/api.py | 14 +++++++------- 1 file changed, 7 insertions(+), 7 deletions(-) diff --git a/superset/sqllab/api.py b/superset/sqllab/api.py index 82812248a361..f81031a601c7 100644 --- a/superset/sqllab/api.py +++ b/superset/sqllab/api.py @@ -248,18 +248,18 @@ def format_sql(self) -> FlaskResponse: if isinstance(template_params, str) else template_params ) + if template_params: + template_processor = get_template_processor( + database=database + ) + sql = template_processor.process_template( + sql, **template_params + ) except json.JSONDecodeError: logger.warning( "Invalid template parameter %s. Skipping processing", str(template_params), ) - template_params = {} - - if template_params: - template_processor = get_template_processor(database=database) - sql = template_processor.process_template( - sql, **template_params - ) result = SQLScript(sql, model.get("engine")).format() return self.response(200, result=result) From fb5243b7ab92d90f3b5a7a5adf6722726548ffde Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Luis=20S=C3=A1nchez?= Date: Fri, 28 Nov 2025 16:12:11 -0300 Subject: [PATCH 4/4] chore(SqlLab): prettier applied due to pre-commit hooks failing. --- superset-frontend/src/SqlLab/actions/sqlLab.js | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/superset-frontend/src/SqlLab/actions/sqlLab.js b/superset-frontend/src/SqlLab/actions/sqlLab.js index 230bb956aa7e..fd23cb57538b 100644 --- a/superset-frontend/src/SqlLab/actions/sqlLab.js +++ b/superset-frontend/src/SqlLab/actions/sqlLab.js @@ -914,7 +914,10 @@ export function queryEditorSetAndSaveSql(targetQueryEditor, sql, queryId) { export function formatQuery(queryEditor) { return function (dispatch, getState) { - const { sql, dbId, templateParams } = getUpToDateQuery(getState(), queryEditor); + const { sql, dbId, templateParams } = getUpToDateQuery( + getState(), + queryEditor, + ); const body = { sql }; // Include database_id and template_params if available for Jinja processing