Add flatten method and extend filterAnd for nested expressions

.changeset/heavy-pans-change.md

@@ -0,0 +1,5 @@
+---
+'druid-query-toolkit': patch
+---
+
+Add flatten method and extend filterAnd for nested expressions

.gitignore

@@ -16,6 +16,7 @@
 /src/sql/parser/index.ts
 /src/sql/parser/.DS_Store
 CLAUDE.md
+.claude/
 
 # TypeScript cache
 *.tsbuildinfo

src/sql/sql-expression.spec.ts

@@ -419,6 +419,16 @@ describe('SqlExpression', () => {
         String(SqlExpression.parse('a AND b And c').filterAnd(ex => ex.toString() !== 'b')),
       ).toEqual('a And c');
     });
+
+    it('works with nested AND', () => {
+      expect(
+        String(
+          SqlExpression.parse('a AND b And (x AND y AND b) AND (z AND b)').filterAnd(
+            ex => ex.toString() !== 'b',
+          ),
+        ),
+      ).toEqual('a And (x AND y) AND z');
+    });
   });
 
   describe('extreme', () => {
@@ -444,6 +454,43 @@ describe('SqlExpression', () => {
     });
   });
 
+  describe('#flatten', () => {
+    it('returns the expression as-is for non-multi expressions', () => {
+      const expr = SqlExpression.parse('a = 1');
+      expect(expr.flatten()).toBe(expr);
+    });
+
+    it('flattens nested AND expressions', () => {
+      const expr = SqlExpression.parse('a AND (b AND c)');
+      const flattened = expr.flatten();
+      expect(String(flattened)).toEqual('a AND b AND c');
+    });
+
+    it('flattens nested OR expressions', () => {
+      const expr = SqlExpression.parse('a OR (b OR c)');
+      const flattened = expr.flatten();
+      expect(String(flattened)).toEqual('a OR b OR c');
+    });
+
+    it('does not flatten when flatteningOp does not match', () => {
+      const expr = SqlExpression.parse('a AND b');
+      const flattened = expr.flatten('OR');
+      expect(flattened).toBe(expr);
+    });
+
+    it('flattens deeply nested expressions', () => {
+      const expr = SqlExpression.parse('a AND (b AND (c AND d))');
+      const flattened = expr.flatten();
+      expect(String(flattened)).toEqual('a AND b AND c AND d');
+    });
+
+    it('flattens mixed nested expressions', () => {
+      const expr = SqlExpression.parse('(a OR b) AND ((c OR d) AND e)');
+      const flattened = expr.flatten();
+      expect(String(flattened)).toEqual('(a OR b) AND (c OR d) AND e');
+    });
+  });
+
   describe('#fillPlaceholders', () => {
     it('works in basic case', () => {
       expect(
@@ -536,4 +583,62 @@ describe('SqlExpression', () => {
       );
     });
   });
+
+  describe('#removeColumnFromAnd', () => {
+    it('remove from single expression not AND', () => {
+      const sql = `A > 1`;
+
+      expect(String(SqlExpression.parse(sql).removeColumnFromAnd('A'))).toEqual('undefined');
+      expect(String(SqlExpression.parse(sql).removeColumnFromAnd('B'))).toEqual('A > 1');
+    });
+
+    it('remove from simple AND', () => {
+      const sql = `A AND B`;
+
+      expect(String(SqlExpression.parse(sql).removeColumnFromAnd('A'))).toEqual('B');
+    });
+
+    it('remove from single expression type multiple', () => {
+      const sql = `A AND B AND C`;
+
+      expect(String(SqlExpression.parse(sql).removeColumnFromAnd('A'))).toEqual('B AND C');
+    });
+
+    it('remove from more complex AND', () => {
+      const sql = `A AND B > 1 AND C`;
+
+      expect(String(SqlExpression.parse(sql).removeColumnFromAnd('C'))).toEqual('A AND B > 1');
+    });
+
+    it('handles nested AND comparison expression', () => {
+      const sql = `(A > 1 AND D) AND B AND C`;
+
+      expect(String(SqlExpression.parse(sql).removeColumnFromAnd('A'))).toEqual('D AND B AND C');
+    });
+
+    it('remove nested comparison expression', () => {
+      const sql = `(A > 1 OR D) AND B AND C`;
+
+      expect(String(SqlExpression.parse(sql).removeColumnFromAnd('A'))).toEqual('B AND C');
+    });
+
+    it.each([
+      'A',
+      'B',
+      'A AND B',
+      'A AND B AND C',
+      '(A AND B) AND C',
+      'A AND (B AND C)',
+      'A AND ((B AND C) AND D)',
+      '(A OR a) AND (((B OR b) AND (C OR c)) AND (D OR d))',
+    ])('invariants hold on: %s', sql => {
+      const ex = SqlExpression.parse(sql);
+
+      expect(String(ex.removeColumnFromAnd('X'))).toEqual(sql);
+
+      expect(String(ex.flatten('AND').removeColumnFromAnd('A'))).toEqual(
+        String(ex.removeColumnFromAnd('A')?.flatten('AND')),
+      );
+    });
+  });
 });

