From 59b23ea8b9b0af979fbddbe85d8e29a099cff84c Mon Sep 17 00:00:00 2001
From: Chris Swenson <cswenson@meta.com>
Date: Thu, 16 Jan 2025 12:08:09 -0600
Subject: [PATCH] Add `defaultRowLimit` parameter when compiling queries
 (#2097)

---
 packages/malloy/src/malloy.ts             |  6 +++-
 packages/malloy/src/model/malloy_query.ts | 21 ++++++++++++
 packages/malloy/src/model/malloy_types.ts |  1 +
 test/src/core/api.spec.ts                 | 40 +++++++++++++++++++++++
 4 files changed, 67 insertions(+), 1 deletion(-)

diff --git a/packages/malloy/src/malloy.ts b/packages/malloy/src/malloy.ts
index a680c4daf..f482b22e3 100644
--- a/packages/malloy/src/malloy.ts
+++ b/packages/malloy/src/malloy.ts
@@ -117,6 +117,7 @@ interface CompileQueryOptions {
   replaceMaterializedReferences?: boolean;
   materializedTablePrefix?: string;
   eventStream?: EventStream;
+  defaultRowLimit?: number;
 }
 
 export class Malloy {
@@ -443,7 +444,10 @@ export class Malloy {
         }
         const compiledSql = queryModel.compileQuery(
           segment,
-          options,
+          {
+            ...options,
+            defaultRowLimit: undefined,
+          },
           false
         ).sql;
         selectStr += parenAlready ? compiledSql : `(${compiledSql})`;
diff --git a/packages/malloy/src/model/malloy_query.ts b/packages/malloy/src/model/malloy_query.ts
index 0f7ce42ae..6fc2a498e 100644
--- a/packages/malloy/src/model/malloy_query.ts
+++ b/packages/malloy/src/model/malloy_query.ts
@@ -5004,12 +5004,33 @@ export class QueryModel {
     };
   }
 
+  addDefaultRowLimit(query: Query, defaultRowLimit?: number): Query {
+    if (defaultRowLimit === undefined) return query;
+    const lastSegment = query.pipeline[query.pipeline.length - 1];
+    if (lastSegment.type === 'raw') return query;
+    if (lastSegment.limit !== undefined) return query;
+    return {
+      ...query,
+      pipeline: [
+        ...query.pipeline.slice(0, -1),
+        {
+          ...lastSegment,
+          limit: defaultRowLimit,
+        },
+      ],
+    };
+  }
+
   compileQuery(
     query: Query,
     prepareResultOptions?: PrepareResultOptions,
     finalize = true
   ): CompiledQuery {
     let newModel: QueryModel | undefined;
+    query = this.addDefaultRowLimit(
+      query,
+      prepareResultOptions?.defaultRowLimit
+    );
     const m = newModel || this;
     const ret = m.loadQuery(
       query,
diff --git a/packages/malloy/src/model/malloy_types.ts b/packages/malloy/src/model/malloy_types.ts
index c3afd0b84..0e55cc970 100644
--- a/packages/malloy/src/model/malloy_types.ts
+++ b/packages/malloy/src/model/malloy_types.ts
@@ -1620,6 +1620,7 @@ export interface SearchValueMapResult {
 export interface PrepareResultOptions {
   replaceMaterializedReferences?: boolean;
   materializedTablePrefix?: string;
+  defaultRowLimit?: number;
 }
 
 type UTD =
diff --git a/test/src/core/api.spec.ts b/test/src/core/api.spec.ts
index e08d4d896..b9e5f7681 100644
--- a/test/src/core/api.spec.ts
+++ b/test/src/core/api.spec.ts
@@ -145,4 +145,44 @@ describe('tags', () => {
   });
 });
 
+describe('default row limit', () => {
+  test('default row limit gets added to query', async () => {
+    const query = runtime.loadQuery(
+      "run: duckdb.table('malloytest.aircraft') -> { group_by: state }"
+    );
+    const result = await query.run({defaultRowLimit: 1});
+    expect(result.data.rowCount).toBe(1);
+  });
+
+  test('default row limit does not get added to raw query', async () => {
+    const query = runtime.loadQuery(
+      `
+        run: duckdb.sql("""
+          SELECT 1 as value
+          UNION ALL
+          SELECT 2 as value
+        """)
+      `
+    );
+    const result = await query.run({defaultRowLimit: 1});
+    expect(result.data.rowCount).toBe(2);
+  });
+
+  test('default row limit does not override existing limit', async () => {
+    const query = runtime.loadQuery(
+      "run: duckdb.table('malloytest.aircraft') -> { group_by: state; limit: 2 }"
+    );
+    const result = await query.run({defaultRowLimit: 1});
+    expect(result.data.rowCount).toBe(2);
+  });
+
+  test('no default row limit does not add row limit', async () => {
+    const query = runtime.loadQuery(
+      "run: duckdb.table('malloytest.aircraft') -> { group_by: state }"
+    );
+    const result = await query.run();
+    expect(result.data.rowCount).toBe(10); // Weird that there are only 10 states in this table?
+  });
+});
+
 afterAll(async () => await runtime.connection.close());