chore(api-gateway): add limit for sql-runner method (#6050)

Author: MikeNitsenko

packages/cubejs-api-gateway/src/gateway.ts

@@ -78,6 +78,7 @@ import { SQLServer } from './sql-server';
 import { makeSchema } from './graphql';
 import { ConfigItem, prepareAnnotation } from './helpers/prepareAnnotation';
 import transformData from './helpers/transformData';
+import { shouldAddLimit } from './helpers/shouldAddLimit';
 import {
   transformCube,
   transformMeasure,
@@ -1042,26 +1043,55 @@ class ApiGateway {
     const requestStarted = new Date();
     try {
       if (!query) {
-        throw new UserError(
-          'A user\'s query must contain a body'
-        );
+        throw new UserError('A user\'s query must contain a body');
       }
 
       if (!Array.isArray(query) && !query.query) {
         throw new UserError(
           'A user\'s query must contain at least one query param.'
         );
       }
-
+      
       query = {
         ...query,
-        requestId: context.requestId
+        requestId: context.requestId,
       };
 
-      if (query.resultFilter?.objectTypes && !Array.isArray(query.resultFilter.objectTypes)) {
-        throw new UserError(
-          'A query.resultFilter.objectTypes must be an array of strings'
-        );
+      const orchestratorApi = this.getAdapterApi(context);
+
+      if (shouldAddLimit(query.query)) {
+        if (
+          !query.limit ||
+          !Number.isInteger(query.limit) ||
+          query.limit < 0
+        ) {
+          throw new UserError(
+            'A user\'s query must contain limit query param and it must be positive number'
+          );
+        }
+
+        if (query.limit > getEnv('dbQueryDefaultLimit')) {
+          throw new UserError('The query limit has been exceeded');
+        }
+  
+        if (
+          query.resultFilter?.objectTypes &&
+          !Array.isArray(query.resultFilter.objectTypes)
+        ) {
+          throw new UserError(
+            'A query.resultFilter.objectTypes must be an array of strings'
+          );
+        }
+        
+        query.query = query.query.trim();
+        if (query.query.charAt(query.query.length - 1) === ';') {
+          query.query = query.query.slice(0, -1);
+        }
+
+        const driver = await orchestratorApi
+          .driverFactory(query.dataSource || 'default');
+
+        driver.wrapQueryWithLimit(query);
       }
 
       this.log(
@@ -1072,7 +1102,8 @@ class ApiGateway {
         context
       );
 
-      const result = await this.getAdapterApi(context).executeQuery(query);
+      const result = await orchestratorApi.executeQuery(query);
+
       if (result.data.length) {
         const objectLimit = Number(query.resultFilter?.objectLimit) || 100;
         const stringLimit = Number(query.resultFilter?.stringLimit) || 100;

packages/cubejs-api-gateway/src/helpers/shouldAddLimit.ts

@@ -0,0 +1,17 @@
+import { UserError } from '../UserError';
+
+// Determine SQL query type and add LIMIT clause if needed
+export const shouldAddLimit = (sql: string): Boolean => {
+  // TODO: Enhance the way we determine query type
+  const ddlRegex = /^(CREATE|ALTER|DROP|TRUNCATE)\b/i;
+  const dmlRegex = /^(INSERT|UPDATE|DELETE)\b/i;
+  const selectRegex = /^(SELECT|WITH)\b/i;
+
+  if (ddlRegex.test(sql) || dmlRegex.test(sql)) {
+    return false;
+  } else if (selectRegex.test(sql)) {
+    return true;
+  }
+
+  throw new UserError('Invalid SQL query');
+};

packages/cubejs-api-gateway/test/helpers/shouldAddLimit.test.ts

@@ -0,0 +1,55 @@
+import { shouldAddLimit } from '../../src/helpers/shouldAddLimit';
+
+describe('shouldAddLimit', () => {
+  test('returns false for CREATE query', () => {
+    const sql = 'CREATE TABLE users (id INT, name VARCHAR(255))';
+    expect(shouldAddLimit(sql)).toBe(false);
+  });
+
+  test('returns false for ALTER query', () => {
+    const sql = 'ALTER TABLE users ADD COLUMN email VARCHAR(255)';
+    expect(shouldAddLimit(sql)).toBe(false);
+  });
+
+  test('returns false for DROP query', () => {
+    const sql = 'DROP TABLE users';
+    expect(shouldAddLimit(sql)).toBe(false);
+  });
+
+  test('returns false for TRUNCATE query', () => {
+    const sql = 'TRUNCATE TABLE users';
+    expect(shouldAddLimit(sql)).toBe(false);
+  });
+
+  test('returns false for INSERT query', () => {
+    const sql = 'INSERT INTO users (id, name) VALUES (1, "John")';
+    expect(shouldAddLimit(sql)).toBe(false);
+  });
+
+  test('returns false for UPDATE query', () => {
+    const sql = 'UPDATE users SET name="Jane" WHERE id=1';
+    expect(shouldAddLimit(sql)).toBe(false);
+  });
+
+  test('returns false for DELETE query', () => {
+    const sql = 'DELETE FROM users WHERE id=1';
+    expect(shouldAddLimit(sql)).toBe(false);
+  });
+
+  test('returns true for SELECT query', () => {
+    const sql = 'SELECT * FROM users';
+    expect(shouldAddLimit(sql)).toBe(true);
+  });
+
+  test('returns true for WITH query', () => {
+    const sql = 'WITH cte AS (SELECT id, name FROM users) SELECT * FROM cte';
+    expect(shouldAddLimit(sql)).toBe(true);
+  });
+
+  test('throw UserError for invalid SQL query', () => {
+    const sql = 'INVALID QUERY';
+    expect(() => {
+      shouldAddLimit(sql);
+    }).toThrow('Invalid SQL query');
+  });
+});

packages/cubejs-api-gateway/test/index.test.ts

@@ -556,7 +556,8 @@ describe('API Gateway', () => {
         method: 'post',
         successBody: {
           query: {
-            query: 'SELECT * FROM sql-runner',
+            query: 'SELECT * FROM sql-runner; ',
+            limit: 10000,
             resultFilter: {
               objectLimit: 2,
               stringLimit: 2,
@@ -579,6 +580,24 @@ describe('API Gateway', () => {
             error: 'A user\'s query must contain at least one query param.'
           },
           body: { query: {} }
+        }, {
+          result: {
+            status: 400,
+            error: 'A user\'s query must contain limit query param and it must be positive number'
+          },
+          body: { query: { query: 'SELECT * FROM sql-runner' } }
+        }, {
+          result: {
+            status: 400,
+            error: 'The query limit has been exceeded'
+          },
+          body: { query: { query: 'SELECT * FROM sql-runner', limit: 60000 } }
+        }, {
+          result: {
+            status: 400,
+            error: 'A user\'s query must contain limit query param and it must be positive number'
+          },
+          body: { query: { query: 'SELECT * FROM sql-runner', limit: -1 } }
         }, {
           result: {
             status: 400,
@@ -587,6 +606,7 @@ describe('API Gateway', () => {
           body: {
             query: {
               query: 'SELECT * FROM sql-runner',
+              limit: 10000,
               resultFilter: {
                 objectTypes: 'text'
               },

packages/cubejs-api-gateway/test/mocks.ts

@@ -71,6 +71,10 @@ export const compilerApi = jest.fn().mockImplementation(() => ({
     };
   },
 
+  async getDbType() {
+    return 'postgres';
+  },
+
   async metaConfig() {
     return [
       {
@@ -190,7 +194,7 @@ export class AdapterApiMock {
   }
 
   public async executeQuery(query) {
-    if (query?.query === 'SELECT * FROM sql-runner') {
+    if (query?.query.includes('SELECT * FROM sql-runner')) {
       return {
         data: [
           { skip: 'skip' },
@@ -204,6 +208,14 @@ export class AdapterApiMock {
     };
   }
 
+  public driverFactory() {
+    return {
+      wrapQueryWithLimit(query: { query: string; limit: number }) {
+        query.query = `SELECT * FROM (${query.query}) AS t LIMIT ${query.limit}`;
+      },
+    };
+  }
+
   public addDataSeenSource() {
     return undefined;
   }

packages/cubejs-base-driver/src/BaseDriver.ts

@@ -438,4 +438,8 @@ export abstract class BaseDriver implements DriverInterface {
   public nowTimestamp() {
     return Date.now();
   }
+
+  public wrapQueryWithLimit(query: { query: string, limit: number}) {
+    query.query = `SELECT * FROM (${query.query}) AS t LIMIT ${query.limit}`;
+  }
 }

packages/cubejs-base-driver/test/unit/BaseDriver.test.ts

@@ -56,4 +56,14 @@ describe('BaseDriver', () => {
       { name: 'string', type: 'string' }
     ]);
   });
+  
+  test('wrapQueryWithLimit wraps the query with a limit', () => {
+    const driver = new BaseDriverImplementedMock({});
+    const query = { query: 'SELECT * FROM users', limit: 10 };
+    driver.wrapQueryWithLimit(query);
+    expect(query).toEqual({
+      query: 'SELECT * FROM (SELECT * FROM users) AS t LIMIT 10',
+      limit: 10,
+    });
+  });
 });

packages/cubejs-mssql-driver/driver/MSSqlDriver.js

@@ -181,6 +181,10 @@ class MSSqlDriver extends BaseDriver {
   readOnly() {
     return !!this.config.readOnly;
   }
+
+  wrapQueryWithLimit(query) {
+    query.query = `SELECT TOP ${query.limit} * FROM (${query.query}) AS t`;
+  }
 }
 
 module.exports = MSSqlDriver;

packages/cubejs-oracle-driver/driver/OracleDriver.js

@@ -139,6 +139,10 @@ class OracleDriver extends BaseDriver {
   readOnly() {
     return true;
   }
+
+  wrapQueryWithLimit(query) {
+    query.query = `SELECT * FROM (${query.query}) AS t WHERE ROWNUM <= ${query.limit}`;
+  }
 }
 
 module.exports = OracleDriver;