src/sql/sql-expression.ts

@@ -12,7 +12,7 @@
  * limitations under the License.
  */
 
-import type { LiteralValue, RefName, SqlOrderByDirection, SqlType } from '.';
+import type { LiteralValue, RefName, SqlMultiOp, SqlOrderByDirection, SqlType } from '.';
 import {
   SqlAlias,
   SqlColumn,
@@ -360,6 +360,10 @@ export abstract class SqlExpression extends SqlBase {
     return SqlExpression.divide(this, ...expressions);
   }
 
+  public flatten(_flatteningOp?: SqlMultiOp): SqlExpression {
+    return this;
+  }
+
   public decomposeViaAnd(_options?: DecomposeViaOptions): SqlExpression[] {
     return [this];
   }

src/sql/sql-multi/sql-multi.spec.ts

@@ -1946,39 +1946,6 @@ describe('Brackets', () => {
   });
 });
 
-describe('#removeColumnFromAnd', () => {
-  it('remove from single expression not AND', () => {
-    const sql = `A > 1`;
-
-    expect(String(SqlExpression.parse(sql).removeColumnFromAnd('A'))).toEqual('undefined');
-    expect(String(SqlExpression.parse(sql).removeColumnFromAnd('B'))).toEqual('A > 1');
-  });
-
-  it('remove from simple AND', () => {
-    const sql = `A AND B`;
-
-    expect(String(SqlExpression.parse(sql).removeColumnFromAnd('A'))).toEqual('B');
-  });
-
-  it('remove from single expression type multiple', () => {
-    const sql = `A AND B AND C`;
-
-    expect(String(SqlExpression.parse(sql).removeColumnFromAnd('A'))).toEqual('B AND C');
-  });
-
-  it('remove from more complex AND', () => {
-    const sql = `A AND B > 1 AND C`;
-
-    expect(String(SqlExpression.parse(sql).removeColumnFromAnd('C'))).toEqual('A AND B > 1');
-  });
-
-  it('remove nested comparison expression', () => {
-    const sql = `(A > 1 OR D) AND B AND C`;
-
-    expect(String(SqlExpression.parse(sql).removeColumnFromAnd('A'))).toEqual('B AND C');
-  });
-});
-
 describe('containsColumn', () => {
   it('nested expression', () => {
     const sql = `A > 1 AND D OR B OR C`;

src/sql/sql-multi/sql-multi.ts

@@ -132,6 +132,15 @@ export class SqlMulti extends SqlExpression {
     return SqlBase.fromValue(value);
   }
 
+  public flatten(flatteningOp?: SqlMultiOp): SqlExpression {
+    const { op, args } = this;
+    if (flatteningOp && op !== flatteningOp) return this;
+    return SqlMulti.create(
+      op,
+      args.values.flatMap(v => v.flatten(op)),
+    );
+  }
+
   public decomposeViaAnd(options: DecomposeViaOptions = {}): SqlExpression[] {
     const { op, args } = this;
     if (op !== 'AND') return super.decomposeViaAnd(options);
@@ -149,7 +158,7 @@ export class SqlMulti extends SqlExpression {
       return super.filterAnd(fn);
     }
 
-    const args = this.args.filter(fn);
+    const args = this.args.filterMap(a => a.filterAnd(fn));
     if (!args) return;
 
     if (args.length() === 1) {