Continuation of the work on the SET statement (#91)

* allow parsing of context only

* more tools around adding context to random text

* fix set quotes

* add tests

* changeset

* add changeUnderlyingExpression

* use it.each

* use it.each everywhere

* failing JSON_OBJECT tests

* add key value parsing

* add tests

* rewrite more tests to use it.each

* add more tests

* typo

* add more tests

* add IDE setting

* handle nested JSON_OBJECTS

Author: vogievetsky

.changeset/lemon-hairs-appear.md

@@ -0,0 +1,5 @@
+---
+'druid-query-toolkit': minor
+---
+
+Allow parsing SET statements in an otherwise unparsable string and more integrations"

.gitignore

@@ -9,13 +9,13 @@
 
 /node_modules/
 /coverage/
+/coverage_old/
 /build/
 /dist/
 /types/
 /src/sql/parser/index.ts
 /src/sql/parser/.DS_Store
+CLAUDE.md
 
 # TypeScript cache
 *.tsbuildinfo
-
-_old/

.idea/google-java-format.xml

@@ -75,14 +75,13 @@ ORDER BY 5 DESC
  */
 ```
 
-For more examples check out the unit tests.
+For more examples, check out the unit tests.
 
 #### ToDo
 
 Not every valid DruidSQL construct can currently be parsed, the following snippets are not currently supported:
 
 - `(a, b) IN (subquery)`
-- Support `FROM "wikipedia_k" USING (k)`
 
 ## License
 

README.md

@@ -24,6 +24,7 @@ try {
   parser = peg.generate(header + '\n\n' + rules, {
     output: 'source',
     plugins: [tspegjs],
+    allowedStartRules: ['Start', 'StartSetStatementsOnly'],
   });
 } catch (e) {
   console.error('Failed to compile');

script/compile-peg.js

@@ -36,8 +36,8 @@ function backAndForthNotCustom(expression: string): void {
 }
 
 describe('filter-pattern', () => {
-  it('fixed points', () => {
-    const expressions: string[] = [
+  describe('fixed point expressions', () => {
+    it.each([
       `"lol" = 'hello'`,
       `"lol" <> 'hello'`,
       `"lol" IN ('hello', 'goodbye')`,
@@ -67,28 +67,19 @@ describe('filter-pattern', () => {
       `TIMESTAMP '2022-06-30 22:56:14.123' <= "__time" AND "__time" <= TIMESTAMP '2022-06-30 22:56:15.923'`,
       `TIMESTAMP '2022-06-30 22:56:14.123' < "__time" AND "__time" <= TIMESTAMP '2022-06-30 22:56:15.923'`,
       `(TIME_FLOOR(MAX_DATA_TIME(), 'P3M', NULL, 'Etc/UTC') <= "DIM:__time" AND "DIM:__time" < TIME_SHIFT(TIME_FLOOR(MAX_DATA_TIME(), 'P3M', NULL, 'Etc/UTC'), 'P1D', 1, 'Etc/UTC'))`,
-    ];
-
-    for (const expression of expressions) {
-      try {
-        backAndForthNotCustom(expression);
-      } catch (e) {
-        console.log(`Problem with: \`${expression}\``);
-        throw e;
-      }
-    }
+    ])('correctly handles expression: %s', expression => {
+      backAndForthNotCustom(expression);
+    });
   });
 
-  it('invalid expressions', () => {
-    const expressions: string[] = [
+  describe('invalid expressions', () => {
+    it.each([
       `"__time" >= TIMESTAMP '2022-06-30 22:56:15.923' AND TIMESTAMP '2021-06-30 22:56:14.123' >= "__time"`,
       `TIMESTAMP '2021-06-30 22:56:14.123' >= "__time" AND "__time" >= TIMESTAMP '2022-06-30 22:56:15.923'`,
-    ];
-
-    for (const expression of expressions) {
+    ])('correctly handles invalid expression: %s', expression => {
       const pattern = fitFilterPattern(SqlExpression.parse(expression));
       expect(pattern.type).toEqual('custom');
-    }
+    });
   });
 
   describe('fitFilterPattern', () => {

src/filter-pattern/unify.spec.ts

@@ -40,7 +40,7 @@ export class Introspect {
   }
 
   static getQueryColumnIntrospectionQuery(query: SqlQuery | SqlTable): SqlQuery {
-    return SqlQuery.create(query).changeLimitValue(0);
+    return SqlQuery.selectStarFrom(query).changeLimitValue(0);
   }
 
   static getQueryColumnIntrospectionPayload(

src/introspect/introspect.ts

@@ -38,6 +38,7 @@ export * from './sql-case/sql-when-then-part';
 export * from './sql-case/sql-case';
 export * from './sql-alias/sql-alias';
 export * from './sql-labeled-expression/sql-labeled-expression';
+export * from './sql-key-value/sql-key-value';
 export * from './sql-window-spec/sql-window-spec';
 export * from './sql-window-spec/sql-frame-bound';
 export * from './sql-set-statement/sql-set-statement';

src/sql/index.ts

@@ -12,13 +12,31 @@
  * limitations under the License.
  */
 
-Start = initial:_? thing:(SqlQueryWithPossibleContext / SqlAlias) final:_sc?
+Start = initial:_ thing:(SqlQueryWithPossibleContext / SqlAlias) final:_sc
 {
   if (initial) thing = thing.changeSpace('initial', initial);
   if (final) thing = thing.changeSpace('final', final);
   return thing;
 }
 
+StartSetStatementsOnly = spaceBefore:_ statements:(SqlSetStatement _sc)* rest:$(.*)
+{
+  let ret = {
+    spaceBefore: spaceBefore,
+    rest: rest
+  }
+
+  if (statements.length) {
+    ret.contextStatements = new S.SeparatedArray(
+      statements.map(function(x) { return x[0] }),
+      statements.map(function(x) { return x[1] }).slice(0, statements.length - 1)
+    );
+    ret.spaceAfter = statements[statements.length - 1][1];
+  }
+
+  return ret;
+}
+
 // ------------------------------
 
 SqlAlias = expression:Expression alias:((_ AsToken)? _ RefNameAlias)? columns:(_ SqlColumnList)?
@@ -62,6 +80,38 @@ SqlLabeledExpression = label:RefNameAlias preArrow:_ "=>" postArrow:_ expression
   });
 }
 
+SqlKeyValue = LongKeyValueForm / ShortKeyValueForm
+
+LongKeyValueForm = keyToken:KeyToken postKey:_ key:Expression postKeyExpression:_ valueToken:ValueToken preValueExpression:_ value:Expression
+{
+  return new S.SqlKeyValue({
+    key: key,
+    value: value,
+    spacing: {
+      postKey: postKey,
+      postKeyExpression: postKeyExpression,
+      preValueExpression: preValueExpression
+    },
+    keywords: {
+      key: keyToken,
+      value: valueToken
+    }
+  });
+}
+
+ShortKeyValueForm = key:Expression postKeyExpression:_ ":" preValueExpression:_ value:Expression
+{
+  return new S.SqlKeyValue({
+    key: key,
+    value: value,
+    short: true,
+    spacing: {
+      postKeyExpression: postKeyExpression,
+      preValueExpression: preValueExpression
+    }
+  });
+}
+
 SqlExtendClause =
   extend:(ExtendToken _)?
   OpenParen
@@ -1010,6 +1060,7 @@ Function =
 / TimestampAddDiffFunction
 / PositionFunction
 / JsonValueReturningFunction
+/ JsonObjectFunction
 / ArrayFunction
 / NakedFunction
 
@@ -1171,6 +1222,32 @@ JsonValueReturningFunction =
   });
 }
 
