From 9ba1d66e4ef43481387c60aa0c301edea6ad667d Mon Sep 17 00:00:00 2001 From: Pavel Tiunov Date: Mon, 15 May 2023 18:26:08 -0700 Subject: [PATCH] fix: Limit override isn't respected in queryRewrite --- packages/cubejs-api-gateway/src/gateway.ts | 65 +++++------ packages/cubejs-api-gateway/src/query.js | 64 ++++++++--- .../cubejs-api-gateway/src/types/query.ts | 3 +- .../cubejs-api-gateway/test/index.test.ts | 108 ++++++++++++++++++ 4 files changed, 185 insertions(+), 55 deletions(-) diff --git a/packages/cubejs-api-gateway/src/gateway.ts b/packages/cubejs-api-gateway/src/gateway.ts index 2fb07f4a3979f..55f9a5eae1e6b 100644 --- a/packages/cubejs-api-gateway/src/gateway.ts +++ b/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(); diff --git a/packages/cubejs-api-gateway/src/query.js b/packages/cubejs-api-gateway/src/query.js index 7532da3390b21..8e6c74e22e8a1 100644 --- a/packages/cubejs-api-gateway/src/query.js +++ b/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, }; diff --git a/packages/cubejs-api-gateway/src/types/query.ts b/packages/cubejs-api-gateway/src/types/query.ts index 5180412637b05..9f231a35cd651 100644 --- a/packages/cubejs-api-gateway/src/types/query.ts +++ b/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 { diff --git a/packages/cubejs-api-gateway/test/index.test.ts b/packages/cubejs-api-gateway/test/index.test.ts index 69e952f54516b..717ff949f465d 100644 --- a/packages/cubejs-api-gateway/test/index.test.ts +++ b/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'