feat: PatchMeasure member expression (cube-js#9218)

Add new member expression: `PatchMeasure`.
It allows to change aggregation type and add measure filters for existing measure. This is to allow queries like `MIN(avgMeasure)` and `SUM(CASE dim = 'foo' WHEN TRUE THEN sumMeasure ELSE NULL END)` in SQL API.

Under the hood
* SQL API now allows mismatched aggregation functions and aggregation on top of `CASE` in SQL pushdown rules, and rewrites that to special expression `__patch_measure(measure_column, 'new_agg_type', boolean_expression)`
* This expression is syntactic-only, it should never execute, instead it would be handled by wrapper SQL generation
* Generated member expression is passed to JS-side
* `api-gateway` evaluates it from JSON from JS functions
* `BaseQuery` and `BaseMeasure` now can handle new form: generate patched measure definition and use original measure during traversing and collecting

Supporting changes
* Changed rewrites for usual measures in SQL pushdown from just `measure_column` to `MEASURE(measure_column)`, to make rewrites uniform
* Turn member expressions to tagged union in JSON representation

Author: mcheshkov

packages/cubejs-api-gateway/package.json

@@ -30,6 +30,7 @@
     "@cubejs-backend/native": "1.2.30",
     "@cubejs-backend/shared": "1.2.30",
     "@ungap/structured-clone": "^0.3.4",
+    "assert-never": "^1.4.0",
     "body-parser": "^1.19.0",
     "chrono-node": "^2.6.2",
     "express": "^4.21.1",

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

@@ -1,5 +1,6 @@
 /* eslint-disable no-restricted-syntax */
 import * as stream from 'stream';
+import { assertNever } from 'assert-never';
 import jwt, { Algorithm as JWTAlgorithm } from 'jsonwebtoken';
 import R from 'ramda';
 import bodyParser from 'body-parser';
@@ -83,6 +84,7 @@ import {
   normalizeQueryCancelPreAggregations,
   normalizeQueryPreAggregationPreview,
   normalizeQueryPreAggregations,
+  parseInputMemberExpression,
   preAggsJobsRequestSchema,
   remapToQueryAdapterFormat,
 } from './query';
@@ -1392,30 +1394,46 @@ class ApiGateway {
   }
 
   private parseMemberExpression(memberExpression: string): string | ParsedMemberExpression {
-    try {
-      if (memberExpression.startsWith('{')) {
-        const obj = JSON.parse(memberExpression);
-        const args = obj.cube_params;
-        args.push(`return \`${obj.expr}\``);
+    if (memberExpression.startsWith('{')) {
+      const obj = parseInputMemberExpression(JSON.parse(memberExpression));
+      let expression: ParsedMemberExpression['expression'];
+      switch (obj.expr.type) {
+        case 'SqlFunction':
+          expression = [
+            ...obj.expr.cubeParams,
+            `return \`${obj.expr.sql}\``,
+          ];
+          break;
+        case 'PatchMeasure':
+          expression = {
+            type: 'PatchMeasure',
+            sourceMeasure: obj.expr.sourceMeasure,
+            replaceAggregationType: obj.expr.replaceAggregationType,
+            addFilters: obj.expr.addFilters.map(filter => [
+              ...filter.cubeParams,
+              `return \`${filter.sql}\``,
+            ]),
+          };
+          break;
+        default:
+          assertNever(obj.expr);
+      }
 
-        const groupingSet = obj.grouping_set ? {
-          groupType: obj.grouping_set.group_type,
-          id: obj.grouping_set.id,
-          subId: obj.grouping_set.sub_id ? obj.grouping_set.sub_id : undefined
-        } : undefined;
+      const groupingSet = obj.groupingSet ? {
+        groupType: obj.groupingSet.groupType,
+        id: obj.groupingSet.id,
+        subId: obj.groupingSet.subId ? obj.groupingSet.subId : undefined
+      } : undefined;
 
-        return {
-          cubeName: obj.cube_name,
-          name: obj.alias,
-          expressionName: obj.alias,
-          expression: args,
-          definition: memberExpression,
-          groupingSet,
-        };
-      } else {
-        return memberExpression;
-      }
-    } catch {
+      return {
+        cubeName: obj.cubeName,
+        name: obj.alias,
+        expressionName: obj.alias,
+        expression,
+        definition: memberExpression,
+        groupingSet,
+      };
+    } else {
       return memberExpression;
     }
   }
@@ -1433,14 +1451,31 @@ class ApiGateway {
     };
   }
 
-  private evalMemberExpression(memberExpression: MemberExpression | ParsedMemberExpression): string | MemberExpression {
-    const expression = Array.isArray(memberExpression.expression) ?
-      Function.constructor.apply(null, memberExpression.expression) : memberExpression.expression;
+  private evalMemberExpression(memberExpression: MemberExpression | ParsedMemberExpression): MemberExpression | ParsedMemberExpression {
+    if (typeof memberExpression.expression === 'function') {
+      return memberExpression;
+    }
 
-    return {
-      ...memberExpression,
-      expression,
-    };
+    if (Array.isArray(memberExpression.expression)) {
+      return {
+        ...memberExpression,
+        expression: Function.constructor.apply(null, memberExpression.expression),
+      };
+    }
+
+    if (memberExpression.expression.type === 'PatchMeasure') {
+      return {
+        ...memberExpression,
+        expression: {
+          ...memberExpression.expression,
+          addFilters: memberExpression.expression.addFilters.map(filter => ({
+            sql: Function.constructor.apply(null, filter),
+          })),
+        }
+      };
+    }
+
+    throw new Error(`Unexpected member expression to evaluate: ${memberExpression}`);
   }
 
   public async sqlGenerators({ context, res }: { context: RequestContext, res: ResponseResultFn }) {

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

@@ -39,11 +39,31 @@ const getPivotQuery = (queryType, queries) => {
   return pivotQuery;
 };
 
+const parsedPatchMeasureFilterExpression = Joi.array().items(Joi.string());
+
+const evaluatedPatchMeasureFilterExpression = Joi.object().keys({
+  sql: Joi.func().required(),
+});
+
+const parsedPatchMeasureExpression = Joi.object().keys({
+  type: Joi.valid('PatchMeasure').required(),
+  sourceMeasure: Joi.string().required(),
+  replaceAggregationType: Joi.string().allow(null).required(),
+  addFilters: Joi.array().items(parsedPatchMeasureFilterExpression).required(),
+});
+
+const evaluatedPatchMeasureExpression = parsedPatchMeasureExpression.keys({
+  addFilters: Joi.array().items(evaluatedPatchMeasureFilterExpression).required(),
+});
+
 const id = Joi.string().regex(/^[a-zA-Z0-9_]+\.[a-zA-Z0-9_]+$/);
 const idOrMemberExpressionName = Joi.string().regex(/^[a-zA-Z0-9_]+\.[a-zA-Z0-9_]+$|^[a-zA-Z0-9_]+$/);
 const dimensionWithTime = Joi.string().regex(/^[a-zA-Z0-9_]+\.[a-zA-Z0-9_]+(\.[a-zA-Z0-9_]+)?$/);
 const parsedMemberExpression = Joi.object().keys({
-  expression: Joi.array().items(Joi.string()).min(1).required(),
+  expression: Joi.alternatives(
+    Joi.array().items(Joi.string()).min(1),
+    parsedPatchMeasureExpression,
+  ).required(),
   cubeName: Joi.string().required(),
   name: Joi.string().required(),
   expressionName: Joi.string(),
@@ -55,7 +75,43 @@ const parsedMemberExpression = Joi.object().keys({
   })
 });
 const memberExpression = parsedMemberExpression.keys({
-  expression: Joi.func().required(),
+  expression: Joi.alternatives(
+    Joi.func().required(),
+    evaluatedPatchMeasureExpression,
+  ).required(),
+});
+
+const inputSqlFunction = Joi.object().keys({
+  cubeParams: Joi.array().items(Joi.string()).required(),
+  sql: Joi.string().required(),
+});
+
+// This should be aligned with cubesql side
+const inputMemberExpressionSqlFunction = inputSqlFunction.keys({
+  type: Joi.valid('SqlFunction').required(),
+});
+
+// This should be aligned with cubesql side
+const inputMemberExpressionPatchMeasure = Joi.object().keys({
+  type: Joi.valid('PatchMeasure').required(),
+  sourceMeasure: Joi.string().required(),
+  replaceAggregationType: Joi.string().allow(null).required(),
+  addFilters: Joi.array().items(inputSqlFunction).required(),
+});
+
+// This should be aligned with cubesql side
+const inputMemberExpression = Joi.object().keys({
+  cubeName: Joi.string().required(),
+  alias: Joi.string().required(),
+  expr: Joi.alternatives(
+    inputMemberExpressionSqlFunction,
+    inputMemberExpressionPatchMeasure,
+  ),
+  groupingSet: Joi.object().keys({
+    groupType: Joi.valid('Rollup', 'Cube').required(),
+    id: Joi.number().required(),
+    subId: Joi.number().allow(null),
+  }).allow(null)
 });
 
 const operators = [
@@ -215,6 +271,20 @@ const normalizeQueryFilters = (filter) => (
   })
 );
 
+/**
+ * Parse incoming member expression
+ * @param {unknown} expression
+ * @throws {import('./UserError').UserError}
+ * @returns {import('./types/query').InputMemberExpression}
+ */
+function parseInputMemberExpression(expression) {
+  const { error } = inputMemberExpression.validate(expression);
+  if (error) {
+    throw new UserError(`Invalid member expression format: ${error.message || error.toString()}`);
+  }
+  return expression;
+}
+
 /**
  * Normalize incoming network query.
  * @param {Query} query
@@ -384,5 +454,6 @@ export {
   normalizeQueryPreAggregations,
   normalizeQueryPreAggregationPreview,
   normalizeQueryCancelPreAggregations,
+  parseInputMemberExpression,
   remapToQueryAdapterFormat,
 };

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

@@ -38,14 +38,34 @@ type LogicalOrFilter = {
   or: (QueryFilter | LogicalAndFilter)[]
 };
 
+export type GroupingSetType = 'Rollup' | 'Cube';
+
 type GroupingSet = {
-    groupType: string,
+    groupType: GroupingSetType,
     id: number,
     subId?: null | number
 };
 
+export type EvalPatchMeasureFilterExpression = {
+  sql: Function,
+};
+
+export type PatchMeasureExpression = {
+  type: 'PatchMeasure',
+  sourceMeasure: string,
+  replaceAggregationType: string | null,
+  addFilters: Array<Array<string>>,
+};
+
+export type EvalPatchMeasureExpression = {
+  type: 'PatchMeasure',
+  sourceMeasure: string,
+  replaceAggregationType: string | null,
+  addFilters: Array<EvalPatchMeasureFilterExpression>,
+};
+
 type ParsedMemberExpression = {
-  expression: string[];
+  expression: string[] | PatchMeasureExpression;
   cubeName: string;
   name: string;
   expressionName: string;
@@ -54,7 +74,33 @@ type ParsedMemberExpression = {
 };
 
 type MemberExpression = Omit<ParsedMemberExpression, 'expression'> & {
-  expression: Function;
+  expression: Function | EvalPatchMeasureExpression;
+};
+
+type InputSqlFunction = {
+  cubeParams: Array<string>,
+  sql: string,
+};
+
+export type InputMemberExpressionSqlFunction = {
+  type: 'SqlFunction'
+} & InputSqlFunction;
+
+export type InputMemberExpressionPatchMeasure = {
+  type: 'PatchMeasure',
+  sourceMeasure: string,
+  replaceAggregationType: string | null,
+  addFilters: Array<InputSqlFunction>,
+};
+
+export type InputMemberExpressionExpr = InputMemberExpressionSqlFunction | InputMemberExpressionPatchMeasure;
+
+// This should be aligned with cubesql side
+export type InputMemberExpression = {
+  cubeName: string,
+  alias: string,
+  expr: InputMemberExpressionExpr,
+  groupingSet: GroupingSet | null,
 };
 
 /**

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

@@ -771,6 +771,21 @@ describe('API Gateway', () => {
   });
 
   describe('sql api member expressions evaluations', () => {
+    const query = {
+      measures: [
+        // eslint-disable-next-line no-template-curly-in-string
+        '{"cubeName":"sales","alias":"sum_sales_line_i","expr":{"type":"SqlFunction","cubeParams":["sales"],"sql":"SUM(${sales.line_items_price})"},"groupingSet":null}'
+      ],
+      dimensions: [
+        // eslint-disable-next-line no-template-curly-in-string
+        '{"cubeName":"sales","alias":"users_age","expr":{"type":"SqlFunction","cubeParams":["sales"],"sql":"${sales.users_age}"},"groupingSet":null}',
+        // eslint-disable-next-line no-template-curly-in-string
+        '{"cubeName":"sales","alias":"cast_sales_users","expr":{"type":"SqlFunction","cubeParams":["sales"],"sql":"CAST(${sales.users_first_name} AS TEXT)"},"groupingSet":null}'
+      ],
+      segments: [],
+      order: []
+    };
+
     test('throw error if expressions are not allowed', async () => {
       const { apiGateway } = await createApiGateway();
       const request: QueryRequest = {
@@ -779,20 +794,7 @@ describe('API Gateway', () => {
           const errorMessage = message as { error: string };
           expect(errorMessage.error).toEqual('Error: Expressions are not allowed in this context');
         },
-        query: {
-          measures: [
-            // eslint-disable-next-line no-template-curly-in-string
-            '{"cube_name":"sales","alias":"sum_sales_line_i","cube_params":["sales"],"expr":"SUM(${sales.line_items_price})","grouping_set":null}'
-          ],
-          dimensions: [
-            // eslint-disable-next-line no-template-curly-in-string
-            '{"cube_name":"sales","alias":"users_age","cube_params":["sales"],"expr":"${sales.users_age}","grouping_set":null}',
-            // eslint-disable-next-line no-template-curly-in-string
-            '{"cube_name":"sales","alias":"cast_sales_users","cube_params":["sales"],"expr":"CAST(${sales.users_first_name} AS TEXT)","grouping_set":null}'
-          ],
-          segments: [],
-          order: []
-        },
+        query,
         expressionParams: [],
         exportAnnotatedSql: true,
         memberExpressions: false,
@@ -820,20 +822,7 @@ describe('API Gateway', () => {
         res(message) {
           expect(message.hasOwnProperty('sql')).toBe(true);
         },
-        query: {
-          measures: [
-            // eslint-disable-next-line no-template-curly-in-string
-            '{"cube_name":"sales","alias":"sum_sales_line_i","cube_params":["sales"],"expr":"SUM(${sales.line_items_price})","grouping_set":null}'
-          ],
-          dimensions: [
-            // eslint-disable-next-line no-template-curly-in-string
-            '{"cube_name":"sales","alias":"users_age","cube_params":["sales"],"expr":"${sales.users_age}","grouping_set":null}',
-            // eslint-disable-next-line no-template-curly-in-string
-            '{"cube_name":"sales","alias":"cast_sales_users","cube_params":["sales"],"expr":"CAST(${sales.users_first_name} AS TEXT)","grouping_set":null}'
-          ],
-          segments: [],
-          order: []
-        },
+        query,
         expressionParams: [],
         exportAnnotatedSql: true,
         memberExpressions: true,