feat: Support Ts in subqueries (#128)

## Summary

Author: bartoval

eslint.config.mjs

@@ -6,7 +6,14 @@ import eslintConfigPrettier from 'eslint-config-prettier';
 import { requireLicenseHeader } from './lint-license-rule.mjs';
 
 export default defineConfig([
-  globalIgnores(['lib/', 'node_modules/', 'src/parser/antlr/', '*.js', '*.mjs']),
+  globalIgnores([
+    'lib/',
+    'node_modules/',
+    'src/parser/antlr/',
+    'storybook-static/',
+    '*.js',
+    '*.mjs',
+  ]),
   eslint.configs.recommended,
   ...tseslint.configs.recommended,
   {

src/parser/__tests__/from.test.ts

@@ -528,5 +528,36 @@ describe('FROM', () => {
         },
       });
     });
+
+    it('can parse ts subquery as source', () => {
+      const text = 'FROM (TS k8s)';
+      const { ast, errors } = parse(text);
+
+      expect(errors.length).toBe(0);
+      expect(ast[0].args[0]).toMatchObject({
+        type: 'parens',
+        child: {
+          type: 'query',
+          commands: [{ type: 'command', name: 'ts' }],
+        },
+      });
+    });
+
+    it('can parse ts subquery with pipes', () => {
+      const text = 'FROM (TS k8s | STATS max_bytes = max(bytes) BY cluster)';
+      const { ast, errors } = parse(text);
+
+      expect(errors.length).toBe(0);
+      expect(ast[0].args[0]).toMatchObject({
+        type: 'parens',
+        child: {
+          type: 'query',
+          commands: [
+            { type: 'command', name: 'ts' },
+            { type: 'command', name: 'stats' },
+          ],
+        },
+      });
+    });
   });
 });

src/parser/__tests__/function.test.ts

@@ -584,6 +584,59 @@ describe('function AST nodes', () => {
           ],
         });
       });
+
+      it('IN ts subquery', () => {
+        const query = 'FROM a | WHERE a IN (TS b | KEEP c)';
+        const { root, errors } = parse(query);
+        const expression = Walker.findFunction(root, ({ name }) => name === 'in');
+
+        expect(errors.length).toBe(0);
+        expect(expression?.args.length).toBe(2);
+        expect(expression).toMatchObject({
+          type: 'function',
+          subtype: 'binary-expression',
+          name: 'in',
+          args: [
+            { type: 'column', name: 'a' },
+            {
+              type: 'parens',
+              child: {
+                type: 'query',
+                commands: [
+                  { type: 'command', name: 'ts' },
+                  { type: 'command', name: 'keep' },
+                ],
+              },
+            },
+          ],
+        });
+      });
+
+      it('NOT IN ts subquery', () => {
+        const query = 'FROM a | WHERE a NOT IN (TS b | KEEP c)';
+        const { root, errors } = parse(query);
+        const expression = Walker.findFunction(root, ({ name }) => name === 'not in');
+
+        expect(errors.length).toBe(0);
+        expect(expression).toMatchObject({
+          type: 'function',
+          subtype: 'binary-expression',
+          name: 'not in',
+          args: [
+            { type: 'column', name: 'a' },
+            {
+              type: 'parens',
+              child: {
+                type: 'query',
+                commands: [
+                  { type: 'command', name: 'ts' },
+                  { type: 'command', name: 'keep' },
+                ],
+              },
+            },
+          ],
+        });
+      });
     });
   });
 });

src/parser/core/cst_to_ast_converter.ts

@@ -660,6 +660,7 @@ export class CstToAstConverter {
     if (sourceCommandCtx) {
       const fromCommandCtx = sourceCommandCtx.fromCommand();
       const rowCommandCtx = sourceCommandCtx.rowCommand();
+      const timeSeriesCommandCtx = sourceCommandCtx.timeSeriesCommand();
 
       if (fromCommandCtx) {
         const fromCommand = this.fromFromCommand(fromCommandCtx);
@@ -669,6 +670,8 @@ export class CstToAstConverter {
         }
       } else if (rowCommandCtx) {
         commands.push(this.fromRowCommand(rowCommandCtx));
+      } else if (timeSeriesCommandCtx) {
+        commands.push(this.fromTimeseriesCommand(timeSeriesCommandCtx));
       }
     }
 

src/pretty_print/__tests__/basic_pretty_printer.test.ts

@@ -945,6 +945,12 @@ describe('single line query', () => {
           expect(reprint('FROM a | WHERE b NOT IN (FROM c | KEEP d)').text).toBe(
             'FROM a | WHERE b NOT IN (FROM c | KEEP d)'
           );
+          expect(reprint('FROM a | WHERE b IN (TS c | KEEP d)').text).toBe(
+            'FROM a | WHERE b IN (TS c | KEEP d)'
+          );
+          expect(reprint('FROM a | WHERE b NOT IN (TS c | KEEP d)').text).toBe(
+            'FROM a | WHERE b NOT IN (TS c | KEEP d)'
+          );
         });
       });
     });
@@ -1171,4 +1177,12 @@ describe('subqueries (parens)', () => {
 
     assertReprint(src, expected);
   });
+
+  test('can print ts subquery in FROM', () => {
+    assertReprint('FROM (TS k8s | LIMIT 5)');
+  });
+
+  test('can print mixed FROM subqueries with ts source', () => {
+    assertReprint('FROM idx, (TS k8s | STATS max_bytes = MAX(bytes) BY cluster)');
+  });
 });

src/pretty_print/__tests__/wrapping_pretty_printer.test.ts

@@ -1546,6 +1546,27 @@ describe('subqueries (parens)', () => {
   | SORT max DESC`
     );
   });
+
+  test('can print IN subqueries with ts source', () => {
+    assertReprint(
+      'FROM main_index_with_a_long_name | WHERE main_field_with_a_long_name IN (TS sub_index_with_a_long_name | WHERE sub_field_with_a_long_name > 10 | KEEP sub_field_with_a_long_name)',
+      `FROM main_index_with_a_long_name
+  | WHERE
+      main_field_with_a_long_name IN
+        (TS sub_index_with_a_long_name
+          | WHERE sub_field_with_a_long_name > 10
+          | KEEP sub_field_with_a_long_name)`
+    );
+  });
+
+  test('can print ts subquery in FROM', () => {
+    assertReprint(
+      'FROM index1, (TS k8s_with_a_long_index_name | STATS max_bytes = MAX(bytes) BY cluster)',
+      `FROM
+  index1,
+  (TS k8s_with_a_long_index_name | STATS max_bytes = MAX(bytes) BY cluster)`
+    );
+  });
 });
 
 test.todo('Idempotence on multiple times pretty printing');