Web console: fix SET statements appearing in JSON queries

licenses.yaml

@@ -5762,7 +5762,7 @@ license_category: binary
 module: web-console
 license_name: Apache License version 2.0
 copyright: Imply Data
-version: 1.1.1
+version: 1.1.4
 
 ---
 

web-console/.gitignore

@@ -5,6 +5,7 @@ public/
 private/
 lib/*.css
 coverage/
+.claude/
 
 coordinator-console/
 pages/

web-console/package.json

@@ -71,7 +71,7 @@
     "d3-shape": "^3.2.0",
     "d3-time-format": "^4.1.0",
     "date-fns": "^2.28.0",
-    "druid-query-toolkit": "^1.1.1",
+    "druid-query-toolkit": "^1.1.4",
     "echarts": "^5.5.1",
     "file-saver": "^2.0.5",
     "hjson": "^3.2.2",

web-console/src/druid-models/workbench-query/workbench-query.spec.ts

@@ -469,4 +469,221 @@ describe('WorkbenchQuery', () => {
       ).toEqual("End of input while parsing an object (missing '}') at line 2,9 >>>  lol: 1 ...");
     });
   });
+
+  describe('#changeQueryStringContext', () => {
+    it('modifies SQL query string with SET statements', () => {
+      const workbenchQuery = WorkbenchQuery.blank().changeQueryString('SELECT * FROM wikipedia');
+
+      const newContext = { maxNumTasks: 3, useCache: false };
+      const updatedQuery = workbenchQuery.changeQueryStringContext(newContext);
+
+      expect(updatedQuery.getQueryString()).toContain('SET maxNumTasks = 3');
+      expect(updatedQuery.getQueryString()).toContain('SET useCache = FALSE');
+      expect(updatedQuery.getQueryString()).toContain('SELECT * FROM wikipedia');
+    });
+
+    it('updates existing SET statements in SQL query', () => {
+      const workbenchQuery = WorkbenchQuery.blank().changeQueryString(sane`
+        SET maxNumTasks = 2;
+        SELECT * FROM wikipedia
+      `);
+
+      const newContext = { maxNumTasks: 5, finalizeAggregations: true };
+      const updatedQuery = workbenchQuery.changeQueryStringContext(newContext);
+
+      expect(updatedQuery.getQueryString()).toContain('SET maxNumTasks = 5');
+      expect(updatedQuery.getQueryString()).toContain('SET finalizeAggregations = TRUE');
+      expect(updatedQuery.getQueryString()).toContain('SELECT * FROM wikipedia');
+    });
+
+    it('works with JSON queries by modifying queryContext instead', () => {
+      const jsonQuery = sane`
+        {
+          "queryType": "topN",
+          "dataSource": "test"
+        }
+      `;
+
+      const workbenchQuery = WorkbenchQuery.blank()
+        .changeQueryString(jsonQuery)
+        .changeQueryContext({ originalContext: true });
+
+      const newContext = { maxNumTasks: 3, useCache: false };
+      const updatedQuery = workbenchQuery.changeQueryStringContext(newContext);
+
+      expect(updatedQuery.queryContext).toEqual(newContext);
+      expect(updatedQuery.getQueryString()).toEqual(jsonQuery);
+    });
+
+    it('handles empty context object', () => {
+      const workbenchQuery = WorkbenchQuery.blank().changeQueryString('SELECT * FROM wikipedia');
+
+      const updatedQuery = workbenchQuery.changeQueryStringContext({});
+
+      expect(updatedQuery.getQueryString()).toBe('SELECT * FROM wikipedia');
+    });
+
+    it('preserves original queryContext when working with JSON', () => {
+      const jsonQuery = '{"queryType": "timeseries"}';
+      const workbenchQuery = WorkbenchQuery.blank()
+        .changeQueryString(jsonQuery)
+        .changeQueryContext({ existing: 'value' });
+
+      const updatedQuery = workbenchQuery.changeQueryStringContext({ new: 'context' });
+
+      expect(updatedQuery.queryContext).toEqual({ new: 'context' });
+    });
+  });
+
+  describe('#getQueryStringContext', () => {
+    it('extracts context from SQL SET statements', () => {
+      const queryWithSets = sane`
+        SET maxNumTasks = 3;
+        SET useCache = false;
+        SET stringParam = 'test';
+        SELECT * FROM wikipedia
+      `;
+
+      const workbenchQuery = WorkbenchQuery.blank().changeQueryString(queryWithSets);
+      const extractedContext = workbenchQuery.getQueryStringContext();
+
+      expect(extractedContext).toEqual({
+        maxNumTasks: 3,
+        useCache: false,
+        stringParam: 'test',
+      });
+    });
+
+    it('returns empty object when no SET statements in SQL', () => {
+      const workbenchQuery = WorkbenchQuery.blank().changeQueryString('SELECT * FROM wikipedia');
+
+      const extractedContext = workbenchQuery.getQueryStringContext();
+
+      expect(extractedContext).toEqual({});
+    });
+
+    it('returns queryContext for JSON queries', () => {
+      const jsonQuery = '{"queryType": "topN"}';
+      const contextValue = { maxNumTasks: 5, useCache: true };
+
+      const workbenchQuery = WorkbenchQuery.blank()
+        .changeQueryString(jsonQuery)
+        .changeQueryContext(contextValue);
+
+      const extractedContext = workbenchQuery.getQueryStringContext();
+
+      expect(extractedContext).toEqual(contextValue);
+    });
+
+    it('handles mixed SET statements and regular SQL', () => {
+      const queryWithMixedContent = sane`
+        -- Comment
+        SET timeout = 30000;
+        SET maxRows = 1000;
+
+        SELECT COUNT(*)
+        FROM wikipedia
+        WHERE channel = 'en'
+      `;
+
+      const workbenchQuery = WorkbenchQuery.blank().changeQueryString(queryWithMixedContent);
+      const extractedContext = workbenchQuery.getQueryStringContext();
+
+      expect(extractedContext).toEqual({
+        timeout: 30000,
+        maxRows: 1000,
+      });
+    });
+
+    it('returns empty object for malformed JSON queries', () => {
+      const malformedJson = '{ "queryType": "topN"';
+      const workbenchQuery = WorkbenchQuery.blank()
+        .changeQueryString(malformedJson)
+        .changeQueryContext({ fallback: true });
+
+      const extractedContext = workbenchQuery.getQueryStringContext();
+
+      expect(extractedContext).toEqual({ fallback: true });
+    });
+
+    it('handles various data types in SET statements', () => {
+      const queryWithVariousTypes = sane`
+        SET stringParam = 'text';
+        SET numberParam = 42;
+        SET booleanParam = TRUE;
+        SET nullParam = NULL;
+        SELECT * FROM test
+      `;
+
+      const workbenchQuery = WorkbenchQuery.blank().changeQueryString(queryWithVariousTypes);
+      const extractedContext = workbenchQuery.getQueryStringContext();
+
+      expect(extractedContext).toEqual({
+        stringParam: 'text',
+        numberParam: 42,
+        booleanParam: true,
+        nullParam: null,
+      });
+    });
+  });
+
+  describe('changeQueryStringContext and getQueryStringContext symmetry', () => {
+    it('maintains symmetry for SQL queries', () => {
+      const originalQuery = 'SELECT * FROM wikipedia';
+      const testContext = { maxNumTasks: 3, useCache: false, timeout: 30000 };
+
+      const workbenchQuery = WorkbenchQuery.blank().changeQueryString(originalQuery);
+      const updatedQuery = workbenchQuery.changeQueryStringContext(testContext);
+      const extractedContext = updatedQuery.getQueryStringContext();
+
+      expect(extractedContext).toEqual(testContext);
+    });
+
+    it('maintains symmetry for JSON queries', () => {
+      const jsonQuery = '{"queryType": "topN", "dataSource": "test"}';
+      const testContext = { maxNumTasks: 5, finalizeAggregations: true };
+
+      const workbenchQuery = WorkbenchQuery.blank().changeQueryString(jsonQuery);
+      const updatedQuery = workbenchQuery.changeQueryStringContext(testContext);
+      const extractedContext = updatedQuery.getQueryStringContext();
+
+      expect(extractedContext).toEqual(testContext);
+    });
+
+    it('roundtrip preserves simple context values for SQL queries', () => {
+      const simpleContext = {
+        stringValue: 'test',
+        numberValue: 42,
+        booleanValue: true,
+        nullValue: null,
+      };
+
+      const workbenchQuery = WorkbenchQuery.blank()
+        .changeQueryString('SELECT * FROM test')
+        .changeQueryStringContext(simpleContext);
+
+      const extractedContext = workbenchQuery.getQueryStringContext();
+
+      expect(extractedContext).toEqual(simpleContext);
+    });
+
+    it('roundtrip preserves complex context values for JSON queries', () => {
+      const complexContext = {
+        stringValue: 'test',
+        numberValue: 42,
+        booleanValue: true,
+        nullValue: null,
+        arrayValue: [1, 2, 3],
+        objectValue: { nested: { key: 'value' } },
+      };
+
+      const workbenchQuery = WorkbenchQuery.blank()
+        .changeQueryString('{"queryType": "topN"}')
+        .changeQueryStringContext(complexContext);
+
+      const extractedContext = workbenchQuery.getQueryStringContext();
+
+      expect(extractedContext).toEqual(complexContext);
+    });
+  });
 });

web-console/src/druid-models/workbench-query/workbench-query.ts

@@ -286,13 +286,19 @@ export class WorkbenchQuery {
   }
 
   public changeQueryStringContext(queryContext: QueryContext): WorkbenchQuery {
-    const { queryString } = this;
-    return this.changeQueryString(SqlSetStatement.setContextInText(queryString, queryContext));
+    if (this.isJsonLike()) {
+      // JSON query: set the inner context instead of modifying the query string
+      return this.changeQueryContext(queryContext);
+    }
+    return this.changeQueryString(SqlSetStatement.setContextInText(this.queryString, queryContext));
   }
 
   public getQueryStringContext(): QueryContext {
-    const { queryString } = this;
-    return SqlSetStatement.getContextFromText(queryString);
+    if (this.isJsonLike()) {
+      // JSON query: return the inner context for symmetry with changeQueryStringContext
+      return this.queryContext;
+    }
+    return SqlSetStatement.getContextFromText(this.queryString);
   }
 
   public changeQueryContext(queryContext: QueryContext): WorkbenchQuery {