adapter: Add query policy admission controls

[aljoscha]

Motivation

Ad hoc queries can create temporary dataflows or read object storage on a serving cluster. Operators need an enforceable admission policy, with an observe-only rollout mode, without restricting the DDL that maintains that cluster.

Description

Implements milestone 1 of query policies:

• Dedicated, durable, environment-wide QUERY POLICY objects with ownership, USAGE grants, CREATE/ALTER/DROP/SHOW, and cluster and effective-role attachments.
• reject rules for query_plan_includes = slow_path_query | persist_read. Policies default to warn; enforce rejects with SQLSTATE 53000, the matching policy/rule/cluster, and an EXPLAIN/index hint.
• Both cluster and role restrictions apply. Attachments are durable IDs, not session settings. Policies cannot be dropped while attached.
• Admission runs on the optimized plan in both frontend and backend paths, before execution. It covers SELECT, COPY TO, SUBSCRIBE, and the reads performed by INSERT SELECT/UPDATE/DELETE. DDL and EXPLAIN are exempt.
• Independent default-off flags: enable_query_policies for management and enable_query_policy_enforcement for warnings/rejections.
• Internal mz_query_policy_queries_total{outcome="warned"|"rejected"} counts each affected query once. Introspection exposes policy definitions and rules.

Duration/heap limits, cancellation, active-statement introspection, and violation logging remain out of scope.

Attachments and authorization

Each cluster and role can have at most one policy attached, and each policy can contain multiple rules. A query can therefore be subject to two distinct policies.

This adds an optional QueryPolicyId to durable cluster configuration and role configuration. It does not add a general attachment framework. SQL resolves the policy name when attaching it; subsequent queries resolve that stable ID to the policy definition, so editing a policy affects all its attachments without copying its rules into clusters or roles.

Admission combines the executing cluster's policy and the current effective role's policy. Both constrain the query, with a shared policy evaluated only once. Role attachments are separate from session defaults: SET, RESET, and RESET ALL cannot change them. Catalog transactions reject dangling attachments and reject dropping an attached policy, including with CASCADE. Existing clusters and roles have no attachment after upgrade.

Creation requires the grantable CREATEQUERYPOLICY system privilege, following CREATENETWORKPOLICY. Owners can alter/drop policies. Cluster attachment requires CREATE on the cluster and USAGE on the policy; role attachment requires CREATEROLE and policy USAGE. This combined cluster/role attachment model is new; the underlying catalog and privilege machinery is reused.

Review guide

The diff is large because a dedicated catalog object requires substantial catalog boilerplate: durable storage and serialization, a versioned catalog snapshot and migration, object IDs and exhaustive matches, ownership/privilege handling, and SQL syntax/planning. Most of that follows the existing flat-object patterns.

Start with the behavior:

• src/adapter/src/query_policy.rs: classification of the optimized plan and combined cluster/role admission decisions.
• src/adapter/src/frontend_peek.rs and src/adapter/src/coord/sequencer/inner/{peek,subscribe}.rs: enforcement before execution in both paths, including synthetic read-then-write peeks.
• src/adapter/src/frontend_read_then_write.rs: admission before frontend OCC writes start their internal subscription. Background maintenance is exempt.
• src/adapter/src/catalog/state.rs: lookup of durable cluster and effective-role attachments.
• src/sql/src/plan/statement/ddl.rs and src/sql/src/rbac.rs: rule validation and attachment authorization.

Then review the catalog lifecycle separately from admission behavior. test/sqllogictest/query_policy.slt describes the main user-visible contracts.

Verification

Adds parser regression coverage and test/sqllogictest/query_policy.slt for admission distinctions, both execution paths, read-then-write rejection, rollout, role/cluster composition, and lifecycle. Reference documentation and a serving-cluster rollout guide are included.

An executed SQL demo with actual results and a runnable script demonstrates both execution paths, warnings/rejections, unchanged rows after rejected writes, role/cluster composition, and the rollout switch. Its Prometheus counter increases match the 16 rejected and 2 warned queries.

