feat: AST SQL rewriting

Signed-off-by: Alex Qyoun-ae <4062971+MazterQyou@users.noreply.github.com>

Author: MazterQyou

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

@@ -21,6 +21,7 @@ import {
   ResultWrapper,
   rowsToColumnar,
 } from '@cubejs-backend/native';
+import type { SqlFilterItem, SqlFiltersResponse } from '@cubejs-backend/native';
 import type {
   Application as ExpressApplication,
   ErrorRequestHandler,
@@ -119,6 +120,14 @@ type HandleErrorOptions = {
     requestStarted?: Date
 };
 
+/**
+ * Upper bound on the number of filters a single SQL filters request may carry,
+ * counting the leaves of filter groups. A batch costs one rewrite of the query
+ * regardless of its size, so the bound is about how large a predicate the
+ * rewrite engine is asked to saturate over.
+ */
+const MAX_SQL_FILTERS = 500;
+
 function userAsyncHandler(handler: (req: Request & { context: ExtendedRequestContext }, res: ExpressResponse) => Promise<void>) {
   return (req: ExpressRequest, res: ExpressResponse, next: NextFunction) => {
     handler(req as any, res).catch(next);
@@ -422,6 +431,26 @@ class ApiGateway {
       });
     }));
 
+    app.get(`${this.basePath}/v1/sql-filters`, userMiddlewares, userAsyncHandler(async (req: any, res) => {
+      await this.getSqlFilters({
+        query: req.query.query,
+        context: req.context,
+        res: this.resToResultFn(res)
+      });
+    }));
+
+    app.post(`${this.basePath}/v1/sql-filters`, jsonParser, userMiddlewares, userAsyncHandler(async (req, res) => {
+      await this.modifySqlFilters({
+        query: req.body.query,
+        add: req.body.add,
+        set: req.body.set,
+        delete: req.body.delete,
+        replace: req.body.replace,
+        context: req.context,
+        res: this.resToResultFn(res)
+      });
+    }));
+
     app.get(`${this.basePath}/v1/dry-run`, userMiddlewares, userAsyncHandler(async (req: any, res) => {
       await this.dryRun({
         query: req.query.query,
@@ -1486,6 +1515,154 @@ class ApiGateway {
     }
   }
 
+  /**
+   * Responds with the result of a SQL filters operation. Planning and
+   * rewriting failures are reported by the native layer in-band as
+   * `{ status: 'error', error }`, so they are mapped onto a 4xx to keep
+   * the endpoint's failures visible to status-code-based clients.
+   */
+  protected async resSqlFilters(result: SqlFiltersResponse, res: ResponseResultFn) {
+    if (result.status === 'error') {
+      await res(result, { status: 400 });
+      return;
+    }
+
+    await res(result);
+  }
+
+  /**
+   * Returns the list of Cube filters of a SQL query in Cube query format,
+   * extracted from the logical plan of the query.
+   */
+  protected async getSqlFilters({
+    query,
+    context,
+    res,
+  }: { query: string } & BaseRequest) {
+    try {
+      await this.assertApiScope('sql', context.securityContext);
+
+      if (typeof query !== 'string' || !query.trim()) {
+        throw new UserError('query parameter must be a non-empty string');
+      }
+
+      const result = await this.sqlServer.getSqlFilters(query, context.securityContext);
+
+      await this.resSqlFilters(result, res);
+    } catch (e: any) {
+      this.handleError({
+        e,
+        context,
+        query,
+        res,
+      });
+    }
+  }
+
+  /**
+   * Modifies the Cube filters of a SQL query. Exactly one of `add`, `set`,
+   * `delete` or `replace` must be provided. `add` adds the requested filters,
+   * `set` replaces all outermost filters with the specified set, `delete`
+   * attempts to delete the requested filters (all occurrences of equal
+   * filters are deleted), and `replace` replaces one exact set of filters
+   * with another (all occurrences of equal filters are replaced). Filters
+   * may be `and`/`or` filter groups; existing filters must match perfectly
+   * to be deleted or replaced. Only the outermost SELECT is modified, and
+   * only dimensions/measures available in the outermost SELECT can be
+   * filtered.
+   */
+  protected async modifySqlFilters({
+    query,
+    add,
+    set,
+    delete: deleteFilters,
+    replace,
+    context,
+    res,
+  }: { query: string, add?: unknown, set?: unknown, delete?: unknown, replace?: unknown } & BaseRequest) {
+    try {
+      await this.assertApiScope('sql', context.securityContext);
+
+      if (typeof query !== 'string' || !query.trim()) {
+        throw new UserError('query parameter must be a non-empty string');
+      }
+
+      const requestedOps = [add, set, deleteFilters, replace].filter((op) => op !== undefined);
+      if (requestedOps.length !== 1) {
+        throw new UserError('Exactly one of add, set, delete or replace parameters is required');
+      }
+
+      // The size of the predicate the rewrite engine ends up with is what
+      // costs, and a filter group nests any number of leaves inside a single
+      // array entry, so the leaves are what is bounded
+      const countFilters = (filters: unknown[]): number => filters.reduce<number>(
+        (total, filter: any) => {
+          // A group is whichever of the two fields holds an array - the other
+          // may be present and not an array, which is for the native layer to
+          // reject rather than something to recurse into
+          const group = Array.isArray(filter?.and) ? filter.and : filter?.or;
+
+          return total + (Array.isArray(group) ? Math.max(countFilters(group), 1) : 1);
+        },
+        0,
+      );
+
+      const assertFilterArray = (filters: unknown, name: string): SqlFilterItem[] => {
+        if (!Array.isArray(filters)) {
+          throw new UserError(`${name} parameter must be an array of filters`);
+        }
+
+        if (countFilters(filters) > MAX_SQL_FILTERS) {
+          throw new UserError(`${name} parameter must contain at most ${MAX_SQL_FILTERS} filters`);
+        }
+
+        return filters;
+      };
+
+      if (add !== undefined) {
+        const result = await this.sqlServer.addSqlFilters(query, assertFilterArray(add, 'add'), context.securityContext);
+
+        await this.resSqlFilters(result, res);
+        return;
+      }
+
+      if (set !== undefined) {
+        const result = await this.sqlServer.setSqlFilters(query, assertFilterArray(set, 'set'), context.securityContext);
+
+        await this.resSqlFilters(result, res);
+        return;
+      }
+
+      if (deleteFilters !== undefined) {
+        const result = await this.sqlServer.deleteSqlFilters(query, assertFilterArray(deleteFilters, 'delete'), context.securityContext);
+
+        await this.resSqlFilters(result, res);
+        return;
+      }
+
+      if (typeof replace !== 'object' || replace === null || Array.isArray(replace)) {
+        throw new UserError('replace parameter must be an object with old and new filter arrays');
+      }
+
+      const { old: oldFilters, new: newFilters } = replace as Record<string, unknown>;
+      const result = await this.sqlServer.replaceSqlFilters(
+        query,
+        assertFilterArray(oldFilters, 'replace.old'),
+        assertFilterArray(newFilters, 'replace.new'),
+        context.securityContext,
+      );
+
+      await this.resSqlFilters(result, res);
+    } catch (e: any) {
+      this.handleError({
+        e,
+        context,
+        query,
+        res,
+      });
+    }
+  }
+
   public async sql({
     query,
     context,

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

@@ -5,11 +5,18 @@ import {
   execSql,
   sql4sql,
   rest4sql,
+  getSqlFilters,
+  addSqlFilters,
+  setSqlFilters,
+  deleteSqlFilters,
+  replaceSqlFilters,
   SqlInterfaceInstance,
   Request as NativeRequest,
   LoadRequestMeta,
   Sql4SqlResponse,
   QueryConvertResponse,
+  SqlFilterItem,
+  SqlFiltersResponse,
 } from '@cubejs-backend/native';
 import type { ShutdownMode } from '@cubejs-backend/native';
 import { displayCLIWarning, getEnv, CacheMode } from '@cubejs-backend/shared';
@@ -88,6 +95,26 @@ export class SQLServer {
     return rest4sql(this.getSqlInterfaceInstance(), sqlQuery, securityContext);
   }
 
+  public async getSqlFilters(sqlQuery: string, securityContext?: unknown): Promise<SqlFiltersResponse> {
+    return getSqlFilters(this.getSqlInterfaceInstance(), sqlQuery, securityContext);
+  }
+
+  public async addSqlFilters(sqlQuery: string, filters: SqlFilterItem[], securityContext?: unknown): Promise<SqlFiltersResponse> {
+    return addSqlFilters(this.getSqlInterfaceInstance(), sqlQuery, filters, securityContext);
+  }
+
+  public async setSqlFilters(sqlQuery: string, filters: SqlFilterItem[], securityContext?: unknown): Promise<SqlFiltersResponse> {
+    return setSqlFilters(this.getSqlInterfaceInstance(), sqlQuery, filters, securityContext);
+  }
+
+  public async deleteSqlFilters(sqlQuery: string, filters: SqlFilterItem[], securityContext?: unknown): Promise<SqlFiltersResponse> {
+    return deleteSqlFilters(this.getSqlInterfaceInstance(), sqlQuery, filters, securityContext);
+  }
+
+  public async replaceSqlFilters(sqlQuery: string, oldFilters: SqlFilterItem[], newFilters: SqlFilterItem[], securityContext?: unknown): Promise<SqlFiltersResponse> {
+    return replaceSqlFilters(this.getSqlInterfaceInstance(), sqlQuery, oldFilters, newFilters, securityContext);
+  }
+
   protected buildCheckSqlAuth(options: SQLServerOptions): CheckSQLAuthFn {
     return (options.checkSqlAuth && this.wrapCheckSqlAuthFn(options.checkSqlAuth))
       || this.createDefaultCheckSqlAuthFn(options);

packages/cubejs-backend-native/js/index.ts

@@ -170,6 +170,21 @@ export type QueryConvertResponse = {
   error?: string;
 };
 
+export type SqlFilterItem = {
+  member?: string;
+  operator?: string;
+  values?: string[];
+  and?: SqlFilterItem[];
+  or?: SqlFilterItem[];
+};
+
+export type SqlFiltersResponse = {
+  status: string;
+  sql?: string;
+  filters?: SqlFilterItem[];
+  error?: string;
+};
+
 let loadedNative: any = null;
 
 export function loadNative() {
@@ -459,6 +474,36 @@ export const rest4sql = async (instance: SqlInterfaceInstance, sqlQuery: string,
   return native.rest4sql(instance, sqlQuery, securityContext ? JSON.stringify(securityContext) : null);
 };
 
+export const getSqlFilters = async (instance: SqlInterfaceInstance, sqlQuery: string, securityContext?: unknown): Promise<SqlFiltersResponse> => {
+  const native = loadNative();
+
+  return native.getSqlFilters(instance, sqlQuery, securityContext ? JSON.stringify(securityContext) : null);
+};
+
+export const addSqlFilters = async (instance: SqlInterfaceInstance, sqlQuery: string, filters: SqlFilterItem[], securityContext?: unknown): Promise<SqlFiltersResponse> => {
+  const native = loadNative();
+
+  return native.addSqlFilters(instance, sqlQuery, JSON.stringify(filters), securityContext ? JSON.stringify(securityContext) : null);
+};
+
+export const setSqlFilters = async (instance: SqlInterfaceInstance, sqlQuery: string, filters: SqlFilterItem[], securityContext?: unknown): Promise<SqlFiltersResponse> => {
+  const native = loadNative();
+
+  return native.setSqlFilters(instance, sqlQuery, JSON.stringify(filters), securityContext ? JSON.stringify(securityContext) : null);
+};
+
+export const deleteSqlFilters = async (instance: SqlInterfaceInstance, sqlQuery: string, filters: SqlFilterItem[], securityContext?: unknown): Promise<SqlFiltersResponse> => {
+  const native = loadNative();
+
+  return native.deleteSqlFilters(instance, sqlQuery, JSON.stringify(filters), securityContext ? JSON.stringify(securityContext) : null);
+};
+
+export const replaceSqlFilters = async (instance: SqlInterfaceInstance, sqlQuery: string, oldFilters: SqlFilterItem[], newFilters: SqlFilterItem[], securityContext?: unknown): Promise<SqlFiltersResponse> => {
+  const native = loadNative();
+
+  return native.replaceSqlFilters(instance, sqlQuery, JSON.stringify(oldFilters), JSON.stringify(newFilters), securityContext ? JSON.stringify(securityContext) : null);
+};
+
 export const buildSqlAndParams = (cubeEvaluator: any): any[] => {
   const native = loadNative();
   return native.buildSqlAndParams(cubeEvaluator);

packages/cubejs-backend-native/src/lib.rs

@@ -20,6 +20,7 @@ pub mod orchestrator;
 pub mod python;
 pub mod rest4sql;
 pub mod sql4sql;
+pub mod sql_filters;
 pub mod stream;
 pub mod template;
 pub mod transport;

packages/cubejs-backend-native/src/node_export.rs

@@ -18,6 +18,9 @@ use crate::cubesql_utils::with_session;
 use crate::logger::NodeBridgeLogger;
 use crate::rest4sql::rest4sql;
 use crate::sql4sql::sql4sql;
+use crate::sql_filters::{
+    add_sql_filters, delete_sql_filters, get_sql_filters, replace_sql_filters, set_sql_filters,
+};
 use crate::stream::{OnCloseHandler, OnDrainHandler};
 use crate::tokio_runtime_node;
 use crate::transport::NodeBridgeTransport;
@@ -845,6 +848,11 @@ pub fn register_module_exports<C: NodeConfiguration + 'static>(
     cx.export_function("execSql", exec_sql)?;
     cx.export_function("sql4sql", sql4sql)?;
     cx.export_function("rest4sql", rest4sql)?;
+    cx.export_function("getSqlFilters", get_sql_filters)?;
+    cx.export_function("addSqlFilters", add_sql_filters)?;
+    cx.export_function("setSqlFilters", set_sql_filters)?;
+    cx.export_function("deleteSqlFilters", delete_sql_filters)?;
+    cx.export_function("replaceSqlFilters", replace_sql_filters)?;
     cx.export_function("isFallbackBuild", is_fallback_build)?;
     cx.export_function("__js_to_clrepr_to_js", debug_js_to_clrepr_to_js)?;
 

packages/cubejs-backend-native/src/rest4sql.rs

@@ -14,7 +14,7 @@ use cubesql::compile::engine::df::scan::CubeScanNode;
 use cubesql::transport::TransportLoadRequestQuery;
 use cubesql::CubeError;
 
-fn json_value_to_js<'ctx>(
+pub(crate) fn json_value_to_js<'ctx>(
     cx: &mut impl Context<'ctx>,
     value: &serde_json::Value,
 ) -> JsResult<'ctx, JsValue> {