fix(schema-compiler): fix DAP with query_rewrite and python config

Author: bsod90

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

@@ -1196,18 +1196,20 @@ class ApiGateway {
             currentQuery = this.parseMemberExpressionsInQuery(currentQuery);
           }
 
-          let normalizedQuery = normalizeQuery(currentQuery, persistent);
+          const normalizedQuery = normalizeQuery(currentQuery, persistent);
+          let evaluatedQuery = normalizedQuery;
 
           if (hasExpressionsInQuery) {
             // We need to parse/eval all member expressions early as applyRowLevelSecurity
             // needs to access the full SQL query in order to evaluate rules
-            normalizedQuery =
+            evaluatedQuery =
               this.evalMemberExpressionsInQuery(normalizedQuery);
           }
 
           // First apply cube/view level security policies
           const queryWithRlsFilters = await compilerApi.applyRowLevelSecurity(
             normalizedQuery,
+            evaluatedQuery,
             context
           );
           // Then apply user-supplied queryRewrite
@@ -1219,7 +1221,7 @@ class ApiGateway {
           // applyRowLevelSecurity may add new filters which may contain raw member expressions
           // if that's the case, we should run an extra pass of parsing here to make sure
           // nothing breaks down the road
-          if (this.hasExpressionsInQuery(rewrittenQuery)) {
+          if (hasExpressionsInQuery || this.hasExpressionsInQuery(rewrittenQuery)) {
             rewrittenQuery = this.parseMemberExpressionsInQuery(rewrittenQuery);
             rewrittenQuery = this.evalMemberExpressionsInQuery(rewrittenQuery);
           }

packages/cubejs-server-core/src/core/CompilerApi.js

@@ -271,15 +271,15 @@ export class CompilerApi {
    * - combining all filters for different roles with OR
    * - combining cube and view filters with AND
   */
-  async applyRowLevelSecurity(query, context) {
+  async applyRowLevelSecurity(query, evaluatedQuery, context) {
     const compilers = await this.getCompilers({ requestId: context.requestId });
     const { cubeEvaluator } = compilers;
 
     if (!cubeEvaluator.isRbacEnabled()) {
       return query;
     }
 
-    const queryCubes = await this.getCubesFromQuery(query, context);
+    const queryCubes = await this.getCubesFromQuery(evaluatedQuery, context);
 
     // We collect Cube and View filters separately because they have to be
     // applied in "two layers": first Cube filters, then View filters on top

packages/cubejs-testing/birdbox-fixtures/rbac-python/cube.py

@@ -0,0 +1,62 @@
+# Cube configuration options: https://cube.dev/docs/config
+
+from cube import config
+
+
+@config('context_to_roles')
+def context_to_roles(context):
+    return context.get("securityContext", {}).get("auth", {}).get("roles", [])
+
+
+def extract_matching_dicts(data):
+    matching_dicts = []
+    keys = ['values', 'member', 'operator']
+
+    # Recursive function to traverse through the list or dictionary
+    def traverse(element):
+        if isinstance(element, dict):
+            # Check if any of the specified keys are in the dictionary
+            if any(key in element for key in keys):
+                matching_dicts.append(element)
+            # Traverse the dictionary values
+            for value in element.values():
+                traverse(value)
+        elif isinstance(element, list):
+            # Traverse the list items
+            for item in element:
+                traverse(item)
+
+    traverse(data)
+    return matching_dicts
+
+
+@config('query_rewrite')
+def query_rewrite(query: dict, ctx: dict) -> dict:
+    filters = extract_matching_dicts(query.get('filters'))
+
+    for value in range(len(query['timeDimensions'])):
+        filters.append(query['timeDimensions'][value]['dateRange'])
+
+    if not filters or None in filters:
+        raise Exception("Queries can't be run without a filter")
+    return query
+
+
+@config('check_sql_auth')
+def check_sql_auth(query: dict, username: str, password: str) -> dict:
+    if username == 'admin':
+        return {
+            'username': 'admin',
+            'password': password,
+            'securityContext': {
+                'auth': {
+                    'username': 'admin',
+                    'userAttributes': {
+                        'canHaveAdmin': True,
+                        'city': 'New York'
+                    },
+                    'roles': ['admin']
+                }
+            }
+        }
+    raise Exception("Invalid username or password")

packages/cubejs-testing/birdbox-fixtures/rbac-python/model/cubes/users.yaml

@@ -0,0 +1,54 @@
+cubes:
+  - name: users
+    sql_table: users
+
+    measures:
+      - name: count
+        sql: id
+        type: count
+
+    dimensions:
+      - name: city
+        sql: city
+        type: string
+
+      - name: id
+        sql: id
+        type: number
+        primary_key: true
+
+    access_policy:
+      - role: "*"
+        row_level:
+          filters:
+            - member: "{CUBE}.city"
+              operator: equals
+              values: ["{ security_context.auth.userAttributes.city }"]
+      - role: admin
+        conditions:
+          # This thing will fail if there's no auth info in the context
+          # Unfortunately, as of now, there's no way to write more complex expressions
+          # that would allow us to check for the existence of the auth object
+          - if: "{ security_context.auth.userAttributes.canHaveAdmin }"
+        row_level:
+          filters:
+            - or:
+              - and:
+                - member: "{CUBE}.city"
+                  operator: notStartsWith
+                  values:
+                    - London
+                    - "{ security_context.auth.userAttributes.city }"
+                    # mixing string, dynamic values, integers and bools should not
+                    # cause any compilation issues
+                    - 4
+                    - true
+                - member: "city"
+                  operator: notEquals
+                  values:
+                    - 'San Francisco'
+              - member: "{CUBE}.city"
+                operator: equals
+                values:
+                  - "New York"
+

packages/cubejs-testing/src/birdbox.ts

@@ -262,7 +262,7 @@ export async function startBirdBoxFromContainer(
         if (pid !== null) {
           process.kill(pid, signal);
         } else {
-          process.stdout.write(`[Birdbox] Cannot kill Cube instance running in TEST_CUBE_HOST mode without TEST_CUBE_PID defined\n`);
+          process.stdout.write('[Birdbox] Cannot kill Cube instance running in TEST_CUBE_HOST mode without TEST_CUBE_PID defined\n');
           throw new Error('Attempted to use killCube while running with TEST_CUBE_HOST');
         }
       },
@@ -541,9 +541,15 @@ export async function startBirdBoxFromCli(
   }
 
   if (options.cubejsConfig) {
+    const configType = options.cubejsConfig.split('.').at(-1);
+    for (const configFile of ['cube.js', 'cube.py']) {
+      if (fs.existsSync(path.join(testDir, configFile))) {
+        fs.removeSync(path.join(testDir, configFile));
+      }
+    }
     fs.copySync(
       path.join(process.cwd(), 'birdbox-fixtures', options.cubejsConfig),
-      path.join(testDir, 'cube.js')
+      path.join(testDir, `cube.${configType}`)
     );
   }
 

packages/cubejs-testing/test/__snapshots__/smoke-rbac.test.ts.snap

@@ -1,5 +1,13 @@
 // Jest Snapshot v1, https://goo.gl/fbAQLP
 
+exports[`Cube RBAC Engine [Python config] RBAC via SQL API [python config] SELECT * from users: users_python 1`] = `
+Array [
+  Object {
+    "count": "551",
+  },
+]
+`;
+
 exports[`Cube RBAC Engine RBAC via REST API line_items hidden price_dim: line_items_view_no_policy_rest 1`] = `
 Array [
   Object {

packages/cubejs-testing/test/smoke-rbac.test.ts

@@ -13,40 +13,40 @@ import {
   JEST_BEFORE_ALL_DEFAULT_TIMEOUT,
 } from './smoke-tests';
 
+const PG_PORT = 5656;
+let connectionId = 0;
+
+async function createPostgresClient(user: string, password: string) {
+  connectionId++;
+  const currentConnId = connectionId;
+
+  console.debug(`[pg] new connection ${currentConnId}`);
+
+  const conn = new PgClient({
+    database: 'db',
+    port: PG_PORT,
+    host: '127.0.0.1',
+    user,
+    password,
+    ssl: false,
+  });
+  conn.on('error', (err) => {
+    console.log(err);
+  });
+  conn.on('end', () => {
+    console.debug(`[pg] end ${currentConnId}`);
+  });
+
+  await conn.connect();
+
+  return conn;
+}
+
 describe('Cube RBAC Engine', () => {
   jest.setTimeout(60 * 5 * 1000);
   let db: StartedTestContainer;
   let birdbox: BirdBox;
 
-  const pgPort = 5656;
-  let connectionId = 0;
-
-  async function createPostgresClient(user: string, password: string) {
-    connectionId++;
-    const currentConnId = connectionId;
-
-    console.debug(`[pg] new connection ${currentConnId}`);
-
-    const conn = new PgClient({
-      database: 'db',
-      port: pgPort,
-      host: '127.0.0.1',
-      user,
-      password,
-      ssl: false,
-    });
-    conn.on('error', (err) => {
-      console.log(err);
-    });
-    conn.on('end', () => {
-      console.debug(`[pg] end ${currentConnId}`);
-    });
-
-    await conn.connect();
-
-    return conn;
-  }
-
   beforeAll(async () => {
     db = await PostgresDBRunner.startContainer({});
     await PostgresDBRunner.loadEcom(db);
@@ -64,7 +64,7 @@ describe('Cube RBAC Engine', () => {
         CUBEJS_DB_USER: 'test',
         CUBEJS_DB_PASS: 'test',
         //
-        CUBEJS_PG_SQL_PORT: `${pgPort}`,
+        CUBEJS_PG_SQL_PORT: `${PG_PORT}`,
       },
       {
         schemaDir: 'rbac/model',
@@ -345,3 +345,60 @@ describe('Cube RBAC Engine [dev mode]', () => {
     }
   });
 });
+
+describe('Cube RBAC Engine [Python config]', () => {
+  jest.setTimeout(60 * 5 * 1000);
+  let db: StartedTestContainer;
+  let birdbox: BirdBox;
+
+  beforeAll(async () => {
+    db = await PostgresDBRunner.startContainer({});
+    await PostgresDBRunner.loadEcom(db);
+    birdbox = await getBirdbox(
+      'postgres',
+      {
+        ...DEFAULT_CONFIG,
+        CUBEJS_DEV_MODE: 'false',
+        NODE_ENV: 'production',
+        //
+        CUBEJS_DB_TYPE: 'postgres',
+        CUBEJS_DB_HOST: db.getHost(),
+        CUBEJS_DB_PORT: `${db.getMappedPort(5432)}`,
+        CUBEJS_DB_NAME: 'test',
+        CUBEJS_DB_USER: 'test',
+        CUBEJS_DB_PASS: 'test',
+        //
+        CUBEJS_PG_SQL_PORT: `${PG_PORT}`,
+      },
+      {
+        schemaDir: 'rbac-python/model',
+        cubejsConfig: 'rbac-python/cube.py',
+      }
+    );
+  }, JEST_BEFORE_ALL_DEFAULT_TIMEOUT);
+
+  afterAll(async () => {
+    await birdbox.stop();
+    await db.stop();
+  }, JEST_AFTER_ALL_DEFAULT_TIMEOUT);
+
+  describe('RBAC via SQL API [python config]', () => {
+    let connection: PgClient;
+
+    beforeAll(async () => {
+      connection = await createPostgresClient('admin', 'admin_password');
+    });
+
+    afterAll(async () => {
+      await connection.end();
+    }, JEST_AFTER_ALL_DEFAULT_TIMEOUT);
+
+    test('SELECT * from users', async () => {
+      const res = await connection.query('SELECT COUNT(city) as count from "users" HAVING (COUNT(1) > 0)');
+      // const res = await connection.query('SELECT * FROM users limit 10');
+      // This query should return all rows because of the `allow_all` statement
+      // It should also exclude the `created_at` dimension as per memberLevel policy
+      expect(res.rows).toMatchSnapshot('users_python');
+    });
+  });
+});