Release note

None. Query policies are experimental and disabled by default.

[aljoscha]

Executed SQL demo

The demo ran against bin/environmentd --optimized at 6ed0597fc5. All demo assertions passed. It opens a fresh connection after each peek-routing flag change, exercising both frontend and backend peeks. Writes in this server use the legacy read-then-write path. Frontend OCC write admission is covered by query_policy.slt in CI.

Both peek paths produced warnings with correct result rows in warn mode, and SQLSTATE 53000 in enforce mode. Constant/index peeks remained allowed. Rejected INSERT SELECT, UPDATE, and DELETE left the tables unchanged. The run also exercised persist reads, SUBSCRIBE, combined role/cluster restrictions, delegated creation, DDL exemptions, introspection, detach/drop, and disabling enforcement.

The internal Prometheus counters were absent before the demo and reported 16 rejected queries and 2 warned queries afterward, matching the transcript:

mz_query_policy_queries_total{outcome="rejected"} 16
mz_query_policy_queries_total{outcome="warned"} 2

System configuration and system/quickstart grants use the mz_system connection. Policy creation is attempted as qp_creator before and after the grant. The composition queries run as qp_agent. Other SQL uses the bootstrap materialize role.

The transcript contains the actual SQL, result rows, notices, and errors. The script asserts expected rows, warning delivery, and SQLSTATEs.

SQL and actual output

> ALTER SYSTEM SET enable_query_policies = true
ALTER SYSTEM

> ALTER SYSTEM SET enable_query_policy_enforcement = true
ALTER SYSTEM

> ALTER SYSTEM SET persist_fast_path_limit = 1000
ALTER SYSTEM

> SET auto_route_catalog_queries = false
SET

## Delegated policy creation

> CREATE ROLE qp_creator
CREATE ROLE

> CREATE QUERY POLICY qp_delegated (RULES (r(action='reject', metric='query_plan_includes', value='persist_read')))
ERROR [42501]: permission denied for SYSTEM
DETAIL: The 'qp_creator' role needs CREATEQUERYPOLICY privileges on SYSTEM

> GRANT CREATEQUERYPOLICY ON SYSTEM TO qp_creator
GRANT

> CREATE QUERY POLICY qp_delegated (RULES (r(action='reject', metric='query_plan_includes', value='persist_read')))
CREATE QUERY POLICY

> DROP QUERY POLICY qp_delegated
DROP QUERY POLICY

> REVOKE CREATEQUERYPOLICY ON SYSTEM FROM qp_creator
REVOKE

> DROP ROLE qp_creator
DROP ROLE

> CREATE TABLE qp_indexed (a int, b int)
CREATE TABLE

> INSERT INTO qp_indexed VALUES (1, 10), (1, 30), (2, 20)
INSERT 0 3

> CREATE DEFAULT INDEX ON qp_indexed
CREATE INDEX

> CREATE MATERIALIZED VIEW qp_raw AS SELECT * FROM qp_indexed
CREATE MATERIALIZED VIEW

> CREATE TABLE qp_target (a int, b int)
CREATE TABLE

> CREATE DEFAULT INDEX ON qp_target
CREATE INDEX

> CREATE QUERY POLICY qp_fast (RULES (no_dataflows (action='reject', metric='query_plan_includes', value='slow_path_query')))
CREATE QUERY POLICY

> CREATE QUERY POLICY qp_memory (MODE='enforce', RULES (no_storage (action='reject', metric='query_plan_includes', value='persist_read')))
CREATE QUERY POLICY

## Backend peek path

> ALTER SYSTEM SET enable_frontend_peek_sequencing = false
ALTER SYSTEM

> SET auto_route_catalog_queries = false
SET

> ALTER QUERY POLICY qp_fast SET (MODE='warn')
ALTER QUERY POLICY

> ALTER CLUSTER quickstart SET (QUERY POLICY = qp_fast)
ALTER CLUSTER

