fix(SqlLab): enhance SQL formatting with Jinja template support. (#36277)

superset-frontend/src/SqlLab/actions/sqlLab.js

@@ -914,11 +914,27 @@ export function queryEditorSetAndSaveSql(targetQueryEditor, sql, queryId) {
 
 export function formatQuery(queryEditor) {
   return function (dispatch, getState) {
-    const { sql } = getUpToDateQuery(getState(), queryEditor);
+    const { sql, dbId, templateParams } = getUpToDateQuery(
+      getState(),
+      queryEditor,
+    );
+    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));

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