+JsonObjectFunction =
+  functionName:JsonObjectToken
+  preLeftParen:_
+  OpenParen
+  postLeftParen:_
+  head:SqlKeyValue?
+  tail:(CommaSeparator SqlKeyValue)*
+  postArguments:_
+  CloseParen
+{
+  var value = {
+    functionName: makeFunctionName(functionName)
+  };
+  var spacing = value.spacing = {
+    preLeftParen: preLeftParen,
+    postLeftParen: postLeftParen
+  };
+
+  if (head) {
+    value.args = makeSeparatedArray(head, tail);
+    spacing.postArguments = postArguments;
+  }
+
+  return new S.SqlFunction(value);
+}
+
 ExtractFunction =
   functionName:(ExtractToken / ('"' ExtractToken '"'))
   preLeftParen:_
@@ -1881,6 +1958,9 @@ IntoToken = $("INTO"i !IdentifierPart)
 IsToken = $("IS"i !IdentifierPart)
 JoinToken = $("JOIN"i !IdentifierPart)
 JsonValueToken = $("JSON_VALUE"i !IdentifierPart)
+JsonObjectToken = $("JSON_OBJECT"i !IdentifierPart)
+KeyToken = $("KEY"i !IdentifierPart)
+ValueToken = $("VALUE"i !IdentifierPart)
 LeadingToken = $("LEADING"i !IdentifierPart)
 LikeToken = $("LIKE"i !IdentifierPart)
 LimitToken = $("LIMIT"i !IdentifierPart)

src/sql/parser/druidsql.pegjs

@@ -172,9 +172,49 @@ describe('SqlAlias', () => {
   });
 
   describe('.create', () => {
-    expect(
-      SqlAlias.create(SqlAlias.create(SqlColumn.create('X'), 'name1'), 'name2').toString(),
-    ).toEqual('"X" AS "name2"');
+    it('overwrites existing alias when aliasing an already aliased expression', () => {
+      expect(
+        SqlAlias.create(SqlAlias.create(SqlColumn.create('X'), 'name1'), 'name2').toString(),
+      ).toEqual('"X" AS "name2"');
+    });
+
+    it('creates a simple alias with string column and string alias', () => {
+      expect(SqlAlias.create(SqlColumn.create('col1'), 'alias1').toString()).toEqual(
+        '"col1" AS "alias1"',
+      );
+    });
+
+    it('creates an alias with RefName object as alias', () => {
+      const refName = RefName.create('myAlias', true);
+      expect(SqlAlias.create(SqlColumn.create('col1'), refName).toString()).toEqual(
+        '"col1" AS "myAlias"',
+      );
+    });
+
+    it('auto-quotes aliases that are reserved keywords', () => {
+      expect(SqlAlias.create(SqlColumn.create('col1'), 'select').toString()).toEqual(
+        '"col1" AS "select"',
+      );
+    });
+
+    it('forces quotes when forceQuotes is true', () => {
+      expect(SqlAlias.create(SqlColumn.create('col1'), 'normal', true).toString()).toEqual(
+        '"col1" AS "normal"',
+      );
+    });
+
+    it('adds parentheses to SqlQuery expressions', () => {
+      const query = SqlQuery.create('tbl');
+      const aliasedQuery = SqlAlias.create(query, 'subq');
+      const result = aliasedQuery.toString();
+
+      // Check that the result contains the main components rather than exact formatting
+      expect(result).toContain('(');
+      expect(result).toContain(')');
+      expect(result).toContain('SELECT');
+      expect(result).toContain('FROM "tbl"');
+      expect(result).toContain('AS "subq"');
+    });
   });
 
   describe('#changeAlias', () => {