Add WITH ROLLUP (#83)

* Add WITH ROLLUP * build
lucasavila00 · Aug 12, 2024 · b569a4d · b569a4d
1 parent 2d1dffe
commit b569a4d
Show file tree

Hide file tree

Showing 5 changed files with 81 additions and 3 deletions.
diff --git a/docs/api/classes/select-statement.ts.md b/docs/api/classes/select-statement.ts.md
@@ -201,6 +201,7 @@ export declare class SelectStatement<Selection, Alias, Scope, FlatScope> {
       readonly ctes: ReadonlyArray<CTE>;
       readonly alias?: string;
       readonly scope: ScopeStorage;
+      readonly rollup: boolean;
     }
   );
 }
@@ -241,6 +242,7 @@ clickhouse: {
       f: Record<Selection | FlatScope, SafeString> & SelectionOfScope<Scope> & NoSelectFieldsCompileError
     ) => ReplaceT<Selection>
   ) => SelectStatement<Selection, Alias, Scope, FlatScope>
+  withRollup: () => SelectStatement<Selection, Alias, Scope, FlatScope>
 }
 ```
 

diff --git a/src/classes/select-statement.ts b/src/classes/select-statement.ts
@@ -78,6 +78,7 @@ export class SelectStatement<
             readonly ctes: ReadonlyArray<CTE>;
             readonly alias?: string;
             readonly scope: ScopeStorage;
+            readonly rollup: boolean;
         }
     ) {}
 
@@ -106,6 +107,7 @@ export class SelectStatement<
             ctes: [],
             alias,
             scope,
+            rollup: false,
         });
 
     /**
@@ -133,6 +135,7 @@ export class SelectStatement<
             ctes: [],
             alias,
             scope,
+            rollup: false,
         });
 
     /**
@@ -159,6 +162,7 @@ export class SelectStatement<
                 ctes: [],
                 scope: {},
                 alias: undefined,
+                rollup: false,
             }
         );
 
@@ -176,6 +180,14 @@ export class SelectStatement<
         return this;
     };
 
+    private setRollup = (rollup: boolean): this => {
+        this.__props = {
+            ...this.__props,
+            rollup,
+        };
+        return this;
+    };
+
     private setReplace = (replace: ReplaceT<Selection>): this => {
         this.__props = {
             ...this.__props,
@@ -350,6 +362,9 @@ export class SelectStatement<
                     this.__props.scope
                 ) as any),
             ]),
+
+        withRollup: (): SelectStatement<Selection, Alias, Scope, FlatScope> =>
+            this.copy().setRollup(true),
     };
 
     /**

diff --git a/src/print.ts b/src/print.ts
@@ -21,9 +21,13 @@ const printOrderBy = (orderBy: ReadonlyArray<SafeString>): string =>
     orderBy.length > 0
         ? `ORDER BY ${orderBy.map((it) => it.content).join(", ")}`
         : "";
-const printGroupBy = (orderBy: ReadonlyArray<SafeString>): string =>
+const printGroupBy = (
+    orderBy: ReadonlyArray<SafeString>,
+    withRollup: boolean
+): string =>
     orderBy.length > 0
-        ? `GROUP BY ${orderBy.map((it) => it.content).join(", ")}`
+        ? `GROUP BY ${orderBy.map((it) => it.content).join(", ")}` +
+          (withRollup ? " WITH ROLLUP" : "")
         : "";
 const printLimit = (limit: number | SafeString | null): string =>
     limit == null
@@ -250,7 +254,10 @@ export const printSelectStatementInternal = (
         from,
         prewhere,
         where,
-        printGroupBy(selectStatement.__props.groupBy),
+        printGroupBy(
+            selectStatement.__props.groupBy,
+            selectStatement.__props.rollup
+        ),
         having,
         printOrderBy(selectStatement.__props.orderBy),
         printLimit(selectStatement.__props.limit),

diff --git a/tests/clickhouse-suite/rollup.test.ts b/tests/clickhouse-suite/rollup.test.ts
@@ -0,0 +1,28 @@
+import { table } from "../../src";
+import { configureClickhouse } from "../utils";
+import { addSimpleStringSerializer } from "../utils";
+addSimpleStringSerializer();
+
+describe("clickhouse final", () => {
+    const t1 = table(["x", "y"], "t11_clickhouse").clickhouse.final();
+    const { run } = configureClickhouse();
+
+    beforeAll(async () => {
+        await run(`DROP TABLE IF EXISTS t11_clickhouse`);
+        await run(
+            `CREATE TABLE IF NOT EXISTS t11_clickhouse(x Int64, y Int64) ENGINE = AggregatingMergeTree() ORDER BY y`
+        );
+    });
+
+    it("no alias", async () => {
+        const q = t1
+            .selectStar()
+            .groupBy((f) => [f.x, f.y])
+            .clickhouse.withRollup()
+            .stringify();
+        expect(q).toMatchInlineSnapshot(
+            `SELECT * FROM \`t11_clickhouse\` FINAL GROUP BY \`x\`, \`y\` WITH ROLLUP`
+        );
+        expect(await run(q)).toMatchInlineSnapshot(`Array []`);
+    });
+});
diff --git a/tests/joinTable.test.ts b/tests/joinTable.test.ts
@@ -44,6 +44,32 @@ describe("join", () => {
         );
     });
 
+    it("table repro", async () => {
+        const q = t1
+            .select((f) => ({ a: f.a }))
+            .as("main_alias")
+            .join("inner", t2.select((f) => ({ b: f.b })).as("alias2"))
+            .on((f) => equals(f.main_alias.a, f.alias2.b))
+            .select((f) => ({ x: f.a }))
+            .stringify();
+
+        expect(q).toMatchInlineSnapshot(
+            `SELECT \`a\` AS \`x\` FROM (SELECT \`a\` AS \`a\` FROM \`t1\`) AS \`main_alias\` inner JOIN (SELECT \`b\` AS \`b\` FROM \`t2\`) AS \`alias2\` ON \`main_alias\`.\`a\` = \`alias2\`.\`b\``
+        );
+        const q2 = t1
+            .select((f) => ({ a: f.a }))
+            .as("main_alias")
+            .join("inner", t2.select((f) => ({ b: f.b })).as("alias2"))
+            .on((f) => equals(f.main_alias.a, f.alias2.b))
+            .selectStar()
+            .select((f) => ({ x: f.a }))
+            .stringify();
+
+        expect(q2).toMatchInlineSnapshot(
+            `SELECT \`a\` AS \`x\` FROM (SELECT * FROM (SELECT \`a\` AS \`a\` FROM \`t1\`) AS \`main_alias\` inner JOIN (SELECT \`b\` AS \`b\` FROM \`t2\`) AS \`alias2\` ON \`main_alias\`.\`a\` = \`alias2\`.\`b\`)`
+        );
+    });
+
     it("table -> table -- select", async () => {
         const q = t1
             .join("NATURAL", t2)