fix: Limit override isn't respected in queryRewrite (#6605)

Author: paveltiunov

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

@@ -71,8 +71,7 @@ import {
   normalizeQuery,
   normalizeQueryCancelPreAggregations,
   normalizeQueryPreAggregationPreview,
-  normalizeQueryPreAggregations,
-  validatePostRewrite,
+  normalizeQueryPreAggregations, remapToQueryAdapterFormat,
 } from './query';
 import { cachedHandler } from './cached-handler';
 import { createJWKsFetcher } from './jwk';
@@ -1101,20 +1100,39 @@ class ApiGateway {
 
     const queries = Array.isArray(query) ? query : [query];
 
-    const normalizedQueries: NormalizedQuery[] = await Promise.all(
+    this.log({
+      type: 'Query Rewrite',
+      query
+    }, context);
+
+    let duration = 0;
+
+    let normalizedQueries: NormalizedQuery[] = await Promise.all(
       queries.map(
-        async (currentQuery) => validatePostRewrite(
-          await this.queryRewrite(
-            normalizeQuery(currentQuery),
-            context
-          )
-        )
+        async (currentQuery) => {
+          const normalizedQuery = normalizeQuery(currentQuery, persistent);
+          const startTime = new Date().getTime();
+          const rewrite = await this.queryRewrite(
+            normalizedQuery,
+            context,
+          );
+          duration += new Date().getTime() - startTime;
+          return normalizeQuery(
+            rewrite,
+            persistent,
+          );
+        }
       )
     );
 
-    normalizedQueries.forEach((q) => {
-      this.processQueryLimit(q, persistent);
-    });
+    this.log({
+      type: 'Query Rewrite completed',
+      normalizedQueries,
+      duration,
+      query
+    }, context);
+
+    normalizedQueries = normalizedQueries.map(q => remapToQueryAdapterFormat(q));
 
     if (normalizedQueries.find((currentQuery) => !currentQuery)) {
       throw new Error('queryTransformer returned null query. Please check your queryTransformer implementation');
@@ -1134,29 +1152,6 @@ class ApiGateway {
     return [queryType, normalizedQueries];
   }
 
-  /**
-   * Asserts query limit, sets the default value if neccessary.
-   *
-   * @throw {Error}
-   */
-  public processQueryLimit(query: NormalizedQuery, persistent = false): void {
-    const def = getEnv('dbQueryDefaultLimit') <= getEnv('dbQueryLimit')
-      ? getEnv('dbQueryDefaultLimit')
-      : getEnv('dbQueryLimit');
-
-    if (!persistent) {
-      if (
-        typeof query.rowLimit === 'number' &&
-        query.rowLimit > getEnv('dbQueryLimit')
-      ) {
-        throw new Error('The query limit has been exceeded.');
-      }
-      query.rowLimit = typeof query.rowLimit === 'number'
-        ? query.rowLimit
-        : def;
-    }
-  }
-
   public async sql({ query, context, res }: QueryRequest) {
     const requestStarted = new Date();
 

packages/cubejs-api-gateway/src/query.js

@@ -1,6 +1,7 @@
 import R from 'ramda';
 import moment from 'moment';
 import Joi from 'joi';
+import { getEnv } from '@cubejs-backend/shared';
 
 import { UserError } from './UserError';
 import { dateParser } from './dateParser';
@@ -106,10 +107,7 @@ const querySchema = Joi.object().keys({
 
 const normalizeQueryOrder = order => {
   let result = [];
-  const normalizeOrderItem = (k, direction) => ({
-    id: k,
-    desc: direction === 'desc'
-  });
+  const normalizeOrderItem = (k, direction) => ([k, direction]);
   if (order) {
     result = Array.isArray(order) ?
       order.map(([k, direction]) => normalizeOrderItem(k, direction)) :
@@ -148,25 +146,13 @@ const checkQueryFilters = (filter) => {
   return true;
 };
 
-const validatePostRewrite = (query) => {
-  const validQuery = query.measures && query.measures.length ||
-    query.dimensions && query.dimensions.length ||
-    query.timeDimensions && query.timeDimensions.filter(td => !!td.granularity).length;
-  if (!validQuery) {
-    throw new UserError(
-      'Query should contain either measures, dimensions or timeDimensions with granularities in order to be valid'
-    );
-  }
-  return query;
-};
-
 /**
  * Normalize incoming network query.
  * @param {Query} query
  * @throws {UserError}
  * @returns {NormalizedQuery}
  */
-const normalizeQuery = (query) => {
+const normalizeQuery = (query, persistent) => {
   const { error } = querySchema.validate(query);
   if (error) {
     throw new UserError(`Invalid query format: ${error.message || error.toString()}`);
@@ -187,9 +173,29 @@ const normalizeQuery = (query) => {
     granularity: d.split('.')[2]
   }));
   const timezone = query.timezone || 'UTC';
+
+  const def = getEnv('dbQueryDefaultLimit') <= getEnv('dbQueryLimit')
+    ? getEnv('dbQueryDefaultLimit')
+    : getEnv('dbQueryLimit');
+
+  let newLimit;
+  if (!persistent) {
+    if (
+      typeof query.limit === 'number' &&
+      query.limit > getEnv('dbQueryLimit')
+    ) {
+      throw new Error('The query limit has been exceeded.');
+    }
+    newLimit = typeof query.limit === 'number'
+      ? query.limit
+      : def;
+  } else {
+    newLimit = query.limit;
+  }
+
   return {
     ...query,
-    rowLimit: query.rowLimit || query.limit,
+    limit: newLimit,
     timezone,
     order: normalizeQueryOrder(query.order),
     filters: (query.filters || []).map(f => {
@@ -233,6 +239,26 @@ const normalizeQuery = (query) => {
   };
 };
 
+const remapQueryOrder = order => {
+  let result = [];
+  const normalizeOrderItem = (k, direction) => ({
+    id: k,
+    desc: direction === 'desc'
+  });
+  if (order) {
+    result = Array.isArray(order) ?
+      order.map(([k, direction]) => normalizeOrderItem(k, direction)) :
+      Object.keys(order).map(k => normalizeOrderItem(k, order[k]));
+  }
+  return result;
+};
+
+const remapToQueryAdapterFormat = (query) => (query ? {
+  ...query,
+  rowLimit: query.limit,
+  order: remapQueryOrder(query.order),
+} : query);
+
 const queryPreAggregationsSchema = Joi.object().keys({
   metadata: Joi.object(),
   timezone: Joi.string(),
@@ -297,9 +323,9 @@ const normalizeQueryCancelPreAggregations = query => {
 export {
   getQueryGranularity,
   getPivotQuery,
-  validatePostRewrite,
   normalizeQuery,
   normalizeQueryPreAggregations,
   normalizeQueryPreAggregationPreview,
   normalizeQueryCancelPreAggregations,
+  remapToQueryAdapterFormat,
 };

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

@@ -61,7 +61,7 @@ interface Query {
   offset?: number;
   total?: boolean;
   totalQuery?: boolean;
-  order?: QueryOrderType;
+  order?: any;
   timezone?: string;
   renewQuery?: boolean;
   ungrouped?: boolean;
@@ -81,6 +81,7 @@ interface NormalizedQueryFilter extends QueryFilter {
 interface NormalizedQuery extends Query {
   filters?: NormalizedQueryFilter[];
   rowLimit?: null | number;
+  order?: [{ id: string; desc: boolean }];
 }
 
 export {

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

@@ -260,6 +260,7 @@ describe('API Gateway', () => {
               order: [],
               filters: [],
               rowLimit: 10000,
+              limit: 10000,
               dimensions: [],
               timeDimensions: [],
               queryType: 'regularQuery'
@@ -272,6 +273,113 @@ describe('API Gateway', () => {
             order: [],
             filters: [],
             rowLimit: 10000,
+            limit: 10000,
+            dimensions: [],
+            timeDimensions: [],
+            queryType: 'regularQuery'
+          },
+          transformedQueries: [null]
+        });
+      }
+    );
+  });
+
+  test('normalize queryRewrite limit', async () => {
+    const { app } = createApiGateway(
+      new AdapterApiMock(),
+      new DataSourceStorageMock(),
+      {
+        checkAuth: (req: Request, authorization) => {
+          if (authorization) {
+            jwt.verify(authorization, API_SECRET);
+            req.authInfo = authorization;
+          }
+        },
+        queryRewrite: async (query, context) => {
+          query.limit = 2;
+          return query;
+        }
+      }
+    );
+
+    const query = {
+      measures: ['Foo.bar']
+    };
+
+    return requestBothGetAndPost(
+      app,
+      { url: '/cubejs-api/v1/dry-run', query: { query: JSON.stringify(query) }, body: { query } },
+      (res) => {
+        expect(res.body).toStrictEqual({
+          queryType: 'regularQuery',
+          normalizedQueries: [
+            {
+              measures: ['Foo.bar'],
+              timezone: 'UTC',
+              order: [],
+              filters: [],
+              rowLimit: 2,
+              limit: 2,
+              dimensions: [],
+              timeDimensions: [],
+              queryType: 'regularQuery'
+            }
+          ],
+          queryOrder: [{ id: 'desc' }],
+          pivotQuery: {
+            measures: ['Foo.bar'],
+            timezone: 'UTC',
+            order: [],
+            filters: [],
+            rowLimit: 2,
+            limit: 2,
+            dimensions: [],
+            timeDimensions: [],
+            queryType: 'regularQuery'
+          },
+          transformedQueries: [null]
+        });
+      }
+    );
+  });
+
+  test('normalize order', async () => {
+    const { app } = createApiGateway();
+
+    const query = {
+      measures: ['Foo.bar'],
+      order: {
+        'Foo.bar': 'desc'
+      }
+    };
+
+    return requestBothGetAndPost(
+      app,
+      { url: '/cubejs-api/v1/dry-run', query: { query: JSON.stringify(query) }, body: { query } },
+      (res) => {
+        expect(res.body).toStrictEqual({
+          queryType: 'regularQuery',
+          normalizedQueries: [
+            {
+              measures: ['Foo.bar'],
+              order: [{ id: 'Foo.bar', desc: true }],
+              timezone: 'UTC',
+              filters: [],
+              rowLimit: 10000,
+              limit: 10000,
+              dimensions: [],
+              timeDimensions: [],
+              queryType: 'regularQuery'
+            }
+          ],
+          queryOrder: [{ id: 'desc' }],
+          pivotQuery: {
+            measures: ['Foo.bar'],
+            order: [{ id: 'Foo.bar', desc: true }],
+            timezone: 'UTC',
+            filters: [],
+            rowLimit: 10000,
+            limit: 10000,
             dimensions: [],
             timeDimensions: [],
             queryType: 'regularQuery'