> SELECT a, sum(b) FROM qp_indexed GROUP BY a ORDER BY a
NOTICE [00000]: query would be rejected under query policy "qp_fast" (rule "no_dataflows"): this query would build a temporary dataflow
DETAIL: Query policy "qp_fast" rule "no_dataflows" applies to this query on cluster "quickstart". Policy mode is 'warn', so this rule does not reject the query.
HINT: Run EXPLAIN to inspect the query plan. Create an index that serves the query on this cluster, or use a cluster and role whose query policies permit this plan.
1 | 40
2 | 20

> ALTER QUERY POLICY qp_fast SET (MODE='enforce')
ALTER QUERY POLICY

> SELECT 7
7

> SELECT * FROM qp_indexed ORDER BY a, b
1 | 10
1 | 30
2 | 20

> SELECT a, sum(b) FROM qp_indexed GROUP BY a
ERROR [53000]: query rejected: this query would build a temporary dataflow
DETAIL: Query policy "qp_fast" rule "no_dataflows" applies to this query on cluster "quickstart".
HINT: Run EXPLAIN to inspect the query plan. Create an index that serves the query on this cluster, or use a cluster and role whose query policies permit this plan.

> EXPLAIN SELECT a, sum(b) FROM qp_indexed GROUP BY a
Explained Query:
  →Accumulable GroupAggregate
    Simple aggregates: sum(#1{b})
    Key:
      Project: #0
    →Arranged materialize.public.qp_indexed

Used Indexes:
  - materialize.public.qp_indexed_primary_idx (*** full scan ***)

Target cluster: quickstart


> INSERT INTO qp_target SELECT a, sum(b)::int FROM qp_indexed GROUP BY a
ERROR [53000]: query rejected: this query would build a temporary dataflow
DETAIL: Query policy "qp_fast" rule "no_dataflows" applies to this query on cluster "quickstart".
HINT: Run EXPLAIN to inspect the query plan. Create an index that serves the query on this cluster, or use a cluster and role whose query policies permit this plan.

> SELECT * FROM qp_target

> UPDATE qp_indexed SET b = b + 1 WHERE b < (SELECT avg(b) FROM qp_indexed)
ERROR [53000]: query rejected: this query would build a temporary dataflow
DETAIL: Query policy "qp_fast" rule "no_dataflows" applies to this query on cluster "quickstart".
HINT: Run EXPLAIN to inspect the query plan. Create an index that serves the query on this cluster, or use a cluster and role whose query policies permit this plan.

> DELETE FROM qp_indexed WHERE b < (SELECT avg(b) FROM qp_indexed)
ERROR [53000]: query rejected: this query would build a temporary dataflow
DETAIL: Query policy "qp_fast" rule "no_dataflows" applies to this query on cluster "quickstart".
HINT: Run EXPLAIN to inspect the query plan. Create an index that serves the query on this cluster, or use a cluster and role whose query policies permit this plan.

> SELECT * FROM qp_indexed ORDER BY a, b
1 | 10
1 | 30
2 | 20

> SUBSCRIBE qp_indexed
ERROR [53000]: query rejected: this query would build a temporary dataflow
DETAIL: Query policy "qp_fast" rule "no_dataflows" applies to this query on cluster "quickstart".
HINT: Run EXPLAIN to inspect the query plan. Create an index that serves the query on this cluster, or use a cluster and role whose query policies permit this plan.

> ALTER CLUSTER quickstart SET (QUERY POLICY = qp_memory)
ALTER CLUSTER

> SELECT a, sum(b) FROM qp_indexed GROUP BY a ORDER BY a
1 | 40
2 | 20

> SELECT * FROM qp_raw LIMIT 10
ERROR [53000]: query rejected: this query would read from object storage
DETAIL: Query policy "qp_memory" rule "no_storage" applies to this query on cluster "quickstart".
HINT: Run EXPLAIN to inspect the query plan. Create an index that serves the query on this cluster, or use a cluster and role whose query policies permit this plan.

> SELECT a, sum(b) FROM qp_raw GROUP BY a
ERROR [53000]: query rejected: this query would read from object storage
DETAIL: Query policy "qp_memory" rule "no_storage" applies to this query on cluster "quickstart".
HINT: Run EXPLAIN to inspect the query plan. Create an index that serves the query on this cluster, or use a cluster and role whose query policies permit this plan.

## Frontend peek path

> ALTER SYSTEM SET enable_frontend_peek_sequencing = true
ALTER SYSTEM

> SET auto_route_catalog_queries = false
SET

> ALTER QUERY POLICY qp_fast SET (MODE='warn')
ALTER QUERY POLICY

> ALTER CLUSTER quickstart SET (QUERY POLICY = qp_fast)
ALTER CLUSTER

> SELECT a, sum(b) FROM qp_indexed GROUP BY a ORDER BY a
NOTICE [00000]: query would be rejected under query policy "qp_fast" (rule "no_dataflows"): this query would build a temporary dataflow
DETAIL: Query policy "qp_fast" rule "no_dataflows" applies to this query on cluster "quickstart". Policy mode is 'warn', so this rule does not reject the query.
HINT: Run EXPLAIN to inspect the query plan. Create an index that serves the query on this cluster, or use a cluster and role whose query policies permit this plan.
1 | 40
2 | 20

> ALTER QUERY POLICY qp_fast SET (MODE='enforce')
ALTER QUERY POLICY

> SELECT 7
7

> SELECT * FROM qp_indexed ORDER BY a, b
1 | 10
1 | 30
2 | 20

> SELECT a, sum(b) FROM qp_indexed GROUP BY a
ERROR [53000]: query rejected: this query would build a temporary dataflow
DETAIL: Query policy "qp_fast" rule "no_dataflows" applies to this query on cluster "quickstart".
HINT: Run EXPLAIN to inspect the query plan. Create an index that serves the query on this cluster, or use a cluster and role whose query policies permit this plan.

> EXPLAIN SELECT a, sum(b) FROM qp_indexed GROUP BY a
Explained Query:
  →Accumulable GroupAggregate
    Simple aggregates: sum(#1{b})
    Key:
      Project: #0
    →Arranged materialize.public.qp_indexed

Used Indexes:
  - materialize.public.qp_indexed_primary_idx (*** full scan ***)

Target cluster: quickstart


> INSERT INTO qp_target SELECT a, sum(b)::int FROM qp_indexed GROUP BY a
ERROR [53000]: query rejected: this query would build a temporary dataflow
DETAIL: Query policy "qp_fast" rule "no_dataflows" applies to this query on cluster "quickstart".
HINT: Run EXPLAIN to inspect the query plan. Create an index that serves the query on this cluster, or use a cluster and role whose query policies permit this plan.

> SELECT * FROM qp_target

> UPDATE qp_indexed SET b = b + 1 WHERE b < (SELECT avg(b) FROM qp_indexed)
ERROR [53000]: query rejected: this query would build a temporary dataflow
DETAIL: Query policy "qp_fast" rule "no_dataflows" applies to this query on cluster "quickstart".
HINT: Run EXPLAIN to inspect the query plan. Create an index that serves the query on this cluster, or use a cluster and role whose query policies permit this plan.

> DELETE FROM qp_indexed WHERE b < (SELECT avg(b) FROM qp_indexed)
ERROR [53000]: query rejected: this query would build a temporary dataflow
DETAIL: Query policy "qp_fast" rule "no_dataflows" applies to this query on cluster "quickstart".
HINT: Run EXPLAIN to inspect the query plan. Create an index that serves the query on this cluster, or use a cluster and role whose query policies permit this plan.

> SELECT * FROM qp_indexed ORDER BY a, b
1 | 10
1 | 30
2 | 20

> SUBSCRIBE qp_indexed
ERROR [53000]: query rejected: this query would build a temporary dataflow
DETAIL: Query policy "qp_fast" rule "no_dataflows" applies to this query on cluster "quickstart".
HINT: Run EXPLAIN to inspect the query plan. Create an index that serves the query on this cluster, or use a cluster and role whose query policies permit this plan.

> ALTER CLUSTER quickstart SET (QUERY POLICY = qp_memory)
ALTER CLUSTER

> SELECT a, sum(b) FROM qp_indexed GROUP BY a ORDER BY a
1 | 40
2 | 20

> SELECT * FROM qp_raw LIMIT 10
ERROR [53000]: query rejected: this query would read from object storage
DETAIL: Query policy "qp_memory" rule "no_storage" applies to this query on cluster "quickstart".
HINT: Run EXPLAIN to inspect the query plan. Create an index that serves the query on this cluster, or use a cluster and role whose query policies permit this plan.

> SELECT a, sum(b) FROM qp_raw GROUP BY a
ERROR [53000]: query rejected: this query would read from object storage
DETAIL: Query policy "qp_memory" rule "no_storage" applies to this query on cluster "quickstart".
HINT: Run EXPLAIN to inspect the query plan. Create an index that serves the query on this cluster, or use a cluster and role whose query policies permit this plan.

## Cluster and role restrictions compose

> CREATE ROLE qp_agent
CREATE ROLE

> GRANT USAGE ON CLUSTER quickstart TO qp_agent
GRANT

> GRANT SELECT ON qp_indexed, qp_raw TO qp_agent
GRANT

> GRANT USAGE ON QUERY POLICY qp_memory TO qp_agent
GRANT

> ALTER ROLE qp_agent SET query_policy = qp_memory
ALTER ROLE

> ALTER CLUSTER quickstart SET (QUERY POLICY = qp_fast)
ALTER CLUSTER

> SELECT a, sum(b) FROM qp_indexed GROUP BY a
ERROR [53000]: query rejected: this query would build a temporary dataflow
DETAIL: Query policy "qp_fast" rule "no_dataflows" applies to this query on cluster "quickstart".
HINT: Run EXPLAIN to inspect the query plan. Create an index that serves the query on this cluster, or use a cluster and role whose query policies permit this plan.

> SELECT * FROM qp_raw LIMIT 10
ERROR [53000]: query rejected: this query would read from object storage
DETAIL: Query policy "qp_memory" rule "no_storage" applies to this query on cluster "quickstart".
HINT: Run EXPLAIN to inspect the query plan. Create an index that serves the query on this cluster, or use a cluster and role whose query policies permit this plan.

## DDL is exempt

> CREATE MATERIALIZED VIEW qp_totals AS SELECT a, sum(b) FROM qp_indexed GROUP BY a
CREATE MATERIALIZED VIEW

> CREATE DEFAULT INDEX ON qp_totals
CREATE INDEX

> SELECT * FROM qp_totals ORDER BY a
1 | 40
2 | 20

## Enforcement rollback switch

> ALTER SYSTEM SET enable_query_policy_enforcement = false
ALTER SYSTEM

> SELECT a, sum(b) FROM qp_indexed GROUP BY a ORDER BY a
1 | 40
2 | 20

> SELECT name, mode FROM mz_internal.mz_query_policies ORDER BY name
qp_fast | enforce
qp_memory | enforce

> SELECT name, action, metric, value FROM mz_internal.mz_query_policy_rules ORDER BY name
no_dataflows | reject | query_plan_includes | slow_path_query
no_storage | reject | query_plan_includes | persist_read

> ALTER ROLE qp_agent RESET query_policy
ALTER ROLE

> ALTER CLUSTER quickstart RESET (QUERY POLICY)
ALTER CLUSTER

> DROP QUERY POLICY qp_fast, qp_memory
DROP QUERY POLICY

Demo assertions passed.

Runnable Python script (requires psycopg)

import os

import psycopg

notices = []


def notice(diag):
    notices.append(diag.message_primary)
    print(f"NOTICE [{diag.sqlstate}]: {diag.message_primary}")
    if diag.message_detail:
        print(f"DETAIL: {diag.message_detail}")
    if diag.message_hint:
        print(f"HINT: {diag.message_hint}")


def run(conn, sql, *, error=None, rows=None, warning=None):
    print(f"\n> {sql}", flush=True)
    before_notices = len(notices)
    try:
        cur = conn.execute(sql)
        actual = cur.fetchall() if cur.description else None
        if actual is not None:
            for row in actual:
                print(" | ".join(str(v) for v in row))
        else:
            print(cur.statusmessage)
        assert error is None, f"expected SQLSTATE {error}"
        if rows is not None:
            assert actual == rows, (actual, rows)
        if warning is not None:
            assert any(warning in text for text in notices[before_notices:]), notices[before_notices:]
    except psycopg.Error as exc:
        print(f"ERROR [{exc.sqlstate}]: {exc.diag.message_primary}")
        if exc.diag.message_detail:
            print(f"DETAIL: {exc.diag.message_detail}")
        if exc.diag.message_hint:
            print(f"HINT: {exc.diag.message_hint}")
        assert error is not None and exc.sqlstate == error, str(exc)


admin = psycopg.connect(
    os.getenv("QP_SYSTEM_URL", "postgresql://mz_system@127.0.0.1:6877/materialize"),
    autocommit=True,
)
client = psycopg.connect(
    os.getenv("QP_SQL_URL", "postgresql://materialize@127.0.0.1:6875/materialize"),
    autocommit=True,
)
client.add_notice_handler(notice)
for setting in ["enable_query_policies", "enable_query_policy_enforcement"]:
    run(admin, f"ALTER SYSTEM SET {setting} = true")
run(admin, "ALTER SYSTEM SET persist_fast_path_limit = 1000")
run(client, "SET auto_route_catalog_queries = false")
print("\n## Delegated policy creation", flush=True)
run(client, "CREATE ROLE qp_creator")
creator = psycopg.connect(client.info.dsn, user="qp_creator", autocommit=True)
run(creator, "CREATE QUERY POLICY qp_delegated (RULES (r(action='reject', metric='query_plan_includes', value='persist_read')))", error="42501")
run(admin, "GRANT CREATEQUERYPOLICY ON SYSTEM TO qp_creator")
run(creator, "CREATE QUERY POLICY qp_delegated (RULES (r(action='reject', metric='query_plan_includes', value='persist_read')))")
run(creator, "DROP QUERY POLICY qp_delegated")
creator.close()
run(admin, "REVOKE CREATEQUERYPOLICY ON SYSTEM FROM qp_creator")
run(client, "DROP ROLE qp_creator")
run(client, "CREATE TABLE qp_indexed (a int, b int)")
run(client, "INSERT INTO qp_indexed VALUES (1, 10), (1, 30), (2, 20)")
run(client, "CREATE DEFAULT INDEX ON qp_indexed")
run(client, "CREATE MATERIALIZED VIEW qp_raw AS SELECT * FROM qp_indexed")
run(client, "CREATE TABLE qp_target (a int, b int)")
run(client, "CREATE DEFAULT INDEX ON qp_target")
run(client, "CREATE QUERY POLICY qp_fast (RULES (no_dataflows (action='reject', metric='query_plan_includes', value='slow_path_query')))")
run(client, "CREATE QUERY POLICY qp_memory (MODE='enforce', RULES (no_storage (action='reject', metric='query_plan_includes', value='persist_read')))")

for frontend in [False, True]:
    print(f"\n## {'Frontend' if frontend else 'Backend'} peek path", flush=True)
    run(admin, f"ALTER SYSTEM SET enable_frontend_peek_sequencing = {str(frontend).lower()}")
    dsn = client.info.dsn
    client.close()
    client = psycopg.connect(dsn, autocommit=True)
    client.add_notice_handler(notice)
    run(client, "SET auto_route_catalog_queries = false")
    run(client, "ALTER QUERY POLICY qp_fast SET (MODE='warn')")
    run(client, "ALTER CLUSTER quickstart SET (QUERY POLICY = qp_fast)")
    run(client, "SELECT a, sum(b) FROM qp_indexed GROUP BY a ORDER BY a", rows=[(1, 40), (2, 20)], warning="qp_fast")
    run(client, "ALTER QUERY POLICY qp_fast SET (MODE='enforce')")
    run(client, "SELECT 7", rows=[(7,)])
    run(client, "SELECT * FROM qp_indexed ORDER BY a, b", rows=[(1, 10), (1, 30), (2, 20)])
    run(client, "SELECT a, sum(b) FROM qp_indexed GROUP BY a", error="53000")
    run(client, "EXPLAIN SELECT a, sum(b) FROM qp_indexed GROUP BY a")
    run(client, "INSERT INTO qp_target SELECT a, sum(b)::int FROM qp_indexed GROUP BY a", error="53000")
    run(client, "SELECT * FROM qp_target", rows=[])
    run(client, "UPDATE qp_indexed SET b = b + 1 WHERE b < (SELECT avg(b) FROM qp_indexed)", error="53000")
    run(client, "DELETE FROM qp_indexed WHERE b < (SELECT avg(b) FROM qp_indexed)", error="53000")
    run(client, "SELECT * FROM qp_indexed ORDER BY a, b", rows=[(1, 10), (1, 30), (2, 20)])
    run(client, "SUBSCRIBE qp_indexed", error="53000")
    run(client, "ALTER CLUSTER quickstart SET (QUERY POLICY = qp_memory)")
    run(client, "SELECT a, sum(b) FROM qp_indexed GROUP BY a ORDER BY a", rows=[(1, 40), (2, 20)])
    run(client, "SELECT * FROM qp_raw LIMIT 10", error="53000")
    run(client, "SELECT a, sum(b) FROM qp_raw GROUP BY a", error="53000")

print("\n## Cluster and role restrictions compose", flush=True)
run(client, "CREATE ROLE qp_agent")
run(admin, "GRANT USAGE ON CLUSTER quickstart TO qp_agent")
run(client, "GRANT SELECT ON qp_indexed, qp_raw TO qp_agent")
run(client, "GRANT USAGE ON QUERY POLICY qp_memory TO qp_agent")
run(client, "ALTER ROLE qp_agent SET query_policy = qp_memory")
run(client, "ALTER CLUSTER quickstart SET (QUERY POLICY = qp_fast)")
agent = psycopg.connect(client.info.dsn, user="qp_agent", autocommit=True)
run(agent, "SELECT a, sum(b) FROM qp_indexed GROUP BY a", error="53000")
run(agent, "SELECT * FROM qp_raw LIMIT 10", error="53000")
agent.close()

print("\n## DDL is exempt", flush=True)
run(client, "CREATE MATERIALIZED VIEW qp_totals AS SELECT a, sum(b) FROM qp_indexed GROUP BY a")
run(client, "CREATE DEFAULT INDEX ON qp_totals")
run(client, "SELECT * FROM qp_totals ORDER BY a", rows=[(1, 40), (2, 20)])

print("\n## Enforcement rollback switch", flush=True)
run(admin, "ALTER SYSTEM SET enable_query_policy_enforcement = false")
run(client, "SELECT a, sum(b) FROM qp_indexed GROUP BY a ORDER BY a", rows=[(1, 40), (2, 20)])
run(client, "SELECT name, mode FROM mz_internal.mz_query_policies ORDER BY name")
run(client, "SELECT name, action, metric, value FROM mz_internal.mz_query_policy_rules ORDER BY name")
run(client, "ALTER ROLE qp_agent RESET query_policy")
run(client, "ALTER CLUSTER quickstart RESET (QUERY POLICY)")
run(client, "DROP QUERY POLICY qp_fast, qp_memory")
print("\nDemo assertions passed.", flush=True)