diff --git a/misc/python/materialize/parallel_workload/action.py b/misc/python/materialize/parallel_workload/action.py index 44d6f9887ab6c..d314cc10db9c9 100644 --- a/misc/python/materialize/parallel_workload/action.py +++ b/misc/python/materialize/parallel_workload/action.py @@ -28,10 +28,22 @@ import materialize.parallel_workload.column from materialize.data_ingest.data_type import ( DATA_TYPES, + DATA_TYPES_FOR_COLUMNS, NUMBER_TYPES, + RANGE_TYPES, + UUID, Boolean, + Bytea, + Char, + IntArray, + IntList, + Jsonb, + Oid, Text, TextTextMap, + Timestamp, + TimestampTz, + VarChar, ) from materialize.data_ingest.query_error import QueryError from materialize.data_ingest.row import Operation @@ -43,6 +55,8 @@ Materialized, ) from materialize.mzcompose.services.minio import minio_blob_uri +from materialize.mzcompose.services.mysql import MySql +from materialize.mzcompose.services.sql_server import SqlServer from materialize.parallel_workload.database import ( DB, MAX_CLUSTER_REPLICAS, @@ -53,13 +67,16 @@ MAX_INDEXES, MAX_KAFKA_SINKS, MAX_KAFKA_SOURCES, + MAX_LOADGEN_SOURCES, MAX_MYSQL_SOURCES, + MAX_NETWORK_POLICIES, MAX_POSTGRES_SOURCES, MAX_ROLES, MAX_ROWS, MAX_SCHEMAS, MAX_SQL_SERVER_SOURCES, MAX_TABLES, + MAX_TYPES, MAX_VIEWS, MAX_WEBHOOK_SOURCES, Cluster, @@ -71,14 +88,18 @@ Index, KafkaSink, KafkaSource, + LoadGeneratorSource, + MultiLoadGeneratorSource, MySqlSource, MzTempSchema, + NetworkPolicy, PostgresSource, Role, S3Object, Schema, SqlServerSource, Table, + Type, View, WebhookSource, ) @@ -139,6 +160,67 @@ def ws_connect(ws: websocket.WebSocket, host, port, user: str) -> tuple[int, int return (ws_conn_id, ws_secret_key) +def untrack_objects_in_schemas(exe: Executor, schemas: set[Schema]) -> None: + """Remove every tracked object living in one of the given schemas. + + Called after a CASCADE drop of a schema or database. Source executor + connections are closed. Cross-schema dependents are cascade-dropped + server-side but not pruned here, they later surface as "does not exist", + which is why the CASCADE actions only run in DDL complexity.""" + with exe.db.lock: + exe.db.tables[:] = [t for t in exe.db.tables if t.schema not in schemas] + exe.db.views[:] = [v for v in exe.db.views if v.schema not in schemas] + # NOTE: Removed in place, like the lists above. DropIndexAction + # discards from this set holding only the index's own lock, so + # rebinding the attribute would resurrect an index that was dropped + # between building the replacement set and assigning it. + exe.db.indexes.difference_update( + {i for i in exe.db.indexes if i.schema in schemas} + ) + # Close the executor connections of the ingestion sources being + # untracked so they don't leak. + connected_sources = ( + exe.db.kafka_sources + + exe.db.postgres_sources + + exe.db.mysql_sources + + exe.db.sql_server_sources + ) + for src in connected_sources: + if src.schema in schemas: + try: + src.executor.mz_conn.close() + except Exception: + pass + exe.db.kafka_sources[:] = [ + s for s in exe.db.kafka_sources if s.schema not in schemas + ] + exe.db.postgres_sources[:] = [ + s for s in exe.db.postgres_sources if s.schema not in schemas + ] + exe.db.mysql_sources[:] = [ + s for s in exe.db.mysql_sources if s.schema not in schemas + ] + exe.db.sql_server_sources[:] = [ + s for s in exe.db.sql_server_sources if s.schema not in schemas + ] + exe.db.loadgen_sources[:] = [ + s for s in exe.db.loadgen_sources if s.schema not in schemas + ] + exe.db.multi_loadgen_sources[:] = [ + s for s in exe.db.multi_loadgen_sources if s.schema not in schemas + ] + exe.db.webhook_sources[:] = [ + s for s in exe.db.webhook_sources if s.schema not in schemas + ] + exe.db.kafka_sinks[:] = [ + s for s in exe.db.kafka_sinks if s.schema not in schemas + ] + exe.db.iceberg_sinks[:] = [ + s for s in exe.db.iceberg_sinks if s.schema not in schemas + ] + exe.db.types[:] = [t for t in exe.db.types if t.schema not in schemas] + + # TODO: CASCADE in DROPs, keep track of what will be deleted class Action: rng: random.Random @@ -201,7 +283,14 @@ def errors_to_ignore(self, exe: Executor) -> list[str]: # catalog apply, which surfaces as an unexpected failure. "non-temporary items cannot depend on temporary item", "is not of expected type SqlColumnType", # TODO: Remove when SQL-566 is fixed + # generate_select_query occasionally emits a WITH MUTUALLY + # RECURSIVE body with ERROR AT RECURSION LIMIT, which errors on + # purpose when the iteration outruns the limit. Only our generated + # queries produce this, so ignoring it globally is safe. + "exceeded the recursion limit", ] + if exe.statement_timeout_set: + result.append("canceling statement due to statement timeout") if exe.db.complexity in (Complexity.DDL, Complexity.DDLOnly): result.extend( [ @@ -214,6 +303,12 @@ def errors_to_ignore(self, exe: Executor) -> list[str]: # matching "unknown cluster replica size" errors. "unknown cluster '", "unknown schema", # schema was dropped + # The database was dropped concurrently: DropDatabaseAction + # can land on an emptied database, and + # DropDatabaseCascadeAction (disabled until SQL-518 is + # fixed) drops non-empty ones. + "unknown database", + "invalid database", # CREATE SCHEMA wording for a vanished database "the transaction's active cluster has been dropped", # cluster was dropped "was removed", # dependency was removed, started with moving optimization off main thread, see database-issues#7285 "real-time source dropped before ingesting the upstream system's visible frontier", # Expected, see https://buildkite.com/materialize/nightly/builds/9399#0191be17-1f4c-4321-9b51-edc4b08b71c5 @@ -285,21 +380,83 @@ def errors_to_ignore(self, exe: Executor) -> list[str]: ) if exe.db.scenario == Scenario.Rename: result.extend(["unknown schema", "ambiguous reference to schema name"]) + # RepeatRow is the only scenario that deliberately produces negative + # multiplicities (`repeat_row` with a negative count), and the + # corruption is observed by any later reader of the affected + # collection, not just by the statement that created it. So the whole + # class is tolerated for the entire scenario, not per action. Every + # other scenario stays strict, a negative-accumulation error there is a + # genuine finding. The central list lives in + # `negative_accumulation_errors.py`. if exe.db.scenario == Scenario.RepeatRow: - # Views that use `repeat_row` with a negative count can surface - # negative-accumulation errors both at DDL time (constant folding) - # and at read time (various operators checking multiplicities). - # The central list lives in `negative_accumulation_errors.py` so - # we can update it in one place when the error messages change. result.extend(NEGATIVE_ACCUMULATION_ERRORS) if materialize.parallel_workload.column.NAUGHTY_IDENTIFIERS: result.extend(["identifier length exceeds 255 bytes"]) return result + # MIN/MAX exist for every type except these (bytea, jsonb, map, list, + # array, uuid, oid, and the range types). + _MINMAX_EXCLUDED = ( + Bytea, + Jsonb, + TextTextMap, + IntList, + IntArray, + UUID, + Oid, + ) + tuple(RANGE_TYPES) + + def aggregate_fns(self, column: Column) -> list[str]: + """Aggregate function templates valid for the column's type. + + Used both in window position (OVER ..) and in GROUP BY position. The + collection aggregates (array_agg/list_agg/jsonb_agg/string_agg) + exercise the "collection" reduce rendering, distinct from the + accumulable sum/count path. Type exclusions are empirically derived, + e.g. array_agg rejects char and cannot nest map/list/array.""" + dt = column.data_type + fns = ["COUNT({})"] + if dt in NUMBER_TYPES: + fns.extend( + [ + "SUM({})", + "AVG({})", + "STDDEV({})", + "STDDEV_POP({})", + "STDDEV_SAMP({})", + "VAR_SAMP({})", + "VAR_POP({})", + ] + ) + if dt == Boolean: + fns.extend(["BOOL_AND({})", "BOOL_OR({})"]) + if dt not in self._MINMAX_EXCLUDED: + fns.extend(["MAX({})", "MIN({})"]) + # Collection aggregates. + fns.append("jsonb_agg({})") + if dt != Char: + fns.append("list_agg({})") + if dt not in (Char, TextTextMap, IntList, IntArray): + fns.append("array_agg({})") + if dt in (Text, VarChar, Char): + fns.append("string_agg({}, ',')") + return fns + def generate_select_query(self, exe: Executor, expr_kind: ExprKind) -> str: - obj = self.rng.choice(exe.db.db_objects()) + objects = exe.db.db_objects() + if not objects: + # A concurrent CASCADE drop removed every object. Nothing to query. + return "SELECT 1" + obj = self.rng.choice(objects) column = self.rng.choice(obj.columns) - obj2 = self.rng.choice(exe.db.db_objects_without_views()) + # db_objects_without_views() can momentarily be empty if a CASCADE drop + # removed the last table/source/MV, leaving only plain views. Fall back + # to obj (only used for an optional join, which is skipped for views). + # The list is taken once and reused below: recomputing it can return an + # empty list even though it was non-empty here, and rng.choice on that + # raises IndexError, which fails the whole run. + objs_without_views = exe.db.db_objects_without_views() + obj2 = self.rng.choice(objs_without_views) if objs_without_views else obj obj_name = str(obj) obj2_name = str(obj2) columns = [ @@ -312,105 +469,219 @@ def generate_select_query(self, exe: Executor, expr_kind: ExprKind) -> str: all_columns = list(obj.columns) + list(obj2.columns) if join else obj.columns - column_types = [] - if self.rng.random() < 0.9: - column_types = [ - self.rng.choice(list(DATA_TYPES)) - for i in range(self.rng.randint(1, 10)) + # Self-contained iterative dataflow, cross joined with a real object. + # The iteration bound and the recursion limit are drawn independently, + # so the limit sometimes fires, exercising the RETURN/ERROR AT paths + # on purpose ("exceeded the recursion limit" is an expected error). + if self.rng.random() < 0.05: + limit_kind = self.rng.choice(["RETURN AT", "ERROR AT"]) + expr = expression( + self.rng.choice(list(DATA_TYPES)), obj.columns, self.rng, expr_kind + ) + return ( + f"WITH MUTUALLY RECURSIVE ({limit_kind} RECURSION LIMIT {self.rng.randint(2, 200)}) " + f"cnt (i int8) AS (" + f"SELECT 1 UNION ALL SELECT i + 1 FROM cnt WHERE i < {self.rng.randint(1, 100)}" + f") SELECT i, {expr} FROM cnt, {obj_name}" + f" LIMIT {self.rng.randint(0, 100)}" + ) + + # Explicit LATERAL derived table correlated on a matching-type column. + # Exercises correlated-subquery decorrelation. + if objs_without_views and self.rng.random() < 0.08: + outer_col = self.rng.choice(obj.columns) + inner_obj = self.rng.choice(objs_without_views) + match_cols = [ + c + for c in inner_obj.columns + if c.data_type == outer_col.data_type and c.data_type != TextTextMap ] - expressions = ", ".join( - [ - expression( - column_type, - all_columns, - self.rng, - expr_kind, - ) - for column_type in column_types - ] + corr = "" + if match_cols: + corr = f" WHERE {self.rng.choice(match_cols)} = {outer_col}" + expr = expression( + self.rng.choice(list(DATA_TYPES)), obj.columns, self.rng, expr_kind + ) + return ( + f"SELECT {expr} FROM {obj_name}, LATERAL (" + f"SELECT count(*) AS cnt FROM {inner_obj}{corr}" + f") AS lat LIMIT {self.rng.randint(0, 100)}" + ) + + # Table function in FROM. unnest of a column reference is an + # implicitly lateral call. + if self.rng.random() < 0.1: + int_list_columns = [c for c in obj.columns if c.data_type == IntList] + if int_list_columns and self.rng.choice([True, False]): + func = f"unnest({self.rng.choice(int_list_columns)})" + else: + func = f"generate_series(1, {self.rng.randint(1, 100)})" + expr = expression( + self.rng.choice(list(DATA_TYPES)), obj.columns, self.rng, expr_kind ) if self.rng.choice([True, False]): - column_types = [] - column1 = self.rng.choice(all_columns) - column2 = self.rng.choice(all_columns) - column3 = self.rng.choice(all_columns) - fns = [ - "COUNT({})", - # "LIST_AGG({})", - # "JSONB_AGG({})", - ] - # if column1.data_type == Text: - # fns.extend(["STRING_AGG({}, ',')"]) - # if column1.data_type not in [TextTextMap, IntArray, IntList]: - # fns.extend(["ARRAY_AGG({})"]) - if column1.data_type in NUMBER_TYPES: - fns.extend( - [ - "SUM({})", - "AVG({})", - "MAX({})", - "MIN({})", - "STDDEV({})", - "STDDEV_POP({})", - "STDDEV_SAMP({})", - "VAR_SAMP({})", - "VAR_POP({})", - ] - ) - elif column1.data_type == Boolean: - fns.extend(["BOOL_AND({})", "BOOL_OR({})"]) - window_fn = self.rng.choice(fns) - expressions += f", {window_fn.format(column1)} OVER (PARTITION BY {column2} ORDER BY {column3})" - else: - expressions = "*" + tf = f"{func} WITH ORDINALITY AS tf(x, ord)" + else: + tf = f"{func} AS tf(x)" + return ( + f"SELECT tf.x, {expr} FROM {obj_name}, {tf}" + f" LIMIT {self.rng.randint(0, 100)}" + ) - query = f"SELECT {expressions} FROM {obj_name}" + star = self.rng.random() >= 0.9 + exprs: list[tuple[type, str]] = [] + if not star: + for i in range(self.rng.randint(1, 10)): + dt = self.rng.choice(list(DATA_TYPES)) + exprs.append((dt, expression(dt, all_columns, self.rng, expr_kind))) if join: + join_kind = self.rng.choice( + ["JOIN", "JOIN", "JOIN", "LEFT JOIN", "RIGHT JOIN", "FULL JOIN"] + ) column2 = self.rng.choice(columns) - query += f" JOIN {obj2_name} ON {column} = {column2}" + join_clause = f" {join_kind} {obj2_name} ON {column} = {column2}" + if self.rng.random() < 0.2: + join_clause += ( + f" AND {expression(Boolean, all_columns, self.rng, expr_kind)}" + ) + else: + join_clause = "" - if self.rng.choice([True, False]): - query += f" WHERE {expression(Boolean, all_columns, self.rng, expr_kind)}" - - if bool(column_types) and self.rng.choice([True, False]): - obj3 = self.rng.choice(exe.db.db_objects()) - obj3_name = str(obj3) - column3 = self.rng.choice(obj3.columns) - obj4 = self.rng.choice(exe.db.db_objects()) - obj4_name = str(obj4) - columns_union = [ - c - for c in obj4.columns - if c.data_type == column3.data_type and c.data_type != TextTextMap + def where_clause() -> str: + parts = [] + if self.rng.choice([True, False]): + parts.append(expression(Boolean, all_columns, self.rng, expr_kind)) + if objs_without_views and self.rng.random() < 0.2: + obj3 = self.rng.choice(objs_without_views) + sub_columns = [ + c + for c in obj3.columns + if c.data_type == column.data_type and c.data_type != TextTextMap + ] + sub_kind = self.rng.choice(["exists", "not exists", "in", "scalar"]) + if sub_kind in ("exists", "not exists"): + cond = expression(Boolean, obj3.columns, self.rng, expr_kind) + parts.append( + f"{'NOT ' if sub_kind == 'not exists' else ''}EXISTS (SELECT 1 FROM {obj3} WHERE {cond})" + ) + elif sub_columns: + sub_column = self.rng.choice(sub_columns) + if sub_kind == "in": + parts.append(f"{column} IN (SELECT {sub_column} FROM {obj3})") + else: + parts.append( + f"{column} = (SELECT {sub_column} FROM {obj3} LIMIT 1)" + ) + # Deliberate temporal filter, mz_now() has to be a top-level + # conjunct for the temporal filter machinery to apply. + ts_columns = [ + c for c in all_columns if c.data_type in (Timestamp, TimestampTz) ] - join_union = ( - obj3_name != obj4_name and obj3 not in exe.db.views and columns_union + if ts_columns and expr_kind != ExprKind.WRITE and self.rng.random() < 0.2: + ts_column = self.rng.choice(ts_columns) + op = self.rng.choice(["<=", ">="]) + parts.append( + f"mz_now() {op} {ts_column} + INTERVAL '{self.rng.randint(0, 100000)} seconds'" + ) + if not parts: + return "" + # Generated Boolean expressions can contain OR, so parenthesize + # each part to keep every part as a top-level conjunct. + return " WHERE " + " AND ".join(f"({part})" for part in parts) + + distinct_on_expr = None + group_by = bool(exprs) and self.rng.random() < 0.2 + if group_by: + num_group = self.rng.randint(1, len(exprs)) + select_list = [expr for _, expr in exprs[:num_group]] + # Only the group keys come from exprs. The aggregates that follow + # get None, their result type is not exprs[i]'s. + select_types: list[type | None] = [dt for dt, _ in exprs[:num_group]] + for i in range(self.rng.randint(1, 3)): + agg_column = self.rng.choice(all_columns) + fn = self.rng.choice(self.aggregate_fns(agg_column)) + select_list.append(fn.format(agg_column)) + select_types.append(None) + group_clause = ( + f" GROUP BY {', '.join(str(i + 1) for i in range(num_group))}" ) - all_columns_union = ( - list(obj3.columns) + list(obj4.columns) if join_union else obj3.columns + # Clause order is fixed: HAVING before OPTIONS. + if self.rng.random() < 0.3: + group_clause += f" HAVING count(*) >= {self.rng.randint(0, 10)}" + if self.rng.random() < 0.3: + group_clause += f" OPTIONS (AGGREGATE INPUT GROUP SIZE = {self.rng.randint(1, 1000)})" + distinct_clause = "" + else: + select_list = [expr for _, expr in exprs] if exprs else ["*"] + # `*` expands to every column, so the output types are the columns'. + select_types = ( + [dt for dt, _ in exprs] if exprs else [c.data_type for c in all_columns] ) - expressions3 = ", ".join( - [ - expression( - column_type, - all_columns_union, - self.rng, - expr_kind, - ) - for column_type in column_types - ] + group_clause = "" + distinct_clause = "" + distinct_on_expr = None + if exprs and self.rng.random() < 0.15: + if exprs[0][0] == TextTextMap or self.rng.choice([True, False]): + distinct_clause = "DISTINCT " + else: + # The DISTINCT ON expressions must be a prefix of the + # ORDER BY expressions. The leading ORDER BY item repeats + # this expression verbatim (not its ordinal) so the match + # holds even for bare literals. + distinct_on_expr = exprs[0][1] + distinct_clause = f"DISTINCT ON ({distinct_on_expr}) " + if exprs and self.rng.choice([True, False]): + column1 = self.rng.choice(all_columns) + column2 = self.rng.choice(all_columns) + column3 = self.rng.choice(all_columns) + window_fn = self.rng.choice(self.aggregate_fns(column1)) + select_list.append( + f"{window_fn.format(column1)} OVER (PARTITION BY {column2} ORDER BY {column3})" + ) + select_types.append(None) + + expressions = ", ".join(select_list) + query = f"SELECT {distinct_clause}{expressions} FROM {obj_name}" + query += join_clause + query += where_clause() + query += group_clause + + if not distinct_clause and self.rng.choice([True, False]): + set_op = self.rng.choice( + ["UNION ALL"] * 4 + + ["UNION", "INTERSECT", "INTERSECT ALL", "EXCEPT", "EXCEPT ALL"] ) - query += f" UNION ALL SELECT {expressions3} FROM {obj3_name}" - - if join_union: - column4 = self.rng.choice(columns_union) - query += f" JOIN {obj4_name} ON {column3} = {column4}" - - if self.rng.choice([True, False]): - query += f" WHERE {expression(Boolean, all_columns_union, self.rng, expr_kind)}" + query += f" {set_op} SELECT {expressions} FROM {obj_name}" + query += join_clause + query += where_clause() + query += group_clause + + # Ordinals keep ORDER BY valid across set operations. Map-typed + # output columns are skipped, same as in the join column selection. + # NOTE: select_types is indexed by output position, which is not + # exprs' indexing once a GROUP BY appends aggregates after the keys. + orderable = [i + 1 for i, dt in enumerate(select_types) if dt != TextTextMap] + if distinct_on_expr is not None: + order_by = [distinct_on_expr] + for i in self.rng.sample(orderable, self.rng.randint(0, len(orderable))): + if i != 1: + order_by.append(str(i)) + query += f" ORDER BY {', '.join(order_by)}" + elif exprs and orderable and self.rng.random() < 0.3: + order_by = [] + for i in self.rng.sample(orderable, self.rng.randint(1, len(orderable))): + direction = self.rng.choice(["", " ASC", " DESC"]) + nulls = self.rng.choice(["", " NULLS FIRST", " NULLS LAST"]) + order_by.append(f"{i}{direction}{nulls}") + query += f" ORDER BY {', '.join(order_by)}" query += f" LIMIT {self.rng.randint(0, 100)}" + if self.rng.random() < 0.2: + query += f" OFFSET {self.rng.randint(0, 100)}" + + if self.rng.random() < 0.15: + query = f"WITH cte0 AS ({query}) SELECT * FROM cte0" return query def exe_prepared(self, query: str, stmt_name: str, exe: Executor) -> None: @@ -460,7 +731,13 @@ def run(self, exe: Executor) -> bool: if self.rng.choice([True, False]) else exe.commit(http=Http.NO) ) + # NOTE: A bounded SUBSCRIBE (UP TO) over an object whose as_of has + # advanced to the end of time (e.g. a finished bounded load generator + # source) soft-panics the optimizer + # (https://linear.app/materializeinc/issue/CLU-169). Left out until that + # is fixed. AS OF AT LEAST 0 below is safe (empty until). query = "SUBSCRIBE " + envelope_used = False if self.rng.choice([True, False]): obj = self.rng.choice(exe.db.db_objects()) query += f"{obj}" @@ -470,12 +747,21 @@ def run(self, exe: Executor) -> bool: columns = self.rng.sample(obj.columns, len(obj.columns)) key = ", ".join(column.name(True) for column in columns) query += f" ENVELOPE {envelope} (KEY ({key}))" - - if self.rng.choice([True, False]): - query += " WITH (SNAPSHOT = false)" + envelope_used = True else: query += f"({self.generate_select_query(exe, ExprKind.MATERIALIZABLE)})" + options = [] + if self.rng.choice([True, False]): + options.append("SNAPSHOT = false") + if not envelope_used and self.rng.random() < 0.3: + options.append("PROGRESS") + if options: + query += f" WITH ({', '.join(options)})" + if self.rng.random() < 0.2: + # AT LEAST always plans, no matter how far the since advanced. + query += " AS OF AT LEAST 0" + exe.execute(f"DECLARE c{self.i} CURSOR FOR {query}", http=Http.NO) while True: rows = self.rng.choice(["ALL", self.rng.randrange(1000)]) @@ -690,6 +976,19 @@ def errors_to_ignore(self, exe: Executor) -> list[str]: # roundtrip can produce NULLs for NOT NULL columns. "violates not-null constraint", "timeout: error trying to connect", + # COPY TO CSV writes a large-year date that COPY FROM CSV then + # fails to parse back. + # See https://linear.app/materializeinc/issue/SS-345 + "expected_dur_like_tokens can only be called with", + # TODO: Reenable when SS-361 is fixed. A COPY FROM without a + # column list plans every target column as its DEFAULT, and + # that projection only acts as identity by accident. Expression + # memoization collapses two identical default literals, so any + # table with two same-type nullable columns (their defaults are + # identical typed nulls) decodes a shifted row. Shifts that are + # type-incompatible fail here, type-compatible ones are + # inserted silently. + "failed to decode Row from a record batch", ] ) if exe.db.complexity == Complexity.DDL: @@ -912,6 +1211,17 @@ def run(self, exe: Executor) -> bool: class InsertReturningAction(Action): + def errors_to_ignore(self, exe: Executor) -> list[str]: + result = super().errors_to_ignore(exe) + # The RETURNING expressions re-render the fully-qualified table and + # column names, so a concurrent schema or table rename landing between + # the INSERT target and the RETURNING clause leaves the two referring to + # different names and the query fails with "does not exist". Same class + # as UpdateAction/DeleteAction. + if exe.db.complexity == Complexity.DDL or exe.db.scenario == Scenario.Rename: + result.extend(["does not exist"]) + return result + def run(self, exe: Executor) -> bool: table = None if exe.insert_table is not None: @@ -975,6 +1285,39 @@ def run(self, exe: Executor) -> bool: return True +class CopyToStdoutAction(Action): + def errors_to_ignore(self, exe: Executor) -> list[str]: + result = super().errors_to_ignore(exe) + result.extend( + [ + "in the same timedomain", + # A prior statement in the read transaction routed it to + # mz_catalog_server, where this COPY of a user object is + # rejected. + 'is not allowed from the "mz_catalog_server" cluster', + # BINARY format is deliberately kept in the mix even though + # some types (e.g. map) have no binary output function. + "no binary output function available for type", + ] + ) + if exe.db.complexity == Complexity.DDL: + result.extend( + [ + "does not exist", + ] + ) + return result + + def run(self, exe: Executor) -> bool: + obj = self.rng.choice(exe.db.db_objects()) + query = f"COPY (SELECT * FROM {obj} LIMIT {self.rng.randint(0, 100)}) TO STDOUT" + if self.rng.choice([True, False]): + format = self.rng.choice(["TEXT", "CSV", "BINARY"]) + query += f" WITH (FORMAT {format})" + exe.copy_to_stdout(query) + return True + + class SourceInsertAction(Action): def run(self, exe: Executor) -> bool: with exe.db.lock: @@ -1051,8 +1394,14 @@ def run(self, exe: Executor) -> bool: return False table = self.rng.choice(tables) - column2 = self.rng.choice(table.columns) - query = f"UPDATE {table} SET {column2.name(True)} = {expression(column2.data_type, table.columns, self.rng, kind=ExprKind.WRITE)} WHERE {expression(Boolean, table.columns, self.rng, kind=ExprKind.WRITE)}" + set_columns = self.rng.sample( + table.columns, self.rng.randint(1, len(table.columns)) + ) + set_clause = ", ".join( + f"{c.name(True)} = {expression(c.data_type, table.columns, self.rng, kind=ExprKind.WRITE)}" + for c in set_columns + ) + query = f"UPDATE {table} SET {set_clause} WHERE {expression(Boolean, table.columns, self.rng, kind=ExprKind.WRITE)}" if self.rng.choice([True, False]): self.stmt_id += 1 self.exe_prepared(query, f"update{self.stmt_id}", exe) @@ -1115,7 +1464,28 @@ def run(self, exe: Executor) -> bool: return False table = self.rng.choice(tables) query = f"DELETE FROM {table}" - if self.rng.random() < 0.95: + using_tables = [ + t + for t in exe.db.tables + if t != table and (not t.temp or t in exe.temp_objects) + ] + # TODO: Drop the RepeatRow gate once database-issues#9308 is fixed. + # DELETE .. USING lowers to a semijoin whose DistinctBy can leave the + # target table with a net-negative row, and every later reader of that + # table then surfaces the corruption. Tolerating that class outside + # RepeatRow would blind the DML and DDL complexities to every genuine + # negative-accumulation finding, so the variant only runs in the one + # scenario that already expects negative multiplicities. + if ( + using_tables + and exe.db.scenario == Scenario.RepeatRow + and self.rng.random() < 0.2 + ): + using_table = self.rng.choice(using_tables) + all_columns = list(table.columns) + list(using_table.columns) + query += f" USING {using_table}" + query += f" WHERE {expression(Boolean, all_columns, self.rng, kind=ExprKind.WRITE)}" + elif self.rng.random() < 0.95: query += f" WHERE {expression(Boolean, table.columns, self.rng, kind=ExprKind.WRITE)}" if self.rng.choice([True, False]): self.stmt_id += 1 @@ -1133,15 +1503,83 @@ def run(self, exe: Executor) -> bool: class CommentAction(Action): - def run(self, exe: Executor) -> bool: - table = self.rng.choice(exe.db.tables) + def errors_to_ignore(self, exe: Executor) -> list[str]: + return [ + "unknown role", + ] + super().errors_to_ignore(exe) - if self.rng.choice([True, False]): - column = self.rng.choice(table.columns) - query = f"COMMENT ON COLUMN {column} IS '{Text.random_value(self.rng)}'" - else: - query = f"COMMENT ON TABLE {table} IS '{Text.random_value(self.rng)}'" + def run(self, exe: Executor) -> bool: + # Only the snapshot needs the global lock. Rendering the candidate + # names is O(tracked objects) with a naughtify each, and holding the + # lock across that serializes every other worker behind this action. + with exe.db.lock: + tables = list(exe.db.tables) + views = list(exe.db.views) + direct_sources = exe.db.kafka_sources + exe.db.webhook_sources + table_sources = ( + exe.db.postgres_sources + + exe.db.mysql_sources + + exe.db.sql_server_sources + ) + loadgen_sources = list(exe.db.loadgen_sources) + sinks = exe.db.kafka_sinks + exe.db.iceberg_sinks + indexes = list(exe.db.indexes) + schemas = list(exe.db.schemas) + dbs = list(exe.db.dbs) + clusters = list(exe.db.clusters) + roles = list(exe.db.roles) + + candidates: list[tuple[str, str]] = [] + for table in tables: + candidates.append(("TABLE", str(table))) + # Columns are only ever appended, never removed, so picking one + # off the snapshot cannot index past the end. + candidates.append(("COLUMN", str(self.rng.choice(table.columns)))) + for view in views: + candidates.append( + ( + "MATERIALIZED VIEW" if view.materialized else "VIEW", + str(view), + ) + ) + # Kafka and webhook source objects are readable directly, the + # others follow the source-table model: str(obj) names the table, + # the ingestion source is a separate catalog item. + for source in direct_sources: + candidates.append(("SOURCE", str(source))) + for source in table_sources: + candidates.append(("TABLE", str(source))) + candidates.append( + ( + "SOURCE", + f"{source.schema}.{identifier(source.executor.source)}", + ) + ) + for source in loadgen_sources: + candidates.append(("TABLE", str(source))) + candidates.append( + ("SOURCE", f"{source.schema}.{identifier(source.source_name())}") + ) + for sink in sinks: + candidates.append(("SINK", str(sink))) + for index in indexes: + candidates.append(("INDEX", str(index))) + for schema in schemas: + candidates.append(("SCHEMA", str(schema))) + for db in dbs: + candidates.append(("DATABASE", str(db))) + for cluster in clusters: + candidates.append(("CLUSTER", str(cluster))) + for role in roles: + candidates.append(("ROLE", str(role))) + candidates.append(("SECRET", "materialize.public.pgpass")) + candidates.append(("CONNECTION", "materialize.public.kafka_conn")) + if not candidates: + return False + kind, name = self.rng.choice(candidates) + comment = self.rng.choice([f"'{Text.random_value(self.rng)}'", "NULL"]) + query = f"COMMENT ON {kind} {name} IS {comment}" exe.execute(query, http=Http.RANDOM) return True @@ -1196,10 +1634,12 @@ def run(self, exe: Executor) -> bool: # The indexed object or its schema may have been dropped # concurrently, taking the index with it. Untrack the index # either way so stale entries don't fill up the set and choke - # off CreateIndexAction. - exe.db.indexes.remove(index) + # off CreateIndexAction. Use discard, not remove: a concurrent + # CASCADE drop's untrack_objects_in_schemas may have already + # removed it, and remove would raise KeyError. + exe.db.indexes.discard(index) raise - exe.db.indexes.remove(index) + exe.db.indexes.discard(index) return True @@ -1263,7 +1703,12 @@ def run(self, exe: Executor) -> bool: query = f"DROP TABLE {table}" exe.execute(query, http=Http.RANDOM) - exe.db.tables.remove(table) + # A concurrent CASCADE drop's untrack_objects_in_schemas may have + # already filtered this table out of the list, tolerate that. + try: + exe.db.tables.remove(table) + except ValueError: + pass return True @@ -1425,11 +1870,28 @@ def run(self, exe: Executor) -> bool: return True -class ReplaceMaterializedViewAction(Action): +class CreateReplacementMaterializedViewAction(Action): + """CREATE REPLACEMENT MATERIALIZED VIEW for a random materialized view. + + Deliberately only creates. The replacement stays alive until an + ApplyReplacementMaterializedViewAction or + DropReplacementMaterializedViewAction finds it in the catalog, from any + worker and possibly only after a kill or 0dt deploy. Resolving the + replacement on the creating connection would tie its lifetime to that + connection, so it could never span an envd restart. Incident 1136 + (bootstrap loses the replacement collection's primary link to the + target's shard, a post-restart DROP of the replacement then finalizes the + target's live shard) lives exactly in that gap. + """ + def errors_to_ignore(self, exe: Executor) -> list[str]: errors = [ - # A concurrent or leaked replacement of the same view + # A concurrent replacement of the same view "because it already has a replacement", + # The target's query constant-folded to a plan without live + # inputs and completed, which can_seal cannot know. The + # SealedCollectionCheckAction oracle arbitrates whether a sealed + # shard is legitimate. "is sealed and thus cannot be replaced", ] + super().errors_to_ignore(exe) if exe.db.scenario == Scenario.Rename: @@ -1440,7 +1902,15 @@ def errors_to_ignore(self, exe: Executor) -> list[str]: def run(self, exe: Executor) -> bool: with exe.db.lock: - mvs = [v for v in exe.db.views if v.materialized] + # Skip views whose shard is known to seal legitimately. Sealing + # is transitive: REFRESH AT views seal after their last refresh, + # repeat_row constant views seal on hydration, and any view that + # reads from a sealing input seals too (View.can_seal). A sealed + # target cannot be replaced, so skipping these avoids wasted + # work. can_seal cannot know about constant-folded queries + # though, so "is sealed" errors are still ignored and + # SealedCollectionCheckAction arbitrates damage. + mvs = [v for v in exe.db.views if v.materialized and not v.can_seal] if not mvs: return False view = self.rng.choice(mvs) @@ -1456,17 +1926,298 @@ def run(self, exe: Executor) -> bool: exe.execute( f"CREATE REPLACEMENT MATERIALIZED VIEW {tmp_mv} FOR {view} AS {view.get_select()}", ) - time.sleep(self.rng.random()) + return True + + +class ResolveReplacementMaterializedViewAction(Action): + """Base for actions resolving live replacements. + + Replacements are picked from the catalog instead of tracked state so that + replacements created by other workers or left behind by dead connections, + kills, and 0dt deploys are found too. + """ + + def pick_replacement(self, exe: Executor) -> tuple[str, str] | None: + exe.execute( + "SELECT rd.name, rs.name, r_mv.name, td.name, ts.name, t_mv.name" + " FROM mz_internal.mz_replacements r" + " JOIN mz_catalog.mz_materialized_views r_mv ON r_mv.id = r.id" + " JOIN mz_catalog.mz_schemas rs ON rs.id = r_mv.schema_id" + " JOIN mz_catalog.mz_databases rd ON rd.id = rs.database_id" + " JOIN mz_catalog.mz_materialized_views t_mv ON t_mv.id = r.target_id" + " JOIN mz_catalog.mz_schemas ts ON ts.id = t_mv.schema_id" + " JOIN mz_catalog.mz_databases td ON td.id = ts.database_id", + http=Http.NO, + ) + rows = exe.cur.fetchall() + if not rows: + return None + r_db, r_schema, r_name, t_db, t_schema, t_name = self.rng.choice(rows) + replacement = f"{identifier(r_db)}.{identifier(r_schema)}.{identifier(r_name)}" + target = f"{identifier(t_db)}.{identifier(t_schema)}.{identifier(t_name)}" + return replacement, target + + def errors_to_ignore(self, exe: Executor) -> list[str]: + return [ + # Reading catalog relations inside a txn that already touched + # user objects crosses timedomains. + "in the same timedomain", + ] + super().errors_to_ignore(exe) + + +class ApplyReplacementMaterializedViewAction(ResolveReplacementMaterializedViewAction): + def errors_to_ignore(self, exe: Executor) -> list[str]: + errors = [ + # A concurrent apply or drop won the race for this replacement + "replacement materialized view does not exist", + # The target's shard sealed, which is legitimate for targets + # whose query constant-folded to a plan without live inputs. + # SealedCollectionCheckAction arbitrates whether a sealed shard + # is damage, so it is not treated as a failure here. + "is sealed and thus cannot be replaced", + ] + super().errors_to_ignore(exe) + if exe.db.scenario == Scenario.Rename: + # The target was renamed between picking and applying + errors += ["does not exist"] + return errors + + def run(self, exe: Executor) -> bool: + picked = self.pick_replacement(exe) + if picked is None: + return False + replacement, target = picked + # APPLY REPLACEMENT waits for the target's write frontier to catch up + # to the replacement's. The wait never finishes if the target is + # concurrently cascade-dropped (database-issues#9820), so bound it + # with a statement timeout. + exe.execute("SET statement_timeout = '60s'") + # A timeout is expected for as long as this one is configured. + exe.statement_timeout_set = True try: - exe.execute(f"ALTER MATERIALIZED VIEW {view} APPLY REPLACEMENT {tmp_mv}") - except QueryError: - # Clean up, a leaked replacement blocks all future replacements - # of this view. + exe.execute( + f"ALTER MATERIALIZED VIEW {target} APPLY REPLACEMENT {replacement}" + ) + finally: try: - exe.execute(f"DROP MATERIALIZED VIEW IF EXISTS {tmp_mv}") + exe.execute("RESET statement_timeout") except QueryError: + # The timeout is still configured, so keep tolerating it. + pass + else: + # Back at the default timeout, so a timeout is a genuine + # finding again. The worker's statements are sequential, + # nothing from the window above can still be in flight. + exe.statement_timeout_set = False + return True + + +class DropReplacementMaterializedViewAction(ResolveReplacementMaterializedViewAction): + """Drop a live replacement without applying it. + + This is the incident 1136 trigger statement when it runs after an envd + restart: bootstrap loses the replacement collection's primary link, so + the drop wrongly finalizes the target's live shard. + SealedCollectionCheckAction detects the resulting damage. + """ + + def run(self, exe: Executor) -> bool: + picked = self.pick_replacement(exe) + if picked is None: + return False + replacement, _ = picked + # IF EXISTS, a concurrent apply or drop may have resolved it already. + exe.execute(f"DROP MATERIALIZED VIEW IF EXISTS {replacement}") + return True + + +def plan_user_input_ids(plan: str) -> list[str]: + """User collection ids that an `EXPLAIN .. AS JSON` dataflow plan reads. + + Raises on a plan whose shape it cannot read, rather than returning an empty + list. Constant folding legitimately produces plans that read nothing, so + "reads nothing" and "no longer understands the plan" are indistinguishable + to a caller, and treating the second as the first retires whatever the + caller checks without a trace. Only accepts dataflow (MIR) plans, a + fast-path peek plan carries its inputs elsewhere and is rejected. + """ + try: + parsed = json.loads(plan) + except json.JSONDecodeError as e: + raise ValueError(f"EXPLAIN .. AS JSON output does not parse ({e}): {plan}") + + # `Constant` and `Get` are the only leaves of MirRelationExpr, so a plan + # with neither is not a shape this can read. Global ids outside `Get` are + # not inputs, e.g. the index an `access_strategy` picked. + get_ids: list = [] + constant = False + + def walk(node: object) -> None: + nonlocal constant + if isinstance(node, dict): + get = node.get("Get") + if isinstance(get, dict) and "id" in get: + get_ids.append(get["id"]) + if "Constant" in node: + constant = True + for value in node.values(): + walk(value) + elif isinstance(node, list): + for value in node: + walk(value) + + walk(parsed) + if not get_ids and not constant: + raise ValueError(f"plan has no Get or Constant leaf: {plan}") + + user_ids = set() + for get_id in get_ids: + match get_id: + case {"Global": {"User": int(gid)}}: + user_ids.add(f"u{gid}") + case {"Local": _} | { + "Global": {"System": _} + | {"Transient": _} + | {"IntrospectionSourceIndex": _} + }: + # A Let binding or a non-user namespace, neither is a user + # collection with a frontier to check. pass + case _: + raise ValueError(f"unrecognized Get id {get_id!r} in plan: {plan}") + return sorted(user_ids) + + +class SealedCollectionCheckAction(Action): + """Client-side data-loss oracle for wrongly finalized persist shards. + + A shard's upper only becomes empty when the shard is sealed. That is + legitimate when the writing dataflow completed (a constant plan, inputs + that all sealed, a REFRESH AT view past its last refresh, a bounded load + generator that finished) and destructive when shard finalization ran + against a live collection. Incident 1136 is the known destructive + trigger: dropping a replacement after an envd restart finalizes the + target's shard because bootstrap dropped the replacement collection's + primary link. Tables are covered too because versioned tables + (ALTER TABLE ADD COLUMN) carry the same primary ownership chain. + + A live user table or materialized view with an empty write frontier and + at least one live plan input has no legitimate explanation and is + reported as damage. This oracle is the single arbiter of wrongly sealed + shards, the replacement actions ignore "is sealed" errors instead of + treating them as failures. + """ + + def applicable(self, exe: Executor) -> bool: + # repeat_row constant views seal legitimately on hydration and are + # indistinguishable from damage in the catalog. + return exe.db.scenario != Scenario.RepeatRow + + def errors_to_ignore(self, exe: Executor) -> list[str]: + return [ + "in the same timedomain", + ] + super().errors_to_ignore(exe) + + def completed_legitimately( + self, exe: Executor, database: str, schema: str, name: str + ) -> bool: + """Whether a sealed materialized view's dataflow completed on its own. + + A dataflow completes, sealing its shard, once it can never produce + more output: its optimized plan reads no storage collections at all + (constant folding, e.g. a join ON false, drops inputs the catalog + still records as dependencies), or every storage collection it reads + is itself sealed. Only a sealed view with at least one live plan + input is damage. + """ + try: + exe.execute( + "EXPLAIN OPTIMIZED PLAN AS JSON FOR MATERIALIZED VIEW" + f" {identifier(database)}.{identifier(schema)}.{identifier(name)}", + http=Http.NO, + ) + plan = "\n".join(str(row[0]) for row in exe.cur.fetchall()) + except QueryError as e: + if "unknown catalog item" in str(e) or "does not exist" in str(e): + # Dropped concurrently, so its shard is allowed to seal. + return True raise + input_ids = plan_user_input_ids(plan) + if not input_ids: + return True + id_list = ", ".join(f"'{i}'" for i in input_ids) + exe.execute( + "SELECT count(*) FROM mz_internal.mz_frontiers" + f" WHERE object_id IN ({id_list}) AND write_frontier IS NOT NULL", + http=Http.NO, + ) + return exe.cur.fetchall()[0][0] == 0 + + def run(self, exe: Executor) -> bool: + # Sealing is transitive: a view reading from a REFRESH AT view seals + # once that input seals, even though its own refresh strategy is not + # 'at', and likewise for anything reading a bounded load generator's + # table. Walk the dependency graph up from every REFRESH AT view and + # every load generator source and exclude everything that transitively + # depends on one. The COUNTER sources are UP TO bounded and seal when + # they finish. The multi-subsource generators never seal, so sweeping + # them in only widens the exclusion, erring toward tolerating. repeat_row + # constant views also seal transitively but are not marked in the + # catalog, which is why this oracle is disabled in the RepeatRow + # scenario, see applicable. Plain user tables never seal and need no + # exclusions, and source-fed tables are covered by the walk through + # their source. + def sealed_collections() -> list: + exe.execute( + "WITH MUTUALLY RECURSIVE" + " sealing(id text) AS (" + "SELECT rs.materialized_view_id" + " FROM mz_internal.mz_materialized_view_refresh_strategies rs" + " WHERE rs.type = 'at'" + " UNION" + " SELECT s.id" + " FROM mz_catalog.mz_sources s" + " WHERE s.type = 'load-generator'" + " UNION" + " SELECT d.object_id" + " FROM mz_internal.mz_object_dependencies d" + " JOIN sealing s ON d.referenced_object_id = s.id)" + " SELECT o.id, d.name, sc.name, o.name, o.type" + " FROM mz_catalog.mz_objects o" + " JOIN mz_catalog.mz_schemas sc ON sc.id = o.schema_id" + " JOIN mz_catalog.mz_databases d ON d.id = sc.database_id" + " JOIN mz_internal.mz_frontiers f ON f.object_id = o.id" + " WHERE o.type IN ('materialized-view', 'table')" + " AND o.id LIKE 'u%'" + " AND f.write_frontier IS NULL" + " AND o.id NOT IN (SELECT id FROM sealing)", + http=Http.NO, + ) + return exe.cur.fetchall() + + sealed = sealed_collections() + if not sealed: + return True + # The catalog relations and mz_frontiers are not updated atomically, + # so a query timestamp can land inside a concurrent CREATE or DROP + # where the object is cataloged but its frontier reads as the empty + # antichain. A wrongly finalized shard is permanent, so only report + # objects that are still sealed on a re-check. + time.sleep(5) + sealed_ids = {row[0] for row in sealed} + damaged = [ + row + for row in sealed_collections() + if row[0] in sealed_ids + and not ( + row[4] == "materialized-view" + and self.completed_legitimately(exe, row[1], row[2], row[3]) + ) + ] + if damaged: + raise ValueError( + "Collections with wrongly sealed persist shards" + f" (incident 1136 class): {damaged}" + ) return True @@ -1495,13 +2246,13 @@ def run(self, exe: Executor) -> bool: # Iceberg sinks always have a key, so only allow a conservative # case: all names, types, and nullabilities match, which also # guarantees the key columns exist in the new object. - # TODO: Switch back when SS-324 is fixed to make sure it errors + # TODO: Switch back when SS-344 is fixed to make sure it errors # instead of causing a stall objs = [] old_cols = { c.name(True): (c.data_type, c.nullable) for c in old_object.columns } - for o in exe.db.db_objects_without_views(): + for o in exe.db.db_objects_for_sinks(): if isinstance(old_object, WebhookSource): continue if isinstance(o, WebhookSource): @@ -1557,14 +2308,14 @@ def run(self, exe: Executor) -> bool: # single column formats objs = [ o - for o in exe.db.db_objects_without_views() + for o in exe.db.db_objects_for_sinks() if len(o.columns) == 1 and o.columns[0].data_type == old_object.columns[0].data_type ] elif sink.format in ["FORMAT JSON"]: # We should be able to format all data types as JSON, and they have no # particular backwards-compatiblility requirements. - objs = [o for o in exe.db.db_objects_without_views()] + objs = [o for o in exe.db.db_objects_for_sinks()] else: # Avro schema migration checking can be quite strict, and we need to be not only # compatible with the latest object's schema but all previous schemas. @@ -1577,7 +2328,7 @@ def run(self, exe: Executor) -> bool: old_cols = { c.name(True): (c.data_type, c.nullable) for c in old_object.columns } - for o in exe.db.db_objects_without_views(): + for o in exe.db.db_objects_for_sinks(): if isinstance(old_object, WebhookSource): continue if isinstance(o, WebhookSource): @@ -1805,43 +2556,241 @@ def run(self, exe: Executor) -> bool: return True -class CommitRollbackAction(Action): - def run(self, exe: Executor) -> bool: - if not exe.action_run_since_last_commit_rollback: - return False +class ParameterizedQueryAction(Action): + """PREPARE a query with $1..$n placeholders, then EXECUTE it with values. - if self.rng.random() < 0.7: - exe.commit() - else: - exe.rollback() - exe.action_run_since_last_commit_rollback = False + Exercises the parameter-type-inference and bind/assignment-cast path that + the workload's other prepared statements (no parameters) never reach.""" + + def errors_to_ignore(self, exe: Executor) -> list[str]: + result = super().errors_to_ignore(exe) + result.extend( + [ + "in the same timedomain", + 'is not allowed from the "mz_catalog_server" cluster', + ] + ) + if exe.db.complexity == Complexity.DDL: + result.extend(["does not exist"]) + return result + + def run(self, exe: Executor) -> bool: + obj = self.rng.choice(exe.db.db_objects()) + n = self.rng.randint(1, 4) + # Record and record list are expression-only pseudo types, a + # parameter cast to them fails with "cannot reference pseudo type". + param_types = [self.rng.choice(list(DATA_TYPES_FOR_COLUMNS)) for _ in range(n)] + projection = ", ".join( + f"${i + 1}::{t.name()}" for i, t in enumerate(param_types) + ) + self.stmt_id += 1 + name = f"pq{self.stmt_id}" + query = f"SELECT {projection} FROM {obj} LIMIT {self.rng.randint(0, 10)}" + # Each argument is cast to its parameter's declared type so the + # assignment cast on EXECUTE always succeeds (e.g. a bytea parameter + # rejects a bare text literal). + values = ", ".join( + f"({t.random_value(self.rng, in_query=True)})::{t.name()}" + for t in param_types + ) + # Run sequentially, not in a try/finally: if EXECUTE fails it aborts + # the transaction, and a DEALLOCATE in a finally would then fail with + # "current transaction is aborted", masking the real error. On failure + # the statement leaks on the session, which is harmless: statement + # names are never reused, and the session's prepared statements die + # with it. + exe.execute(f"PREPARE {name} AS {query}", http=Http.NO) + exe.execute(f"EXECUTE {name} ({values})", http=Http.NO, fetch=True) + exe.execute(f"DEALLOCATE {name}", http=Http.NO) return True -class FlipFlagsAction(Action): - def __init__( - self, - rng: random.Random, - composition: Composition | None, - ): - super().__init__(rng, composition) +class BoundedStalenessReadAction(Action): + """A read under `bounded staleness` isolation, a distinct timestamp- + selection path that never blocks and returns 40001 when the freshness + bound cannot be met. The isolation is set transiently around the read and + restored afterwards, since bounded staleness is read-only and would break + writes.""" - BOOLEAN_FLAG_VALUES = ["TRUE", "FALSE"] + def applicable(self, exe: Executor) -> bool: + return exe.db.flags.get("enable_bounded_staleness_isolation", "FALSE") == "TRUE" - self.flags_with_values: dict[str, list[str]] = dict() - self.flags_with_values["persist_blob_target_size"] = ( - # 1 MiB, 16 MiB, 128 MiB - ["1048576", "16777216", "134217728"] - ) - for flag in ["catalog", "source", "snapshot", "txn"]: - self.flags_with_values[f"persist_use_critical_since_{flag}"] = ( - BOOLEAN_FLAG_VALUES - ) - self.flags_with_values["persist_claim_unclaimed_compactions"] = ( - BOOLEAN_FLAG_VALUES + def errors_to_ignore(self, exe: Executor) -> list[str]: + result = super().errors_to_ignore(exe) + result.extend( + [ + # The flag was flipped off between applicable() and run(). + "is not available", + # The freshness bound could not be met. Bounded staleness + # never blocks, it errors instead. + "not been materialized", + "could not find a valid timestamp for the query", + "cannot serve query under bounded staleness", + # A leaked real_time_recency SET on the session (its own reset + # was discarded on a prior query's error path) conflicts with + # bounded staleness. + "cannot be combined with bounded staleness", + "in the same timedomain", + 'is not allowed from the "mz_catalog_server" cluster', + ] ) - self.flags_with_values["persist_optimize_ignored_data_fetch"] = ( - BOOLEAN_FLAG_VALUES + if exe.db.complexity == Complexity.DDL: + result.extend(["does not exist"]) + return result + + def run(self, exe: Executor) -> bool: + bound = self.rng.choice(["1s", "5s", "30s"]) + restore = f"SET TRANSACTION_ISOLATION TO '{exe.isolation}'" + exe.execute( + f"SET TRANSACTION_ISOLATION TO 'bounded staleness {bound}'", + explainable=False, + ) + try: + query = self.generate_select_query(exe, ExprKind.ALL) + exe.execute(query, http=Http.NO, fetch=True) + except QueryError: + # The restore must not raise on top of the read's error, e.g. on a + # killed or cancelled session, that error would replace the read's + # and hide the real failure. A leaked bounded staleness isolation + # makes later actions fail with errors only this action ignores, so + # a failed restore forces a reconnect rather than being swallowed. + try: + exe.execute(restore, explainable=False) + except QueryError: + exe.reconnect_next = True + raise + # No error is being masked here, so a failing restore is a genuine + # failure and propagates. + exe.execute(restore, explainable=False) + return True + + +class ReadOnlyTransactionAction(Action): + """A multi-statement `BEGIN READ ONLY` transaction. All reads run at one + pinned timestamp (one timedomain), and holding the read pins compaction + for the transaction's lifetime.""" + + def errors_to_ignore(self, exe: Executor) -> list[str]: + result = super().errors_to_ignore(exe) + result.extend( + [ + "in the same timedomain", + 'is not allowed from the "mz_catalog_server" cluster', + ] + ) + if exe.db.complexity == Complexity.DDL: + result.extend(["does not exist"]) + return result + + def run(self, exe: Executor) -> bool: + isolation = self.rng.choice( + [ + "", + " ISOLATION LEVEL SERIALIZABLE", + " ISOLATION LEVEL STRICT SERIALIZABLE", + ] + ) + exe.execute(f"BEGIN{isolation} READ ONLY", http=Http.NO) + try: + for _ in range(self.rng.randint(1, 3)): + query = self.generate_select_query(exe, ExprKind.ALL) + exe.execute(query, http=Http.NO, fetch=True) + except QueryError: + # The transaction still has to end, an open one wedges every later + # action on this session, but ending it must not raise over the + # read's error and hide the real failure. COMMIT on a failed + # transaction is turned into a rollback by the coordinator anyway, + # so always rolling back here loses no coverage. + try: + exe.execute("ROLLBACK", http=Http.NO) + except QueryError: + exe.reconnect_next = True + raise + # Outside the error path both endings are worth exercising, and a + # failure of the end statement itself is a genuine one. + end = "COMMIT" if self.rng.choice([True, False]) else "ROLLBACK" + exe.execute(end, http=Http.NO) + return True + + +class DDLTransactionAction(Action): + """A DDL statement inside an explicit `BEGIN`/`COMMIT`. Materialize allows + only a single statement in a DDL transaction, so the value over autocommit + DDL is the open commit window: racing concurrent DDL against it exercises + the "another session modified the catalog while this DDL transaction was + open" serialization path.""" + + def applicable(self, exe: Executor) -> bool: + return exe.db.complexity in (Complexity.DDL, Complexity.DDLOnly) + + def errors_to_ignore(self, exe: Executor) -> list[str]: + return [ + "another session modified the catalog", + "unknown schema", + ] + super().errors_to_ignore(exe) + + def run(self, exe: Executor) -> bool: + if len([t for t in exe.db.tables if not t.temp]) >= MAX_TABLES: + return False + try: + schema = self.rng.choice(exe.db.schemas) + except IndexError: + return False + table_id = exe.db.table_id + exe.db.table_id += 1 + table = Table(self.rng, table_id, schema) + exe.execute("BEGIN", http=Http.NO) + try: + table.create(exe) + exe.execute("COMMIT", http=Http.NO) + except QueryError: + try: + exe.execute("ROLLBACK", http=Http.NO) + except QueryError: + pass + raise + with exe.db.lock: + exe.db.tables.append(table) + return True + + +class CommitRollbackAction(Action): + def run(self, exe: Executor) -> bool: + if not exe.action_run_since_last_commit_rollback: + return False + + if self.rng.random() < 0.7: + exe.commit() + else: + exe.rollback() + exe.action_run_since_last_commit_rollback = False + return True + + +class FlipFlagsAction(Action): + def __init__( + self, + rng: random.Random, + composition: Composition | None, + ): + super().__init__(rng, composition) + + BOOLEAN_FLAG_VALUES = ["TRUE", "FALSE"] + + self.flags_with_values: dict[str, list[str]] = dict() + self.flags_with_values["persist_blob_target_size"] = ( + # 1 MiB, 16 MiB, 128 MiB + ["1048576", "16777216", "134217728"] + ) + for flag in ["catalog", "source", "snapshot", "txn"]: + self.flags_with_values[f"persist_use_critical_since_{flag}"] = ( + BOOLEAN_FLAG_VALUES + ) + self.flags_with_values["persist_claim_unclaimed_compactions"] = ( + BOOLEAN_FLAG_VALUES + ) + self.flags_with_values["persist_optimize_ignored_data_fetch"] = ( + BOOLEAN_FLAG_VALUES ) self.flags_with_values["persist_source_fetch_concurrency"] = [ "1", @@ -1960,6 +2909,9 @@ def __init__( BOOLEAN_FLAG_VALUES ) self.flags_with_values["enable_alter_table_add_column"] = BOOLEAN_FLAG_VALUES + self.flags_with_values["enable_bounded_staleness_isolation"] = ( + BOOLEAN_FLAG_VALUES + ) self.flags_with_values["enable_arrangement_dictionary_compression_alpha"] = ( BOOLEAN_FLAG_VALUES ) @@ -1975,10 +2927,10 @@ def __init__( BOOLEAN_FLAG_VALUES ) self.flags_with_values["cluster"] = ["quickstart", "dont_exist"] - self.flags_with_values["enable_frontend_peek_sequencing"] = [ - "true", - "false", - ] + # NOTE: enable_frontend_peek_sequencing is pinned off in + # ADDITIONAL_SYSTEM_PARAMETER_DEFAULTS (frontend-peek read-hold vs + # compaction race, https://linear.app/materializeinc/issue/SQL-520), so + # it is not flipped here. self.flags_with_values["enable_frontend_subscribes"] = [ "true", "false", @@ -1995,11 +2947,17 @@ def __init__( self.flags_with_values["enable_column_paged_batcher_spill"] = ( BOOLEAN_FLAG_VALUES ) + # Fractions of the *cgroup* memory limit, which under mzcompose's + # process orchestrator is the whole container budget shared by + # environmentd and every replica, not one replica's allowance. A + # process-singleton pool at 0.25 is 6 GiB per clusterd on the 24 GiB CI + # agent, so the values stay small enough that all replicas together + # cannot claim the budget. Small pools also fill up sooner, which + # exercises the paging and compression paths more, not less. self.flags_with_values["column_paged_batcher_budget_fraction"] = [ "0.0", "0.01", - "0.05", - "0.25", + "0.02", ] self.flags_with_values["column_paged_batcher_lz4"] = BOOLEAN_FLAG_VALUES self.flags_with_values["column_paged_batcher_swap_pageout"] = ( @@ -2013,10 +2971,14 @@ def __init__( self.flags_with_values["column_paged_batcher_eager_backing"] = ( BOOLEAN_FLAG_VALUES ) + # Same shared-budget reasoning as the budget fraction above. Sharing its + # value set keeps all three orderings against the budget covered: a + # target of 0 collapses the compressed tier, a target below the budget + # pages out above it, and a target above the budget pages nothing out. self.flags_with_values["column_paged_batcher_pool_rss_target_fraction"] = [ "0.0", - "0.25", - "0.5", + "0.01", + "0.02", ] self.flags_with_values["enable_upsert_paged_spill"] = BOOLEAN_FLAG_VALUES # 0 forces the estimated-size path for every table, the default forces @@ -2329,12 +3291,21 @@ def run(self, exe: Executor) -> bool: flag_value = self.rng.choice(self.flags_with_values[flag_name]) + # Occasionally restore the default instead, a distinct path from + # SET-to-value. + reset = self.rng.random() < 0.1 + conn = None try: conn = self.create_system_connection(exe) if cluster is not None: self.set_cluster_compression(conn, cluster) + elif reset: + self.reset_flag(conn, flag_name) + # Gates reading exe.db.flags fall back to their conservative + # default when the key is absent. + exe.db.flags.pop(flag_name, None) else: self.flip_flag(conn, flag_name, flag_value) exe.db.flags[flag_name] = flag_value @@ -2366,6 +3337,12 @@ def set_cluster_compression(self, conn: Connection, cluster: Cluster) -> None: with conn.cursor() as cur: cur.execute(f"ALTER CLUSTER {cluster} {option};".encode()) + def reset_flag(self, conn: Connection, flag_name: str) -> None: + with conn.cursor() as cur: + cur.execute( + f"ALTER SYSTEM RESET {flag_name};".encode(), + ) + class CreateViewAction(Action): def errors_to_ignore(self, exe: Executor) -> list[str]: @@ -2469,7 +3446,46 @@ def run(self, exe: Executor) -> bool: else: query = f"DROP VIEW {view}" exe.execute(query, http=Http.RANDOM) - exe.db.views.remove(view) + # A concurrent CASCADE drop's untrack_objects_in_schemas may have + # already filtered this view out of the list, tolerate that. + try: + exe.db.views.remove(view) + except ValueError: + pass + return True + + +class CreateOrReplaceViewAction(Action): + """In-place swap of an existing view's definition via CREATE OR REPLACE. + + The body is unchanged, so dependents stay valid, but the coordinator still + tears down and rebuilds the item (and the dataflow, for a materialized + view). Racing that swap against reads and concurrent DDL is the point.""" + + def errors_to_ignore(self, exe: Executor) -> list[str]: + errors = [ + # A dependent references a column the replacement would drop. The + # body is unchanged here, but a concurrent replacement may have + # changed it. + "still depended upon by", + "replica-targeted materialized views is not supported", + "unknown cluster replica", + ] + super().errors_to_ignore(exe) + if exe.db.scenario == Scenario.Rename: + # A base object was renamed, invalidating the captured body. + errors += ["does not exist", "ambiguous reference to schema name"] + return errors + + def run(self, exe: Executor) -> bool: + with exe.db.lock: + views = [view for view in exe.db.views if not view.temp] + if not views: + return False + view = self.rng.choice(views) + with view.lock: + if view not in exe.db.views: + return False + view.create(exe, or_replace=True) return True @@ -2517,6 +3533,48 @@ def run(self, exe: Executor) -> bool: return True +class AlterRoleAction(Action): + """ALTER ROLE ... SET / RESET a default session variable. + + Exercises the per-role session-default path (applied at session init) and + role-name resolution on the ALTER path, neither of which CREATE/DROP ROLE + touch. Part of broadening ALTER coverage (the ALTER paths are where several + catalog bugs have hidden vs the well-worn CREATE/DROP).""" + + def run(self, exe: Executor) -> bool: + with exe.db.lock: + if not exe.db.roles: + return False + role = self.rng.choice(exe.db.roles) + with role.lock: + # Was dropped while we were acquiring the lock. + if role not in exe.db.roles: + return False + var, value = self.rng.choice( + [ + ("cluster", "'quickstart'"), + ("transaction_isolation", "'serializable'"), + ("statement_timeout", "'120s'"), + ("search_path", "public"), + ] + ) + query = ( + f"ALTER ROLE {role} RESET {var}" + if self.rng.choice([True, False]) + else f"ALTER ROLE {role} SET {var} = {value}" + ) + try: + exe.execute(query, http=Http.RANDOM) + except QueryError as e: + # Concurrent DROP ROLE, expected as with DropRoleAction. + if ( + exe.db.scenario not in (Scenario.Kill, Scenario.ZeroDowntimeDeploy) + or "unknown role" not in e.msg + ): + raise e + return True + + class CreateClusterAction(Action): def run(self, exe: Executor) -> bool: with exe.db.lock: @@ -2527,9 +3585,7 @@ def run(self, exe: Executor) -> bool: cluster = Cluster( cluster_id, managed=self.rng.choice([True, False]), - size=self.rng.choice( - ["scale=1,workers=1", "scale=1,workers=4", "scale=2,workers=2"] - ), + size=self.rng.choice(["scale=1,workers=1", "scale=1,workers=2"]), replication_factor=self.rng.choice([1, 2]), introspection_interval="1s", ) @@ -2681,9 +3737,7 @@ def run(self, exe: Executor) -> bool: replica = ClusterReplica( replica_id, - size=self.rng.choice( - ["scale=1,workers=1", "scale=1,workers=4", "scale=2,workers=2"] - ), + size=self.rng.choice(["scale=1,workers=1", "scale=1,workers=2"]), cluster=cluster, ) replica.create(exe) @@ -2729,8 +3783,10 @@ def run(self, exe: Executor) -> bool: class ReconfigureClusterAction(Action): - """Gracefully resize a random managed cluster and assert convergence. + """Gracefully resize a random managed cluster and watch its record settle. + The asserted invariant is that a reconfiguration record never sits + in-progress past its own deadline, whether it cuts over or rolls back. Regression coverage for SQL-530: readiness must not wait for single-replica sources (Postgres, MySQL, SQL Server) hosted on the cluster, which never hydrate on the reconfiguration's target replicas @@ -2781,40 +3837,85 @@ def run(self, exe: Executor) -> bool: # e.g. when another worker put a REFRESH EVERY materialized view # with a far-future refresh on the cluster, which cannot hydrate # on the target replicas before its refresh time. What must never - # happen is the record staying "in-progress" past its deadline, - # which is how a wedged readiness check (SQL-530) manifests. Only - # trust a terminal status once the ALTER's own deadline has - # passed: builtin table reads can lag the catalog, so before that - # a terminal status can only be a stale record from an earlier - # reconfiguration of this cluster. + # happen is a record staying "in-progress" past its own deadline, + # which is how a wedged readiness check (SQL-530) manifests. + # + # `mz_now()` is the timestamp the read itself ran at, and comparing + # it against the record's `deadline` is what makes that assertion + # sound. `mz_clusters` is a builtin table and + # `mz_cluster_reconfigurations` a builtin materialized view + # maintained by mz_catalog_server, and a SERIALIZABLE read of either + # takes the freshest timestamp that does not block, so under load it + # can lag wall clock by minutes. Judging the deadline against the + # test's wall clock therefore calls a correctly rolled-back record + # wedged as soon as the read lags. Both sides coming from the same + # read keeps them in the same frame, stale or not. name_literal = cluster.name().replace("'", "''") query = ( - "SELECT c.size, r.status FROM mz_clusters c " + "SELECT c.size, r.status, r.deadline::text, mz_now()::text " + "FROM mz_clusters c " "LEFT JOIN mz_internal.mz_cluster_reconfigurations r ON r.cluster_id = c.id " f"WHERE c.name = '{name_literal}'" ) - status = None + # The deadline our own ALTER wrote is `now() + 120s` in the same + # millisecond epoch `mz_now()` reports, so a read past this is + # guaranteed to show our record or a later one, never a stale + # terminal record from an earlier reconfiguration of this cluster. + our_deadline_ms = int((alter_started + 120) * 1000) seen = False - deadline = time.time() + 180 - while time.time() < deadline: + # Must stay comfortably under the 300s the run gives all workers to + # join after `end_time` (`parallel_workload.run`). A poll that + # outlives that budget makes a worker doing its job look wedged, and + # the run then hard-exits 0 without its final checks. + poll_until = time.time() + 240 + while time.time() < poll_until: exe.execute(query) rows = exe.cur.fetchall() if not rows: # `mz_clusters` is a builtin table, so the read can lag a # concurrent rename or swap of this cluster and not show - # the polled name yet. Keep polling and let the deadline + # the polled name yet. Keep polling and let the check # below decide, rather than indexing an empty result. time.sleep(1) continue seen = True - size, status = rows[0] + size, status, record_deadline, read_ts = rows[0] + read_ts_ms = int(read_ts) if size == new_size: cluster.size = new_size return True - if time.time() > alter_started + 120 and status in ( - "timed-out", - "resource-exhausted", + # 60s of slack over the controller's 5s tick, which is what + # turns a reached deadline into the terminal status. + if ( + status == "in-progress" + and read_ts_ms > int(record_deadline) + 60_000 ): + raise ValueError( + f"Graceful reconfiguration of cluster {cluster} to size " + f"{new_size} sat in-progress {read_ts_ms - int(record_deadline)}ms " + f"past its deadline {record_deadline}" + ) + if read_ts_ms > our_deadline_ms and status is None: + raise ValueError( + f"Reconfiguration record of cluster {cluster} is gone as of " + f"{read_ts_ms}, past the deadline {our_deadline_ms} of the " + f"reconfiguration to size {new_size} that wrote it" + ) + # Any settled record is a legitimate end state, so stop on every + # status but "in-progress". A rollback at the deadline + # ("timed-out"), a controller that could not get the replicas + # ("resource-exhausted"), a concurrent shape-touching ALTER + # calling the transition off ("cancelled"), or one that + # retargeted it to a shape that then cut over ("finalized" on a + # size that is not ours) all mean the reconfiguration settled. + # FlipFlagsAction alters EXPERIMENTAL ARRANGEMENT COMPRESSION + # without holding the cluster lock, and that is a shape + # dimension (`alter_changes_replica_shape`), so it reaches the + # record this poll is watching. Once ours settles, such an ALTER + # even starts a fresh record, in-progress with the 24h default + # deadline, which is why running out of the budget below on an + # in-progress record is not by itself a failure. + if read_ts_ms > our_deadline_ms and status != "in-progress": return True time.sleep(1) if not seen: @@ -2823,10 +3924,7 @@ def run(self, exe: Executor) -> bool: f"name the reconfiguration to size {new_size} polled for, so the " f"name the workload holds does not match the catalog" ) - raise ValueError( - f"Graceful reconfiguration of cluster {cluster} to size {new_size} " - f"did not complete (reconfiguration status: {status})" - ) + return True class GrantPrivilegesAction(Action): @@ -2835,7 +3933,7 @@ def run(self, exe: Executor) -> bool: if not exe.db.roles: return False role = self.rng.choice(exe.db.roles) - privilege = self.rng.choice(["SELECT", "INSERT", "UPDATE", "ALL"]) + privilege = self.rng.choice(["SELECT", "INSERT", "UPDATE", "DELETE", "ALL"]) tables_views: list[DBObject] = [*exe.db.tables, *exe.db.views] table = self.rng.choice(tables_views) with table.lock, role.lock: @@ -2863,7 +3961,7 @@ def run(self, exe: Executor) -> bool: if not exe.db.roles: return False role = self.rng.choice(exe.db.roles) - privilege = self.rng.choice(["SELECT", "INSERT", "UPDATE", "ALL"]) + privilege = self.rng.choice(["SELECT", "INSERT", "UPDATE", "DELETE", "ALL"]) tables_views: list[DBObject] = [*exe.db.tables, *exe.db.views] table = self.rng.choice(tables_views) with table.lock, role.lock: @@ -2885,68 +3983,597 @@ def run(self, exe: Executor) -> bool: return True -# TODO: Should factor this out so can easily use it without action -class ReconnectAction(Action): - def __init__( - self, - rng: random.Random, - composition: Composition | None, - random_role: bool = True, - ): - super().__init__(rng, composition) - self.random_role = random_role +class GrantRoleAction(Action): + def errors_to_ignore(self, exe: Executor) -> list[str]: + return [ + "unknown role", + # Concurrent memberships can close a cycle, which is rejected. + "is a member of role", + ] + super().errors_to_ignore(exe) def run(self, exe: Executor) -> bool: - exe.mz_service = "materialized" - exe.log("reconnecting") - # The connection's temp objects die with it, drop them from the - # tracked state so other workers stop querying them. - if exe.temp_objects: - with exe.db.lock: - exe.db.tables[:] = [ - t for t in exe.db.tables if t not in exe.temp_objects - ] - exe.db.views[:] = [v for v in exe.db.views if v not in exe.temp_objects] - exe.temp_objects.clear() - host = exe.db.host + with exe.db.lock: + if len(exe.db.roles) < 2: + return False + role1, role2 = self.rng.sample(exe.db.roles, 2) + # Both locks are held across the round trip below, so take them in + # a fixed order. Two workers that sampled the same pair in + # opposite orders would otherwise deadlock. The GRANT operands + # keep the sampled order, so memberships still form in both + # directions, including the cycles this action provokes. + first, second = sorted((role1, role2), key=lambda role: role.role_id) + with first.lock, second.lock: + if role1 not in exe.db.roles or role2 not in exe.db.roles: + return False + exe.execute(f"GRANT {role1} TO {role2}", http=Http.RANDOM) + return True - def pg_port() -> int: - # System workers (e.g. the Cancel worker) live on the internal - # port, everyone else on the external one of the current service. - if exe.user == "mz_system": - return exe.db.ports[ - "mz_system" if exe.mz_service == "materialized" else "mz_system2" - ] - return exe.db.ports[exe.mz_service] +class RevokeRoleAction(Action): + def errors_to_ignore(self, exe: Executor) -> list[str]: + return [ + "unknown role", + ] + super().errors_to_ignore(exe) + + def run(self, exe: Executor) -> bool: with exe.db.lock: - if self.random_role and exe.db.roles: - user = self.rng.choice( - ["materialize", str(self.rng.choice(exe.db.roles))] + if len(exe.db.roles) < 2: + return False + role1, role2 = self.rng.sample(exe.db.roles, 2) + # Fixed lock order, as in GrantRoleAction: the sampled order + # deadlocks two workers that picked the same pair, and the REVOKE + # operands have to stay in the sampled order. + first, second = sorted((role1, role2), key=lambda role: role.role_id) + with first.lock, second.lock: + if role1 not in exe.db.roles or role2 not in exe.db.roles: + return False + exe.execute(f"REVOKE {role1} FROM {role2}", http=Http.RANDOM) + return True + + +class AlterOwnerAction(Action): + def errors_to_ignore(self, exe: Executor) -> list[str]: + result = [ + "unknown role", + "must be a member of", + ] + super().errors_to_ignore(exe) + if exe.db.complexity in (Complexity.DDL, Complexity.DDLOnly): + result.extend(["does not exist"]) + return result + + def run(self, exe: Executor) -> bool: + # Only the snapshot needs the global lock, see CommentAction. + with exe.db.lock: + if not exe.db.roles: + return False + role = self.rng.choice(exe.db.roles) + tables = list(exe.db.tables) + views = list(exe.db.views) + direct_sources = exe.db.kafka_sources + exe.db.webhook_sources + table_sources = ( + exe.db.postgres_sources + + exe.db.mysql_sources + + exe.db.sql_server_sources + ) + loadgen_sources = list(exe.db.loadgen_sources) + sinks = exe.db.kafka_sinks + exe.db.iceberg_sinks + schemas = list(exe.db.schemas) + dbs = list(exe.db.dbs) + clusters = list(exe.db.clusters) + + candidates: list[tuple[str, str]] = [] + # Temp objects cannot change owner, they die with the session. + for table in tables: + if not table.temp: + candidates.append(("TABLE", str(table))) + for view in views: + if not view.temp: + candidates.append( + ( + "MATERIALIZED VIEW" if view.materialized else "VIEW", + str(view), + ) ) - else: - # Keep the executor's original user, e.g. the Cancel worker - # must stay mz_system or its cancels fail with "must be a - # member of" - user = exe.user - conn = exe.cur.connection + # Kafka and webhook source objects are readable directly, the + # others follow the source-table model: str(obj) names the table, + # the ingestion source is a separate catalog item. + for source in direct_sources: + candidates.append(("SOURCE", str(source))) + for source in table_sources: + candidates.append(("TABLE", str(source))) + candidates.append( + ( + "SOURCE", + f"{source.schema}.{identifier(source.executor.source)}", + ) + ) + for source in loadgen_sources: + candidates.append(("TABLE", str(source))) + candidates.append( + ("SOURCE", f"{source.schema}.{identifier(source.source_name())}") + ) + for sink in sinks: + candidates.append(("SINK", str(sink))) + for schema in schemas: + candidates.append(("SCHEMA", str(schema))) + for db in dbs: + candidates.append(("DATABASE", str(db))) + for cluster in clusters: + candidates.append(("CLUSTER", str(cluster))) + candidates.append(("SECRET", "materialize.public.pgpass")) + # NOTE: No CONNECTION target. Changing a connection's owner emits a + # Connection(Altered) implication, which re-alters every dependent + # sink's export connection. That re-alter can fail with InvalidAlter + # and panic the coordinator. + # See https://linear.app/materializeinc/issue/SQL-517 + kind, name = self.rng.choice(candidates) + with role.lock: + if role not in exe.db.roles: + return False + exe.execute(f"ALTER {kind} {name} OWNER TO {role}", http=Http.RANDOM) + return True - if exe.ws and exe.use_ws: - try: - exe.ws.close() - except: - pass - try: - exe.cur.close() - except: - pass - try: - conn.close() - except: - pass +class AlterDefaultPrivilegesAction(Action): + # Privileges per object type. The bool pair is (allows IN SCHEMA, allows IN + # DATABASE): schema-scoped objects accept both, SCHEMAS only IN DATABASE, + # DATABASES and CLUSTERS neither. Mixing e.g. ON DATABASES with IN DATABASE + # is a plan error, not a race, so it is generated out. + OBJECT_TYPES = { + "TABLES": (["SELECT", "INSERT", "UPDATE", "DELETE", "ALL"], True, True), + "TYPES": (["USAGE", "ALL"], True, True), + "SECRETS": (["USAGE", "ALL"], True, True), + "CONNECTIONS": (["USAGE", "ALL"], True, True), + "SCHEMAS": (["USAGE", "CREATE", "ALL"], False, True), + "DATABASES": (["USAGE", "CREATE", "ALL"], False, False), + "CLUSTERS": (["USAGE", "CREATE", "ALL"], False, False), + } - NUM_ATTEMPTS = 20 + def errors_to_ignore(self, exe: Executor) -> list[str]: + result = [ + "unknown role", + "unknown schema", + "unknown database", + "must be a member of", + # FOR ALL ROLES and system-adjacent grants require privileges the + # (possibly reconnected-as-a-random-role) session may lack. + "permission denied to", + ] + super().errors_to_ignore(exe) + return result + + def run(self, exe: Executor) -> bool: + object_type = self.rng.choice(list(self.OBJECT_TYPES.keys())) + privileges, allows_in_schema, allows_in_database = self.OBJECT_TYPES[ + object_type + ] + privilege = self.rng.choice(privileges) + with exe.db.lock: + if not exe.db.roles: + return False + role = self.rng.choice(exe.db.roles) + for_clause = self.rng.choice( + ["FOR ALL ROLES"] + [f"FOR ROLE {r}" for r in exe.db.roles] + ) + in_clause = "" + if allows_in_schema and exe.db.schemas and self.rng.random() < 0.3: + in_clause = f" IN SCHEMA {self.rng.choice(exe.db.schemas)}" + elif allows_in_database and exe.db.dbs and self.rng.random() < 0.3: + in_clause = f" IN DATABASE {self.rng.choice(exe.db.dbs)}" + with role.lock: + if role not in exe.db.roles: + return False + if self.rng.choice([True, False]): + query = f"ALTER DEFAULT PRIVILEGES {for_clause}{in_clause} GRANT {privilege} ON {object_type} TO {role}" + else: + query = f"ALTER DEFAULT PRIVILEGES {for_clause}{in_clause} REVOKE {privilege} ON {object_type} FROM {role}" + exe.execute(query, http=Http.RANDOM) + return True + + +class BroadPrivilegesAction(Action): + """GRANT/REVOKE on object classes beyond the tables and views covered by + Grant/RevokePrivilegesAction.""" + + def errors_to_ignore(self, exe: Executor) -> list[str]: + return [ + "unknown role", + "unknown schema", + "unknown database", + "unknown cluster", + # System privileges require superuser, which a session reconnected + # as a random role does not have. + "permission denied to", + ] + super().errors_to_ignore(exe) + + def run(self, exe: Executor) -> bool: + with exe.db.lock: + if not exe.db.roles: + return False + role = self.rng.choice(exe.db.roles) + targets: list[tuple[str, list[str]]] = [ + ("SYSTEM", ["CREATEDB", "CREATECLUSTER", "CREATEROLE", "ALL"]), + ("SECRET materialize.public.pgpass", ["USAGE", "ALL"]), + # NOTE: No CONNECTION target. GRANT/REVOKE on a connection emits + # a Connection(Altered) implication, which re-alters every + # dependent sink's export connection. That re-alter can fail + # with InvalidAlter and panic the coordinator. + # See https://linear.app/materializeinc/issue/SQL-517 + ] + if exe.db.schemas: + targets.append( + ( + f"SCHEMA {self.rng.choice(exe.db.schemas)}", + ["USAGE", "CREATE", "ALL"], + ) + ) + if exe.db.dbs: + targets.append( + ( + f"DATABASE {self.rng.choice(exe.db.dbs)}", + ["USAGE", "CREATE", "ALL"], + ) + ) + if exe.db.clusters: + targets.append( + ( + f"CLUSTER {self.rng.choice(exe.db.clusters)}", + ["USAGE", "CREATE", "ALL"], + ) + ) + target, privileges = self.rng.choice(targets) + privilege = self.rng.choice(privileges) + with role.lock: + if role not in exe.db.roles: + return False + if self.rng.choice([True, False]): + exe.execute( + f"GRANT {privilege} ON {target} TO {role}", http=Http.RANDOM + ) + else: + exe.execute( + f"REVOKE {privilege} ON {target} FROM {role}", http=Http.RANDOM + ) + return True + + +class ShowAction(Action): + def errors_to_ignore(self, exe: Executor) -> list[str]: + result = super().errors_to_ignore(exe) + result.extend( + [ + # SHOW CREATE CLUSTER only works for managed clusters. We only + # target managed ones, this covers a managed->unmanaged race. + "SHOW CREATE for unmanaged clusters not yet supported", + # With auto_route_catalog_queries off, SHOW compiles to a + # catalog query on the active cluster, so it shares the read + # transaction's timedomain. + "in the same timedomain", + 'is not allowed from the "mz_catalog_server" cluster', + ] + ) + if exe.db.complexity in (Complexity.DDL, Complexity.DDLOnly): + result.extend(["does not exist"]) + return result + + def run(self, exe: Executor) -> bool: + if self.rng.choice([True, False]): + schema_scoped = [ + "SHOW TABLES", + "SHOW VIEWS", + "SHOW MATERIALIZED VIEWS", + "SHOW SOURCES", + "SHOW SINKS", + "SHOW INDEXES", + "SHOW OBJECTS", + "SHOW SECRETS", + "SHOW CONNECTIONS", + "SHOW TYPES", + ] + other = [ + "SHOW CLUSTERS", + "SHOW CLUSTER REPLICAS", + "SHOW DATABASES", + "SHOW SCHEMAS", + "SHOW ROLES", + "SHOW PRIVILEGES", + "SHOW DEFAULT PRIVILEGES", + "SHOW ROLE MEMBERSHIP", + "SHOW ALL", + ] + if self.rng.choice([True, False]): + query = self.rng.choice(schema_scoped) + with exe.db.lock: + if exe.db.schemas and self.rng.choice([True, False]): + query += f" FROM {self.rng.choice(exe.db.schemas)}" + if self.rng.random() < 0.2: + query += " LIKE '%1%'" + else: + query = self.rng.choice(other) + else: + # Only the snapshot needs the global lock, see CommentAction. + with exe.db.lock: + tables = list(exe.db.tables) + views = list(exe.db.views) + direct_sources = exe.db.kafka_sources + exe.db.webhook_sources + table_sources = ( + exe.db.postgres_sources + + exe.db.mysql_sources + + exe.db.sql_server_sources + ) + loadgen_sources = list(exe.db.loadgen_sources) + sinks = exe.db.kafka_sinks + exe.db.iceberg_sinks + indexes = list(exe.db.indexes) + clusters = list(exe.db.clusters) + + candidates: list[tuple[str, str]] = [ + ("CONNECTION", "materialize.public.kafka_conn"), + ("CONNECTION", "materialize.public.csr_conn"), + ] + for table in tables: + candidates.append(("TABLE", str(table))) + for view in views: + candidates.append( + ( + "MATERIALIZED VIEW" if view.materialized else "VIEW", + str(view), + ) + ) + # Kafka and webhook sources are readable directly. The others + # follow the source-table model where str() is the table and + # the ingestion source is a separate catalog item. + for source in direct_sources: + candidates.append(("SOURCE", str(source))) + for source in table_sources: + candidates.append(("TABLE", str(source))) + candidates.append( + ( + "SOURCE", + f"{source.schema}.{identifier(source.executor.source)}", + ) + ) + for source in loadgen_sources: + candidates.append(("TABLE", str(source))) + candidates.append( + ( + "SOURCE", + f"{source.schema}.{identifier(source.source_name())}", + ) + ) + for sink in sinks: + candidates.append(("SINK", str(sink))) + for index in indexes: + candidates.append(("INDEX", str(index))) + for cluster in clusters: + # SHOW CREATE CLUSTER is not supported for unmanaged + # clusters. + if cluster.managed: + candidates.append(("CLUSTER", str(cluster))) + kind, name = self.rng.choice(candidates) + # SHOW REDACTED CREATE CLUSTER is not supported + redacted = ( + "REDACTED " + if kind != "CLUSTER" and self.rng.choice([True, False]) + else "" + ) + query = f"SHOW {redacted}CREATE {kind} {name}" + exe.execute(query, http=Http.RANDOM, fetch=True) + return True + + +class SetSessionVariableAction(Action): + def __init__(self, rng: random.Random, composition: Composition | None): + super().__init__(rng, composition) + self.vars_with_values: dict[str, list[str]] = { + "statement_timeout": ["'30s'", "'60s'", "'0s'"], + "application_name": ["'parallel-workload'", "''"], + "client_min_messages": ["debug1", "info", "notice", "warning", "error"], + "max_query_result_size": ["100000", "1000000", "1000000000"], + "emit_timestamp_notice": ["true", "false"], + "emit_trace_id_notice": ["true", "false"], + # Only UTC is accepted, the rejection of other time zones is + # deliberate error path coverage. + "timezone": ["'UTC'", "'America/New_York'"], + } + + def errors_to_ignore(self, exe: Executor) -> list[str]: + return [ + "invalid value for parameter", + "cannot have value", + "unrecognized configuration parameter", + ] + super().errors_to_ignore(exe) + + def run(self, exe: Executor) -> bool: + var = self.rng.choice(list(self.vars_with_values.keys())) + # statement_timeout is tracked on the executor, and only a statement on + # the worker's own session can change what that session times out at. + # Over the HTTP endpoint the SET or RESET lands on a one-shot session + # instead, leaving the tracking describing a value nothing has. + http = Http.NO if var == "statement_timeout" else Http.RANDOM + if self.rng.random() < 0.2: + exe.execute(f"RESET {var}", http=http) + if var == "statement_timeout": + # Back at the default timeout, so a timeout is a genuine + # finding again. + exe.statement_timeout_set = False + return True + value = self.rng.choice(self.vars_with_values[var]) + local = "LOCAL " if self.rng.random() < 0.1 else "" + exe.execute(f"SET {local}{var} = {value}", http=http) + if var == "statement_timeout": + # A timeout is expected for as long as this one is configured. A + # SET LOCAL keeps this set past the transaction that discards the + # value, which only errs towards tolerating more. + exe.statement_timeout_set = True + return True + + +class DiscardAction(Action): + def run(self, exe: Executor) -> bool: + # The session's temp objects die with DISCARD, drop them from the + # tracked state so other workers stop querying them, mirroring + # ReconnectAction. + if exe.temp_objects: + with exe.db.lock: + exe.db.tables[:] = [ + t for t in exe.db.tables if t not in exe.temp_objects + ] + exe.db.views[:] = [v for v in exe.db.views if v not in exe.temp_objects] + exe.temp_objects.clear() + # Only DISCARD TEMP, not DISCARD ALL. DISCARD ALL (like DEALLOCATE ALL) + # deallocates every prepared statement, including the ones psycopg + # transparently auto-prepares. psycopg's client-side cache would then + # be stale and the next reuse of such a statement fails with + # "prepared statement ... does not exist". DISCARD TEMP still exercises + # the temp-object teardown, which is the interesting path here. + exe.execute("DISCARD TEMP", http=Http.NO) + return True + + +class ValidateConnectionAction(Action): + # aws_conn is excluded, MinIO's STS support for validation is unclear. + CONNECTIONS = [ + "kafka_conn", + "csr_conn", + "postgres_conn", + "mysql_conn", + "sql_server_conn", + "polaris_conn", + ] + + def errors_to_ignore(self, exe: Executor) -> list[str]: + return [ + "timeout: error trying to connect", + # A concurrent ALTER SECRET rotation (the workload only rotates a + # secret to its own value) can transiently expose an empty secret, + # so VALIDATE CONNECTION sends an empty password and the upstream + # system rejects it (secret-rotation atomicity). The same race + # surfaces with each upstream's own wording. + # TODO: Remove when https://linear.app/materializeinc/issue/SS-347 + # is fixed. + "empty password returned by client", # Postgres + "Access denied for user", # MySQL + "Login failed for user", # SQL Server + ] + super().errors_to_ignore(exe) + + def run(self, exe: Executor) -> bool: + name = self.rng.choice(self.CONNECTIONS) + exe.execute(f"VALIDATE CONNECTION materialize.public.{name}", http=Http.NO) + return True + + +class AlterConnectionAction(Action): + # The SET clause per connection, setting the option to the value the + # connection already has. That still exercises the full reconfiguration + # path (restarting dependent sources and sinks) without breaking them. + # NOTE: BROKER takes no `=` (it is parsed specially), HOST/URL do. + SET_CLAUSES = { + "kafka_conn": "BROKER 'kafka:9092'", + "csr_conn": "URL = 'http://schema-registry:8081'", + "postgres_conn": "HOST = 'postgres'", + "mysql_conn": "HOST = 'mysql'", + "sql_server_conn": "HOST = 'sql-server'", + } + + def errors_to_ignore(self, exe: Executor) -> list[str]: + return [ + "timeout: error trying to connect", + # The storage controller can refuse an in-place connection change + # depending on the connection's current state or dependents. + "cannot be altered in the requested way", + ] + super().errors_to_ignore(exe) + + def run(self, exe: Executor) -> bool: + name = self.rng.choice(list(self.SET_CLAUSES.keys())) + set_clause = self.SET_CLAUSES[name] + query = f"ALTER CONNECTION materialize.public.{name} SET ({set_clause})" + if self.rng.choice([True, False]): + validate = self.rng.choice(["true", "false"]) + query += f" WITH (VALIDATE = {validate})" + exe.execute(query, http=Http.RANDOM) + return True + + +class AlterSecretAction(Action): + def run(self, exe: Executor) -> bool: + # Rotate to the same value, exercising the rotation path (including + # dependent connections picking up the new secret version) without + # breaking the credentials. + name, value = self.rng.choice( + [ + ("pgpass", "postgres"), + ("mypass", MySql.DEFAULT_ROOT_PASSWORD), + ("sql_server_pass", SqlServer.DEFAULT_SA_PASSWORD), + ("minio", "minioadmin"), + ] + ) + exe.execute( + f"ALTER SECRET materialize.public.{name} AS '{value}'", http=Http.RANDOM + ) + return True + + +# TODO: Should factor this out so can easily use it without action +class ReconnectAction(Action): + def __init__( + self, + rng: random.Random, + composition: Composition | None, + random_role: bool = True, + ): + super().__init__(rng, composition) + self.random_role = random_role + + def run(self, exe: Executor) -> bool: + exe.mz_service = "materialized" + exe.log("reconnecting") + # The connection's temp objects die with it, drop them from the + # tracked state so other workers stop querying them. + if exe.temp_objects: + with exe.db.lock: + exe.db.tables[:] = [ + t for t in exe.db.tables if t not in exe.temp_objects + ] + exe.db.views[:] = [v for v in exe.db.views if v not in exe.temp_objects] + exe.temp_objects.clear() + host = exe.db.host + + def pg_port() -> int: + # System workers (e.g. the Cancel worker) live on the internal + # port, everyone else on the external one of the current service. + if exe.user == "mz_system": + return exe.db.ports[ + "mz_system" if exe.mz_service == "materialized" else "mz_system2" + ] + return exe.db.ports[exe.mz_service] + + with exe.db.lock: + if self.random_role and exe.db.roles: + user = self.rng.choice( + ["materialize", str(self.rng.choice(exe.db.roles))] + ) + else: + # Keep the executor's original user, e.g. the Cancel worker + # must stay mz_system or its cancels fail with "must be a + # member of" + user = exe.user + conn = exe.cur.connection + + if exe.ws and exe.use_ws: + try: + exe.ws.close() + except: + pass + + try: + exe.cur.close() + except: + pass + try: + conn.close() + except: + pass + + NUM_ATTEMPTS = 20 if exe.ws: for i in range( NUM_ATTEMPTS @@ -2999,6 +4626,10 @@ def pg_port() -> int: exe.set_isolation("SERIALIZABLE") cur.execute("SET auto_route_catalog_queries TO false") conn.autocommit = exe.autocommit + # A statement_timeout set on the old session died with it, so + # timeouts are genuine findings again unless a role default + # carries one over. + exe.statement_timeout_set = False try: cur.execute("SELECT pg_backend_pid()") except Exception as e: @@ -3073,7 +4704,11 @@ def __init__( def run(self, exe: Executor) -> bool: pid = self.rng.choice( - [worker.exe.pg_pid for worker in self.workers if worker.exe and worker.exe.pg_pid != -1] # type: ignore + [ + worker.exe.pg_pid + for worker in self.workers + if worker.exe and worker.exe.pg_pid != -1 + ] # type: ignore ) worker = None for i in range(len(self.workers)): @@ -3303,6 +4938,125 @@ def run(self, exe: Executor) -> bool: return True +class CreateLoadGeneratorSourceAction(Action): + def run(self, exe: Executor) -> bool: + with exe.db.lock: + if len(exe.db.loadgen_sources) >= MAX_LOADGEN_SOURCES: + return False + source_id = exe.db.loadgen_source_id + exe.db.loadgen_source_id += 1 + try: + cluster = self.rng.choice(exe.db.clusters) + schema = self.rng.choice(exe.db.schemas) + except IndexError: + # We mostly prevent index errors, but we don't want to lock too + # much since that would reduce our chance of finding race + # conditions in production code, so ignore the rare case where + # we accidentally removed all objects. + return False + with schema.lock, cluster.lock: + if schema not in exe.db.schemas: + return False + if cluster not in exe.db.clusters: + return False + + source = LoadGeneratorSource(source_id, cluster, schema, self.rng) + source.create(exe) + exe.db.loadgen_sources.append(source) + return True + + +class DropLoadGeneratorSourceAction(Action): + def errors_to_ignore(self, exe: Executor) -> list[str]: + return [ + "still depended upon by", + ] + super().errors_to_ignore(exe) + + def run(self, exe: Executor) -> bool: + with exe.db.lock: + if not exe.db.loadgen_sources: + return False + try: + source = self.rng.choice(exe.db.loadgen_sources) + except IndexError: + # We mostly prevent index errors, but we don't want to lock too + # much since that would reduce our chance of finding race + # conditions in production code, so ignore the rare case where + # we accidentally removed all objects. + return False + with source.lock: + # Was dropped while we were acquiring lock + if source not in exe.db.loadgen_sources: + return False + + exe.execute(f"DROP TABLE IF EXISTS {source}") + exe.execute( + f"DROP SOURCE {source.schema}.{identifier(source.source_name())} CASCADE", + http=Http.RANDOM, + ) + exe.db.loadgen_sources.remove(source) + return True + + +class CreateMultiLoadGeneratorSourceAction(Action): + GENERATORS = ["AUCTION", "TPCH", "MARKETING"] + + def errors_to_ignore(self, exe: Executor) -> list[str]: + # A rare cross-type subsource-name collision, or a concurrent create of + # the same generator type. + return ["already exists"] + super().errors_to_ignore(exe) + + def run(self, exe: Executor) -> bool: + with exe.db.lock: + present = {s.generator for s in exe.db.multi_loadgen_sources} + available = [g for g in self.GENERATORS if g not in present] + if not available: + return False + generator = self.rng.choice(available) + source_id = exe.db.multi_loadgen_source_id + exe.db.multi_loadgen_source_id += 1 + try: + cluster = self.rng.choice(exe.db.clusters) + schema = self.rng.choice(exe.db.schemas) + except IndexError: + return False + with schema.lock, cluster.lock: + if schema not in exe.db.schemas: + return False + if cluster not in exe.db.clusters: + return False + source = MultiLoadGeneratorSource( + source_id, cluster, schema, generator, self.rng + ) + source.create(exe) + # NOTE: No db.lock around the append. Taking it here would nest it + # inside the schema and cluster locks, inverting the db.lock-first + # order every other action uses. list.append is atomic, and a + # concurrent same-generator create is caught server-side ("already + # exists", tolerated above). + exe.db.multi_loadgen_sources.append(source) + return True + + +class DropMultiLoadGeneratorSourceAction(Action): + def errors_to_ignore(self, exe: Executor) -> list[str]: + return [ + "still depended upon by", + ] + super().errors_to_ignore(exe) + + def run(self, exe: Executor) -> bool: + with exe.db.lock: + if not exe.db.multi_loadgen_sources: + return False + source = self.rng.choice(exe.db.multi_loadgen_sources) + with source.lock: + if source not in exe.db.multi_loadgen_sources: + return False + exe.execute(f"DROP SOURCE {source} CASCADE", http=Http.RANDOM) + exe.db.multi_loadgen_sources.remove(source) + return True + + class CreateKafkaSourceAction(Action): def run(self, exe: Executor) -> bool: with exe.db.lock: @@ -3713,7 +5467,7 @@ def run(self, exe: Executor) -> bool: sink_id, cluster, schema, - self.rng.choice(exe.db.db_objects_without_views()), + self.rng.choice(exe.db.db_objects_for_sinks()), self.rng, ) sink.create(exe) @@ -3781,7 +5535,7 @@ def run(self, exe: Executor) -> bool: sink_id, cluster, schema, - self.rng.choice(exe.db.db_objects_without_views()), + self.rng.choice(exe.db.db_objects_for_sinks()), self.rng, ) sink.create(exe) @@ -3823,6 +5577,11 @@ def errors_to_ignore(self, exe: Executor) -> list[str]: result = super().errors_to_ignore(exe) if exe.db.scenario == Scenario.Rename: result.extend(["404: no object was found at the path"]) + # DropSchemaCascadeAction / DropDatabaseCascadeAction can drop a + # webhook source concurrently without taking its per-object lock, so a + # POST that already picked the source can 404. + if exe.db.complexity in (Complexity.DDL, Complexity.DDLOnly): + result.extend(["404: no object was found at the path"]) return result def run(self, exe: Executor) -> bool: @@ -3887,6 +5646,383 @@ def run(self, exe: Executor) -> bool: return True +class AlterClusterSetAction(Action): + """Live reconfigure of a managed cluster (SIZE / REPLICATION FACTOR). + + Resizing or changing the replica count of a cluster hosting indexes, MVs, + sources, and sinks forces rehydration and replica teardown/spin-up under + concurrent DDL and DML.""" + + def errors_to_ignore(self, exe: Executor) -> list[str]: + return [ + # A SET (SIZE) here or a ReconfigureCluster on the same cluster + # leaves a reconfiguration record in flight past the statement + # that started it. Replication factor is folded in at cut-over, + # so changing it meanwhile is refused. + "cannot change replication factor while a reconfiguration is in progress", + ] + super().errors_to_ignore(exe) + + def run(self, exe: Executor) -> bool: + with exe.db.lock: + # Cluster 0 stays fixed, it hosts sources and sinks. + managed = [c for c in exe.db.clusters[1:] if c.managed] + if not managed: + return False + cluster = self.rng.choice(managed) + with cluster.lock: + if cluster not in exe.db.clusters or not cluster.managed: + return False + choice = self.rng.choice(["size", "replication_factor", "reset_rf"]) + if choice == "size": + new_size = self.rng.choice(["scale=1,workers=1", "scale=1,workers=2"]) + exe.execute( + f"ALTER CLUSTER {cluster} SET (SIZE = '{new_size}')", + http=Http.RANDOM, + ) + cluster.size = new_size + for replica in cluster.replicas: + replica.size = new_size + elif choice == "replication_factor": + rf = self.rng.choice([1, 2]) + exe.execute( + f"ALTER CLUSTER {cluster} SET (REPLICATION FACTOR = {rf})", + http=Http.RANDOM, + ) + self._resize_replicas(cluster, rf) + else: + exe.execute( + f"ALTER CLUSTER {cluster} RESET (REPLICATION FACTOR)", + http=Http.RANDOM, + ) + self._resize_replicas(cluster, 1) + return True + + def _resize_replicas(self, cluster: Cluster, count: int) -> None: + # Managed cluster replicas are server-named (r1..rN), so the tracked + # list is rebuilt from scratch to match: the ids are what + # `ClusterReplica.name()` derives those names from. + cluster.replicas = [ + ClusterReplica(i, cluster.size, cluster) for i in range(count) + ] + cluster.replica_id = count + + +class DropSchemaCascadeAction(Action): + """DROP SCHEMA .. CASCADE, an atomic multi-object catalog mutation. + + Only enabled in DDL complexity: cross-schema dependents are cascade-dropped + server-side but stay tracked until they surface as "does not exist", which + DDL complexity ignores.""" + + def applicable(self, exe: Executor) -> bool: + return exe.db.complexity in (Complexity.DDL, Complexity.DDLOnly) + + def run(self, exe: Executor) -> bool: + with exe.db.lock: + if len(exe.db.schemas) <= 1: + return False + schema = self.rng.choice(exe.db.schemas) + # Keep at least two non-temp tables alive outside the dropped + # schema. Query generation picks from the non-view objects, and a + # CASCADE that emptied that set would crash it. This mirrors + # DropTableAction's minimum. + if ( + len([t for t in exe.db.tables if not t.temp and t.schema is not schema]) + < 2 + ): + return False + with schema.lock: + if schema not in exe.db.schemas: + return False + if len(exe.db.schemas) <= 1: + return False + exe.execute(f"DROP SCHEMA {schema} CASCADE", http=Http.RANDOM) + exe.db.schemas.remove(schema) + untrack_objects_in_schemas(exe, {schema}) + return True + + +class DropDatabaseCascadeAction(Action): + """DROP DATABASE .. CASCADE, an atomic multi-object catalog mutation. + + DDL-complexity only, for the same reason as DropSchemaCascadeAction.""" + + def applicable(self, exe: Executor) -> bool: + return exe.db.complexity in (Complexity.DDL, Complexity.DDLOnly) + + def run(self, exe: Executor) -> bool: + with exe.db.lock: + if len(exe.db.dbs) <= 1: + return False + db = self.rng.choice(exe.db.dbs) + # Keep at least two non-temp tables alive outside the dropped + # database, so query generation always has a non-view object. + if ( + len([t for t in exe.db.tables if not t.temp and t.schema.db is not db]) + < 2 + ): + return False + with db.lock: + if db not in exe.db.dbs: + return False + if len(exe.db.dbs) <= 1: + return False + exe.execute(f"DROP DATABASE {db} CASCADE", http=Http.RANDOM) + exe.db.dbs.remove(db) + with exe.db.lock: + dropped = {s for s in exe.db.schemas if s.db is db} + exe.db.schemas[:] = [s for s in exe.db.schemas if s.db is not db] + untrack_objects_in_schemas(exe, dropped) + return True + + +class CreateTypeAction(Action): + def run(self, exe: Executor) -> bool: + with exe.db.lock: + if len(exe.db.types) >= MAX_TYPES: + return False + type_id = exe.db.type_id + exe.db.type_id += 1 + try: + schema = self.rng.choice(exe.db.schemas) + except IndexError: + return False + with schema.lock: + if schema not in exe.db.schemas: + return False + typ = Type(type_id, schema, self.rng) + typ.create(exe) + exe.db.types.append(typ) + return True + + +class DropTypeAction(Action): + def errors_to_ignore(self, exe: Executor) -> list[str]: + return [ + "still depended upon by", + "cannot be dropped", + # Another worker (or a CASCADE drop of the schema/database) can + # drop the type first. + "does not exist", + ] + super().errors_to_ignore(exe) + + def run(self, exe: Executor) -> bool: + with exe.db.lock: + if not exe.db.types: + return False + typ = self.rng.choice(exe.db.types) + with typ.lock: + if typ not in exe.db.types: + return False + exe.execute(f"DROP TYPE {typ}", http=Http.RANDOM) + exe.db.types.remove(typ) + return True + + +class CreateNetworkPolicyAction(Action): + def run(self, exe: Executor) -> bool: + with exe.db.lock: + if len(exe.db.network_policies) >= MAX_NETWORK_POLICIES: + return False + policy_id = exe.db.network_policy_id + exe.db.network_policy_id += 1 + policy = NetworkPolicy(policy_id, self.rng) + policy.create(exe) + with exe.db.lock: + exe.db.network_policies.append(policy) + return True + + +class AlterNetworkPolicyAction(Action): + def run(self, exe: Executor) -> bool: + with exe.db.lock: + if not exe.db.network_policies: + return False + policy = self.rng.choice(exe.db.network_policies) + with policy.lock: + if policy not in exe.db.network_policies: + return False + exe.execute( + f"ALTER NETWORK POLICY {policy} SET ({policy.rules_clause()})", + http=Http.RANDOM, + ) + policy.num_rules = self.rng.randint(1, 3) + return True + + +class DropNetworkPolicyAction(Action): + def errors_to_ignore(self, exe: Executor) -> list[str]: + return [ + # The policy is installed as a default somewhere (should not happen, + # we never install ours, but be safe). + "cannot be dropped", + # Another worker dropped the same policy first. The error carries + # the raw name ('netpol-N', not the quoted form), so DROP resolves + # the name correctly. This is a concurrency race, not the ALTER + # NETWORK POLICY quoted-name bug. + "unknown network policy", + ] + super().errors_to_ignore(exe) + + def run(self, exe: Executor) -> bool: + with exe.db.lock: + if not exe.db.network_policies: + return False + policy = self.rng.choice(exe.db.network_policies) + with policy.lock: + if policy not in exe.db.network_policies: + return False + exe.execute(f"DROP NETWORK POLICY {policy}", http=Http.RANDOM) + exe.db.network_policies.remove(policy) + return True + + +class SystemCatalogReadAction(Action): + """Read a random system-catalog / introspection relation while DDL churns. + + Exercises catalog-read-vs-write consistency and the auto-route path. + Relations that the `materialize` user cannot read (mz_notices, + mz_recent_activity_log_redacted) are left out.""" + + def errors_to_ignore(self, exe: Executor) -> list[str]: + return [ + # An introspection view over many objects can outrun + # statement_timeout, never a real bug in a stress run. + "canceling statement due to statement timeout", + # Reading a system-catalog relation inside a read transaction that + # already touched user objects crosses timedomains. + "in the same timedomain", + 'is not allowed from the "mz_catalog_server" cluster', + ] + super().errors_to_ignore(exe) + + RELATIONS = [ + "mz_catalog.mz_objects", + "mz_catalog.mz_columns", + "mz_catalog.mz_indexes", + "mz_catalog.mz_sources", + "mz_catalog.mz_sinks", + "mz_catalog.mz_materialized_views", + "mz_catalog.mz_views", + "mz_catalog.mz_tables", + "mz_catalog.mz_audit_events", + "mz_catalog.mz_databases", + "mz_catalog.mz_schemas", + "mz_catalog.mz_roles", + "mz_catalog.mz_clusters", + "mz_catalog.mz_cluster_replicas", + "mz_internal.mz_frontiers", + "mz_internal.mz_hydration_statuses", + "mz_internal.mz_compute_dependencies", + "mz_internal.mz_source_statuses", + "mz_internal.mz_sink_statuses", + "mz_internal.mz_materialization_lag", + "mz_internal.mz_wallclock_global_lag_recent_history", + "mz_internal.mz_cluster_replica_statuses", + "mz_internal.mz_object_dependencies", + "mz_internal.mz_object_transitive_dependencies", + "mz_internal.mz_show_all_objects", + "mz_internal.mz_comments", + ] + + def run(self, exe: Executor) -> bool: + relation = self.rng.choice(self.RELATIONS) + exe.execute( + f"SELECT * FROM {relation} LIMIT {self.rng.randint(1, 100)}", + http=Http.RANDOM, + fetch=True, + ) + return True + + +class ExplainAnalyzeAction(Action): + """EXPLAIN ANALYZE against a live materialized view or index dataflow. + + Runs generated introspection queries on the active cluster. Racing it + against drop/replace of the target probes the introspection path.""" + + def errors_to_ignore(self, exe: Executor) -> list[str]: + return [ + # The active cluster has more than one replica. + "log source reads must target a replica", + "does not exist", + "not been hydrated", + "not been materialized", + # A concurrent DROP/reconfigure of the targeted replica retires the + # introspection query. No panic in services.log, just a race. + "target replica failed or was dropped", + # Introspection over a large dataflow can outrun statement_timeout. + "canceling statement due to statement timeout", + # The generated introspection queries can cross timedomains or be + # routed to mz_catalog_server. + "in the same timedomain", + 'is not allowed from the "mz_catalog_server" cluster', + ] + super().errors_to_ignore(exe) + + def run(self, exe: Executor) -> bool: + mvs = [v for v in exe.db.views if v.materialized] + with exe.db.lock: + indexes = list(exe.db.indexes) + candidates: list[tuple[str, DBObject | Index]] = [] + for v in mvs: + candidates.append(("MATERIALIZED VIEW", v)) + for i in indexes: + candidates.append(("INDEX", i)) + if not candidates: + return False + kind, obj = self.rng.choice(candidates) + analysis = self.rng.choice( + ["MEMORY", "CPU", "MEMORY WITH SKEW", "CPU WITH SKEW", "HINTS"] + ) + with obj.lock: + if kind == "MATERIALIZED VIEW" and obj not in exe.db.views: + return False + if kind == "INDEX" and obj not in exe.db.indexes: + return False + exe.execute( + f"EXPLAIN ANALYZE {analysis} FOR {kind} {obj}", + http=Http.NO, + fetch=True, + ) + return True + + +class ExplainFilterPushdownAction(Action): + """EXPLAIN FILTER PUSHDOWN, which inspects durable persist state to compute + which parts a query's filters would read.""" + + def errors_to_ignore(self, exe: Executor) -> list[str]: + result = super().errors_to_ignore(exe) + result.extend( + [ + "in the same timedomain", + 'is not allowed from the "mz_catalog_server" cluster', + # Scanning persist part stats can outrun statement_timeout. + "canceling statement due to statement timeout", + ] + ) + if exe.db.complexity == Complexity.DDL: + result.extend(["does not exist"]) + return result + + def run(self, exe: Executor) -> bool: + mvs = [v for v in exe.db.views if v.materialized] + if mvs and self.rng.choice([True, False]): + view = self.rng.choice(mvs) + with view.lock: + if view not in exe.db.views: + return False + exe.execute( + f"EXPLAIN FILTER PUSHDOWN FOR MATERIALIZED VIEW {view}", + http=Http.NO, + fetch=True, + ) + else: + query = self.generate_select_query(exe, ExprKind.ALL) + exe.execute( + f"EXPLAIN FILTER PUSHDOWN FOR {query}", http=Http.NO, fetch=True + ) + return True + + class SourceSinkStallCheckAction(Action): def applicable(self, exe: Executor) -> bool: return exe.db.scenario not in ( @@ -3924,6 +6060,7 @@ def run(self, exe: Executor) -> bool: ("postgres_sources", exe.db.postgres_sources), ("mysql_sources", exe.db.mysql_sources), ("sql_server_sources", exe.db.sql_server_sources), + ("loadgen_sources", exe.db.loadgen_sources), ("webhook_sources", exe.db.webhook_sources), ]: counts = [] @@ -3935,6 +6072,85 @@ def run(self, exe: Executor) -> bool: return True +class DependencyConsistencyAction(Action): + """Client-side catalog dependency oracle: flag any dependency edge whose + endpoints are not both live objects (a dangling uses/used_by edge). + + PREPARED BUT DISABLED (commented out of read_action_list). This targets the + SQL-521 class: a sink left pointing at a materialized view that was dropped + under cancel, i.e. the same MissingUses inconsistency the coordinator's own + check_consistency asserts on. While SQL-521 is open the workload + legitimately produces such danglers, so enabling this now would re-detect + the known corruption rather than find new bugs. + TODO: enable in read_action_list once SQL-521 is fixed. Verify the + mz_internal.mz_object_dependencies column names (object_id, + referenced_object_id) still hold when enabling.""" + + def errors_to_ignore(self, exe: Executor) -> list[str]: + return [ + # Reading a catalog relation inside a read txn that already touched + # user objects crosses timedomains. + "in the same timedomain", + ] + super().errors_to_ignore(exe) + + def run(self, exe: Executor) -> bool: + exe.execute( + "SELECT d.object_id, d.referenced_object_id " + "FROM mz_internal.mz_object_dependencies d " + "LEFT JOIN mz_objects o1 ON d.object_id = o1.id " + "LEFT JOIN mz_objects o2 ON d.referenced_object_id = o2.id " + "WHERE o1.id IS NULL OR o2.id IS NULL", + http=Http.NO, + ) + dangling = exe.cur.fetchall() + if dangling: + raise ValueError( + f"dangling catalog dependency edges (SQL-521 class): {dangling[:5]}" + ) + return True + + +class SourceReadHoldSweepAction(Action): + """Read a source-backed relation to force read-hold acquisition on the + source's remap shard. + + PREPARED BUT DISABLED (commented out of read_action_list). Intended for the + kill scenario, where racing envd/clusterd restarts stresses read-hold + reinstatement: the SS-346 class (a dependent's read hold on a source's remap + shard not upheld across a restart, so its since advances past the + dependent's upper) and PER-49 (a compute import as_of behind the compacted + since after ALTER TABLE ADD COLUMN + kill). Enabling it now just re-triggers + those known coordinator/compute panics. + TODO: enable in read_action_list once SS-346 and PER-49 are fixed.""" + + def errors_to_ignore(self, exe: Executor) -> list[str]: + return [ + "in the same timedomain", + ] + super().errors_to_ignore(exe) + + def run(self, exe: Executor) -> bool: + with exe.db.lock: + sources = [ + o + for o in exe.db.db_objects() + if isinstance( + o, + LoadGeneratorSource + | KafkaSource + | PostgresSource + | MySqlSource + | SqlServerSource + | WebhookSource, + ) + ] + if not sources: + return False + obj = self.rng.choice(sources) + exe.execute(f"SELECT count(*) FROM {obj}", http=Http.RANDOM) + exe.cur.fetchall() + return True + + class ActionList: action_classes: list[type[Action]] weights: list[float] @@ -3952,15 +6168,31 @@ def __init__( [ (SelectAction, 100), (SelectOneAction, 1), + (ParameterizedQueryAction, 20), # (SQLsmithAction, 30), # Questionable use ( CopyToS3Action, 100, ), + (CopyToStdoutAction, 20), + (ShowAction, 10), + (SystemCatalogReadAction, 10), + # TODO: Reenable once EXPLAIN FILTER PUSHDOWN can no longer panic the + # coordinator when a referenced compute collection is concurrently + # dropped. sequence_explain_pushdown -> acquire_read_holds().expect( + # "missing compute collection") at read_policy.rs:389 (normal peeks and + # EXPLAIN ANALYZE handle the drop gracefully). + # See https://linear.app/materializeinc/issue/SQL-519 + # (ExplainFilterPushdownAction, 5), + # PREPARED BUT DISABLED (see class docstrings): enabling these now just + # re-detects known-unfixed coordinator bugs rather than finding new ones. + # (DependencyConsistencyAction, 5), # TODO: enable once SQL-521 fixed + # (SourceReadHoldSweepAction, 5), # TODO: enable once SS-346 & PER-49 fixed (SetClusterAction, 1), (CommitRollbackAction, 30), (ReconnectAction, 1), (FlipFlagsAction, 2), + (SetSessionVariableAction, 2), ], autocommit=False, ) @@ -4006,6 +6238,8 @@ def __init__( (SetClusterAction, 1), (ReconnectAction, 1), (FlipFlagsAction, 2), + (SetSessionVariableAction, 2), + (DiscardAction, 2), # TODO: Reenable when SS-193 and SS-325 are fixed # (SourceSinkStallCheckAction, 4), # (TransactionIsolationAction, 1), @@ -4021,10 +6255,13 @@ def __init__( (DropTableAction, 2), (CreateViewAction, 8), (DropViewAction, 8), + (CreateOrReplaceViewAction, 4), (CreateRoleAction, 2), (DropRoleAction, 2), + (AlterRoleAction, 2), (CreateClusterAction, 1), (DropClusterAction, 1), + (AlterClusterSetAction, 3), (SwapClusterAction, 10), (CreateClusterReplicaAction, 2), (DropClusterReplicaAction, 2), @@ -4038,8 +6275,13 @@ def __init__( (DropIcebergSinkAction, 4), (CreateKafkaSourceAction, 4), (DropKafkaSourceAction, 4), - (CreateMySqlSourceAction, 4), - (DropMySqlSourceAction, 4), + (CreateLoadGeneratorSourceAction, 4), + (DropLoadGeneratorSourceAction, 4), + (CreateMultiLoadGeneratorSourceAction, 2), + (DropMultiLoadGeneratorSourceAction, 2), + # TODO: Reenable when https://linear.app/materializeinc/issue/SS-307 is fixed + # (CreateMySqlSourceAction, 4), + # (DropMySqlSourceAction, 4), (CreatePostgresSourceAction, 4), (DropPostgresSourceAction, 4), # TODO: Reenable when https://linear.app/materializeinc/issue/SS-290 is fixed @@ -4047,18 +6289,64 @@ def __init__( # (DropSqlServerSourceAction, 4), (GrantPrivilegesAction, 4), (RevokePrivilegesAction, 1), + (GrantRoleAction, 2), + (RevokeRoleAction, 1), + (AlterOwnerAction, 2), + (AlterDefaultPrivilegesAction, 2), + (BroadPrivilegesAction, 2), + (ShowAction, 4), + (ValidateConnectionAction, 2), + # TODO: Reenable once altering a connection that sinks or sources depend + # on can no longer panic the coordinator. Re-altering a dependent sink's + # export connection after the txn fails with InvalidAlter, which + # unwrap_or_terminate turns into a panic. + # See https://linear.app/materializeinc/issue/SQL-517 + # (AlterConnectionAction, 2), + (AlterSecretAction, 2), (ReconnectAction, 1), (CreateDatabaseAction, 1), (DropDatabaseAction, 1), + # TODO: Reenable once a concurrent DROP DATABASE CASCADE can no longer + # panic the coordinator. A staged create (e.g. the source executor's + # CREATE SECRET) whose target database is dropped between staging and + # finish hits resolve_full_name -> get_database (panicking OrdMap index) + # in catalog transact_op. Only CASCADE can drop a non-empty database, + # so this is the precise trigger. + # See https://linear.app/materializeinc/issue/SQL-518 + # (DropDatabaseCascadeAction, 1), (CreateSchemaAction, 1), (DropSchemaAction, 1), + (DropSchemaCascadeAction, 1), + (CreateTypeAction, 2), + (DropTypeAction, 2), + (CreateNetworkPolicyAction, 1), + # TODO: Reenable once ALTER NETWORK POLICY resolves quoted (e.g. + # hyphenated) names. It looks the policy up by its quoted display form, + # so it fails with "unknown network policy" for any name that requires + # quoting, even though CREATE and DROP work. + # See https://linear.app/materializeinc/issue/CLO-143 + # (AlterNetworkPolicyAction, 1), + (DropNetworkPolicyAction, 1), (RenameSchemaAction, 10), (RenameTableAction, 10), (RenameViewAction, 10), (RenameKafkaSinkAction, 10), (RenameIcebergSinkAction, 10), (SwapSchemaAction, 10), - (ReplaceMaterializedViewAction, 20), + (CreateReplacementMaterializedViewAction, 10), + (ApplyReplacementMaterializedViewAction, 5), + (DropReplacementMaterializedViewAction, 5), + (SealedCollectionCheckAction, 2), + (TransactionIsolationAction, 1), + (BoundedStalenessReadAction, 2), + (ReadOnlyTransactionAction, 3), + (DDLTransactionAction, 2), + (SystemCatalogReadAction, 4), + (ExplainAnalyzeAction, 4), + # TODO: Reenable with EXPLAIN FILTER PUSHDOWN's coordinator panic on a + # concurrently-dropped compute collection (read_policy.rs:389). + # See https://linear.app/materializeinc/issue/SQL-519 + # (ExplainFilterPushdownAction, 2), (FlipFlagsAction, 2), # TODO: Reenable when https://linear.app/materializeinc/issue/SQL-405 is fixed. # (AlterTableAddColumnAction, 10), diff --git a/misc/python/materialize/parallel_workload/column.py b/misc/python/materialize/parallel_workload/column.py index 9a874d71443c6..9ec25e15de454 100644 --- a/misc/python/materialize/parallel_workload/column.py +++ b/misc/python/materialize/parallel_workload/column.py @@ -88,7 +88,11 @@ def create(self) -> str: return result -class WebhookColumn(Column): +class SourceColumn(Column): + """Column of a source, taking its name from the upstream schema instead of + generating one. The per-source subclasses only exist to keep the column + lists of the different source kinds distinct in type annotations.""" + def __init__( self, name: str, @@ -105,69 +109,25 @@ def name(self, in_query: bool = False) -> str: return identifier(self.raw_name) if in_query else self.raw_name -class KafkaColumn(Column): - def __init__( - self, - name: str, - data_type: type[DataType], - nullable: bool, - db_object: "DBObject", - ): - self.raw_name = name - self.data_type = data_type - self.nullable = nullable - self.db_object = db_object +class WebhookColumn(SourceColumn): + pass - def name(self, in_query: bool = False) -> str: - return identifier(self.raw_name) if in_query else self.raw_name +class KafkaColumn(SourceColumn): + pass -class MySqlColumn(Column): - def __init__( - self, - name: str, - data_type: type[DataType], - nullable: bool, - db_object: "DBObject", - ): - self.raw_name = name - self.data_type = data_type - self.nullable = nullable - self.db_object = db_object - def name(self, in_query: bool = False) -> str: - return identifier(self.raw_name) if in_query else self.raw_name +class LoadGeneratorColumn(SourceColumn): + pass -class PostgresColumn(Column): - def __init__( - self, - name: str, - data_type: type[DataType], - nullable: bool, - db_object: "DBObject", - ): - self.raw_name = name - self.data_type = data_type - self.nullable = nullable - self.db_object = db_object +class MySqlColumn(SourceColumn): + pass - def name(self, in_query: bool = False) -> str: - return identifier(self.raw_name) if in_query else self.raw_name +class PostgresColumn(SourceColumn): + pass -class SqlServerColumn(Column): - def __init__( - self, - name: str, - data_type: type[DataType], - nullable: bool, - db_object: "DBObject", - ): - self.raw_name = name - self.data_type = data_type - self.nullable = nullable - self.db_object = db_object - def name(self, in_query: bool = False) -> str: - return identifier(self.raw_name) if in_query else self.raw_name +class SqlServerColumn(SourceColumn): + pass diff --git a/misc/python/materialize/parallel_workload/database.py b/misc/python/materialize/parallel_workload/database.py index 8920be86fbad6..b3d72d2fbb281 100644 --- a/misc/python/materialize/parallel_workload/database.py +++ b/misc/python/materialize/parallel_workload/database.py @@ -47,6 +47,7 @@ from materialize.parallel_workload.column import ( Column, KafkaColumn, + LoadGeneratorColumn, MySqlColumn, PostgresColumn, SqlServerColumn, @@ -60,7 +61,18 @@ MAX_COLUMNS = 50 MAX_INCLUDE_HEADERS = 5 -MAX_ROWS = 500 +# The row count a view's join squares. A view over two tables joins them on a +# random boolean predicate, so the join is a cross product and every read of +# that view builds MAX_ROWS**2 wide intermediate rows, whatever LIMIT the +# reading query carries. Measured at 500 with only four columns, one +# `COPY (SELECT * FROM view WHERE .. LIMIT 100) TO 's3://..'` cost ~450 MiB of +# replica RSS, and pw tables carry up to MAX_COLUMNS of them. A handful of such +# reads at once is what grew a clusterd to 19 GiB and had the kernel OOM-killer +# take it, plus environmentd and the Kafka broker, out of the one cgroup all of +# a run's containers share (nightlies 17660-17701). 100 keeps enough rows for +# the DML, persist and index paths to be interesting while cutting the +# intermediate 25-fold. +MAX_ROWS = 100 MAX_CLUSTERS = 4 MAX_CLUSTER_REPLICAS = 2 MAX_DBS = 50 @@ -74,8 +86,11 @@ MAX_MYSQL_SOURCES = 50 MAX_SQL_SERVER_SOURCES = 50 MAX_POSTGRES_SOURCES = 50 +MAX_LOADGEN_SOURCES = 50 MAX_KAFKA_SINKS = 50 MAX_ICEBERG_SINKS = 50 +MAX_TYPES = 50 +MAX_NETWORK_POLICIES = 30 MAX_INITIAL_DBS = 1 MAX_INITIAL_SCHEMAS = 1 @@ -88,6 +103,7 @@ MAX_INITIAL_MYSQL_SOURCES = 1 MAX_INITIAL_SQL_SERVER_SOURCES = 1 MAX_INITIAL_POSTGRES_SOURCES = 1 +MAX_INITIAL_LOADGEN_SOURCES = 1 class BodyFormat(Enum): @@ -165,6 +181,13 @@ def name(self) -> str: class DBObject: columns: list[Column] lock: threading.Lock + # Whether reading from this object can legitimately reach the empty + # (sealed) frontier: its own shard seals in normal operation, or a + # dataflow reading it sees an input frontier that becomes empty. Plain + # tables and unbounded sources never seal, so the default is False. + # Bounded (UP TO) load generators seal once they finish, and views + # propagate sealing from their inputs, materialized or not. + can_seal: bool = False def __init__(self): self.lock = threading.Lock() @@ -283,8 +306,9 @@ def __init__( rng.choice( [ "ON COMMIT", - f"EVERY '{rng.randint(1, 60)} seconds {rng.randint(0, 60)} minutes'", - f"EVERY '{rng.randint(1, 60)} seconds {rng.randint(0, 60)} minutes' ALIGNED TO (mz_now())", + # TODO: Restore minute-scale intervals when CPU-196 is fixed + f"EVERY '{rng.randint(1, 15)} seconds'", + f"EVERY '{rng.randint(1, 15)} seconds' ALIGNED TO (mz_now())", # Always in the future of all refreshes of previously generated MVs "AT mz_now()::string::int8 + 1000", ] @@ -293,6 +317,21 @@ def __init__( else None ) + # A materialized view's shard seals (its write frontier advances to the + # empty frontier) once it can never produce more output. That happens + # for REFRESH AT views after their last refresh, for repeat_row(-1) + # constant views on hydration, and transitively for any view that reads + # from an input that itself seals. The replacement and sealed-shard + # oracles key off this to tell legitimate seals from wrongly finalized + # shards. Unmaterialized views have no shard, but a dataflow reading + # one inlines its inputs, so sealing must propagate through them too. + self.can_seal = ( + (self.refresh or "").startswith("AT") + or self.repeat_row_const + or base_object.can_seal + or (base_object2 is not None and base_object2.can_seal) + ) + if base_object2: self.on_expr = expression( Boolean, all_columns, rng, kind=ExprKind.MATERIALIZABLE @@ -370,8 +409,12 @@ def select_str(exprs: str) -> str: return query - def create(self, exe: Executor) -> None: + def create(self, exe: Executor, or_replace: bool = False) -> None: query = "CREATE " + # OR REPLACE keeps the same catalog item, swapping its definition and + # rebuilding the dataflow. It is incompatible with TEMP. + if or_replace and not self.temp: + query += "OR REPLACE " if self.temp: query += "TEMP " if self.materialized: @@ -576,14 +619,19 @@ def __init__( self.cluster = cluster self.schema = schema self.base_object = base_object - key_cols = [ - column - for column in rng.sample( - base_object.columns, k=rng.randint(1, len(base_object.columns)) - ) - ] - key_col_names = [column.name(True) for column in key_cols] - self.key = f"KEY ({', '.join(key_col_names)}) NOT ENFORCED" + self.mode = rng.choice(["UPSERT", "APPEND"]) + if self.mode == "UPSERT": + key_cols = [ + column + for column in rng.sample( + base_object.columns, k=rng.randint(1, len(base_object.columns)) + ) + ] + key_col_names = [column.name(True) for column in key_cols] + self.key = f"KEY ({', '.join(key_col_names)}) NOT ENFORCED" + else: + # APPEND mode does not permit a KEY. + self.key = "" self.table_name = f"icesink_topic{self.sink_id}_{uuid.uuid4().hex[:8]}" self.rename = 0 @@ -596,7 +644,7 @@ def __str__(self) -> str: return f"{self.schema}.{identifier(self.name())}" def create(self, exe: Executor) -> None: - query = f"CREATE SINK {self} IN CLUSTER {self.cluster} FROM {self.base_object} INTO ICEBERG CATALOG CONNECTION polaris_conn (NAMESPACE 'default_namespace', TABLE '{self.table_name}') USING AWS CONNECTION aws_conn {self.key} MODE UPSERT WITH (COMMIT INTERVAL '1s')" + query = f"CREATE SINK {self} IN CLUSTER {self.cluster} FROM {self.base_object} INTO ICEBERG CATALOG CONNECTION polaris_conn (NAMESPACE 'default_namespace', TABLE '{self.table_name}') USING AWS CONNECTION aws_conn {self.key} MODE {self.mode} WITH (COMMIT INTERVAL '1s')" exe.execute(query) @@ -608,6 +656,8 @@ class KafkaSink(DBObject): base_object: DBObject envelope: str key: str + connection_options: list[str] + no_snapshot: bool def __init__( self, @@ -622,6 +672,19 @@ def __init__( self.cluster = cluster self.schema = schema self.base_object = base_object + self.connection_options = [] + if rng.random() < 0.3: + compression = rng.choice(["none", "gzip", "lz4", "zstd", "snappy"]) + self.connection_options.append(f"COMPRESSION TYPE = '{compression}'") + if rng.random() < 0.2: + self.connection_options.append( + f"TRANSACTIONAL ID PREFIX 'pw-txn-{self.sink_id}'" + ) + if rng.random() < 0.2: + self.connection_options.append( + f"PROGRESS GROUP ID PREFIX 'pw-progress-{self.sink_id}'" + ) + self.no_snapshot = rng.random() < 0.2 universal_formats = [ "FORMAT AVRO USING CONFLUENT SCHEMA REGISTRY CONNECTION csr_conn", "FORMAT JSON", @@ -681,7 +744,10 @@ def create(self, exe: Executor) -> None: if self.partition_count else "" ) - query = f"CREATE SINK {self} IN CLUSTER {self.cluster} FROM {self.base_object} INTO KAFKA CONNECTION kafka_conn (TOPIC {topic}{maybe_partition}) {self.key} {self.format} ENVELOPE {self.envelope}" + options = "".join(f", {option}" for option in self.connection_options) + query = f"CREATE SINK {self} IN CLUSTER {self.cluster} FROM {self.base_object} INTO KAFKA CONNECTION kafka_conn (TOPIC {topic}{maybe_partition}{options}) {self.key} {self.format} ENVELOPE {self.envelope}" + if self.no_snapshot: + query += " WITH (SNAPSHOT = false)" exe.execute(query) @@ -861,6 +927,110 @@ def create(self, exe: Executor) -> None: self.executor.create(logging_exe=exe) +class LoadGeneratorSource(DBObject): + """A COUNTER load generator source. Always bounded by UP TO so it cannot + grow without limit, which also exercises the finished-source lifecycle + state. The readable object (str(self)) is the table created from the + source, matching the source-table model the other sources use.""" + + # A finished (UP TO reached) counter advances the source table's frontier + # to the empty antichain, so anything reading it seals legitimately. + can_seal = True + + source_id: int + cluster: "Cluster" + schema: Schema + columns: list[LoadGeneratorColumn] + tick_interval: str + up_to: int + + def __init__( + self, + source_id: int, + cluster: "Cluster", + schema: Schema, + rng: random.Random, + ): + super().__init__() + self.source_id = source_id + self.cluster = cluster + self.schema = schema + self.tick_interval = rng.choice(["10ms", "100ms", "1s"]) + # Bounded near MAX_ROWS: this source-table is a readable relation that + # views/MVs join over, and unlike user tables it is not capped by + # MAX_ROWS. A large counter feeding a maintained join is a clusterd + # memory risk, so keep it in the same size class as the other data. + self.up_to = rng.randint(1, MAX_ROWS) + self.columns = [LoadGeneratorColumn("counter", Long, False, self)] + + def name(self) -> str: + return naughtify(f"lg-{self.source_id}") + + def source_name(self) -> str: + return naughtify(f"lg-src-{self.source_id}") + + def __str__(self) -> str: + return f"{self.schema}.{identifier(self.name())}" + + def create(self, exe: Executor) -> None: + source = f"{self.schema}.{identifier(self.source_name())}" + exe.execute( + f"CREATE SOURCE {source} IN CLUSTER {self.cluster} FROM LOAD GENERATOR COUNTER (TICK INTERVAL '{self.tick_interval}', UP TO {self.up_to})" + ) + exe.execute(f"CREATE TABLE {self} FROM SOURCE {source}") + + +class MultiLoadGeneratorSource: + """A multi-subsource load generator (AUCTION / TPCH / MARKETING) created + with FOR ALL TABLES. Tracked create/drop-only, not as a readable relation: + its subsources have fixed names and hardcoded schemas, so wiring them into + the general read/expression machinery is out of scope. The value is the + multi-subsource source lifecycle and CASCADE teardown under concurrency. + + At most one source per generator type exists at a time, because FOR ALL + TABLES names its subsources fixedly (a second AUCTION would collide on + `auctions`, `bids`, ...).""" + + source_id: int + cluster: "Cluster" + schema: Schema + generator: str + tick_interval: str + lock: threading.Lock + + def __init__( + self, + source_id: int, + cluster: "Cluster", + schema: Schema, + generator: str, + rng: random.Random, + ): + self.source_id = source_id + self.cluster = cluster + self.schema = schema + self.generator = generator + self.tick_interval = rng.choice(["100ms", "1s"]) + self.lock = threading.Lock() + + def name(self) -> str: + return naughtify(f"mlg-{self.source_id}") + + def __str__(self) -> str: + return f"{self.schema}.{identifier(self.name())}" + + def create(self, exe: Executor) -> None: + options = [f"TICK INTERVAL '{self.tick_interval}'"] + if self.generator == "TPCH": + options.append("SCALE FACTOR 0.0001") + query = ( + f"CREATE SOURCE {self} IN CLUSTER {self.cluster} " + f"FROM LOAD GENERATOR {self.generator} ({', '.join(options)}) " + f"FOR ALL TABLES" + ) + exe.execute(query) + + class S3Object(DBObject): """A COPY TO dump of `table` that CopyFromS3Action can load back into it. @@ -1066,6 +1236,11 @@ def __str__(self) -> str: def create(self, exe: Executor) -> None: exe.execute(f"CREATE ROLE {self}") + # Make the workload's own user (materialize) a member, so it can + # manage the role, e.g. transfer object ownership to it: ALTER .. + # OWNER TO requires the executor to be a member of . + # The creating session has admin on the role it just created. + exe.execute(f"GRANT {self} TO materialize") class ClusterReplica: @@ -1083,6 +1258,12 @@ def __init__(self, replica_id: int, size: str, cluster: "Cluster"): self.lock = threading.Lock() def name(self) -> str: + # A managed cluster's replicas are named by the controller as r1..rN, + # never by us, so neither the rename counter nor naughtify applies. + # Rendering our own name there would make every reference to the + # replica, e.g. a replica-targeted materialized view, fail to resolve. + if self.cluster.managed: + return f"r{self.replica_id+1}" if self.rename: return naughtify(f"r-{self.replica_id+1}-{self.rename}") return naughtify(f"r-{self.replica_id+1}") @@ -1147,6 +1328,78 @@ def create(self, exe: Executor) -> None: exe.execute(query) +class Type: + """A user-defined type (row, list, or map). Schema-scoped.""" + + type_id: int + schema: Schema + kind: str + lock: threading.Lock + rng: random.Random + + def __init__(self, type_id: int, schema: Schema, rng: random.Random): + self.type_id = type_id + self.schema = schema + self.kind = rng.choice(["row", "list", "map"]) + self.rng = rng + self.lock = threading.Lock() + + def name(self) -> str: + return naughtify(f"type-{self.type_id}") + + def __str__(self) -> str: + return f"{self.schema}.{identifier(self.name())}" + + def create(self, exe: Executor) -> None: + if self.kind == "row": + scalar_types = ["int4", "text", "bool", "float8", "timestamp"] + fields = ", ".join( + f"f{i} {self.rng.choice(scalar_types)}" + for i in range(self.rng.randint(1, 4)) + ) + query = f"CREATE TYPE {self} AS ({fields})" + elif self.kind == "list": + element = self.rng.choice(["int4", "text", "float8"]) + query = f"CREATE TYPE {self} AS LIST (ELEMENT TYPE = {element})" + else: + value = self.rng.choice(["int4", "text", "float8"]) + query = f"CREATE TYPE {self} AS MAP (KEY TYPE = text, VALUE TYPE = {value})" + exe.execute(query) + + +class NetworkPolicy: + """A network policy. Top-level (not schema-scoped), like clusters. + + Rules are always allow-all so the workload can never lock itself out of + its own connections. The policy is never installed as the active + `network_policy` system parameter for the same reason.""" + + policy_id: int + num_rules: int + lock: threading.Lock + + def __init__(self, policy_id: int, rng: random.Random): + self.policy_id = policy_id + self.num_rules = rng.randint(1, 3) + self.lock = threading.Lock() + + def name(self) -> str: + return naughtify(f"netpol-{self.policy_id}") + + def __str__(self) -> str: + return identifier(self.name()) + + def rules_clause(self) -> str: + rules = ", ".join( + f"r{i} (action='allow', direction='ingress', address='0.0.0.0/0')" + for i in range(self.num_rules) + ) + return f"RULES ({rules})" + + def create(self, exe: Executor) -> None: + exe.execute(f"CREATE NETWORK POLICY {self} ({self.rules_clause()})") + + # TODO: Can access both databases from same connection! class Database: complexity: Complexity @@ -1176,10 +1429,18 @@ class Database: postgres_source_id: int sql_server_sources: list[SqlServerSource] sql_server_source_id: int + loadgen_sources: list[LoadGeneratorSource] + loadgen_source_id: int + multi_loadgen_sources: list[MultiLoadGeneratorSource] + multi_loadgen_source_id: int iceberg_sinks: list[IcebergSink] iceberg_sink_id: int kafka_sinks: list[KafkaSink] kafka_sink_id: int + types: list[Type] + type_id: int + network_policies: list[NetworkPolicy] + network_policy_id: int s3_path: int s3_objects: list[S3Object] read_then_write_counter: ReadThenWriteCounter @@ -1243,9 +1504,7 @@ def __init__( Cluster( i, managed=rng.choice([True, False]), - size=rng.choice( - ["scale=1,workers=1", "scale=1,workers=4", "scale=2,workers=2"] - ), + size=rng.choice(["scale=1,workers=1", "scale=1,workers=2"]), replication_factor=1, introspection_interval="1s", ) @@ -1262,6 +1521,12 @@ def __init__( self.mysql_sources = [] self.postgres_sources = [] self.sql_server_sources = [] + self.loadgen_sources = [ + LoadGeneratorSource( + i, rng.choice(self.clusters), rng.choice(self.schemas), rng + ) + for i in range(rng.randint(0, MAX_INITIAL_LOADGEN_SOURCES)) + ] self.iceberg_sinks = [] self.kafka_sinks = [] self.s3_objects = [] @@ -1269,9 +1534,16 @@ def __init__( self.mysql_source_id = len(self.mysql_sources) self.postgres_source_id = len(self.postgres_sources) self.sql_server_source_id = len(self.sql_server_sources) + self.loadgen_source_id = len(self.loadgen_sources) + self.multi_loadgen_sources = [] + self.multi_loadgen_source_id = 0 self.iceberg_sink_id = len(self.iceberg_sinks) self.kafka_sink_id = len(self.kafka_sinks) self.read_then_write_counter = ReadThenWriteCounter() + self.types = [] + self.type_id = 0 + self.network_policies = [] + self.network_policy_id = 0 self.lock = threading.Lock() self.sqlsmith_state = "" self.flags = {} @@ -1284,6 +1556,7 @@ def db_objects( | PostgresSource | SqlServerSource | KafkaSource + | LoadGeneratorSource | View | Table ]: @@ -1300,6 +1573,7 @@ def db_objects( + self.mysql_sources + self.postgres_sources + self.sql_server_sources + + self.loadgen_sources + self.webhook_sources ) @@ -1311,6 +1585,7 @@ def db_objects_without_views( | PostgresSource | SqlServerSource | KafkaSource + | LoadGeneratorSource | View | Table ]: @@ -1318,6 +1593,28 @@ def db_objects_without_views( obj for obj in self.db_objects() if type(obj) != View or obj.materialized ] + def db_objects_for_sinks( + self, + ) -> list[ + WebhookSource + | MySqlSource + | PostgresSource + | SqlServerSource + | KafkaSource + | View + | Table + ]: + """Objects usable as a sink's input (base object or ALTER SINK SET + FROM target). Load generator source tables are excluded: an + ALTER SINK .. SET FROM one can trigger the sink stall of + https://linear.app/materializeinc/issue/SS-344, which is worse for a + continuously-producing input.""" + return [ + obj + for obj in self.db_objects_without_views() + if not isinstance(obj, LoadGeneratorSource) + ] + def __iter__(self): """Returns all relations""" return ( @@ -1342,6 +1639,13 @@ def create(self, exe: Executor, composition: Composition) -> None: for row in exe.cur.fetchall(): exe.execute(f"DROP ROLE {identifier(row[0])}") + # Network policies survive restarts and are top-level, so leftovers + # from a killed run have to be swept before recreating. + exe.execute("SELECT name FROM mz_internal.mz_network_policies") + for row in exe.cur.fetchall(): + if row[0].startswith("netpol-"): + exe.execute(f"DROP NETWORK POLICY {identifier(row[0])}") + print("Creating connections") exe.execute( diff --git a/misc/python/materialize/parallel_workload/executor.py b/misc/python/materialize/parallel_workload/executor.py index 052f9c0ab8dc9..11bbaf6c052e0 100644 --- a/misc/python/materialize/parallel_workload/executor.py +++ b/misc/python/materialize/parallel_workload/executor.py @@ -53,10 +53,16 @@ class Executor: reconnect_next: bool rollback_next: bool last_log: str + # "running" exactly while this session waits on the server. Every path that + # makes a round trip has to set it, the end-of-run wedge check reads it to + # tell a server-side hang from a worker stuck in the workload's own code. last_status: str action_run_since_last_commit_rollback: bool autocommit: bool user: str + # The session's transaction isolation, as last set through set_isolation. + # Actions that set an isolation transiently restore this one afterwards. + isolation: str def __init__( self, @@ -84,9 +90,17 @@ def __init__( self.use_ws = self.rng.choice([True, False]) if self.ws else False self.autocommit = cur.connection.autocommit self.mz_service = "materialized" + # Set while a non-default statement_timeout is configured on this + # session, statement timeouts are expected errors then. Cleared again + # on RESET and on reconnect, otherwise a hang in this worker stays + # invisible for the rest of the run. + self.statement_timeout_set = False + # Materialize's default, until the worker sets one on connect. + self.isolation = "STRICT SERIALIZABLE" def set_isolation(self, level: str) -> None: self.execute(f"SET TRANSACTION_ISOLATION TO '{level}'") + self.isolation = level def commit(self, http: Http = Http.RANDOM) -> None: self._end_transaction("commit", http) @@ -97,6 +111,7 @@ def rollback(self, http: Http = Http.RANDOM) -> None: def _end_transaction(self, command: str, http: Http) -> None: self.insert_table = None self.log(command) + self.last_status = "running" ws_error = None try: # When this executor uses the WS session, statements executed with @@ -118,6 +133,8 @@ def _end_transaction(self, command: str, http: Http) -> None: raise except Exception as e: raise QueryError(str(e), command) + finally: + self.last_status = "finished" if ws_error is not None: raise ws_error # TODO(def-): Enable when things are stable @@ -182,6 +199,7 @@ def copy( ) -> None: query += ";" self.log(f"{query} ({rows})") + self.last_status = "running" try: try: @@ -195,6 +213,22 @@ def copy( finally: self.last_status = "finished" + def copy_to_stdout(self, query: str) -> None: + query += ";" + self.log(query) + self.last_status = "running" + try: + try: + with self.cur.copy(query.encode()) as copy: + for _ in copy: + pass + except Exception as e: + raise QueryError(str(e), query) + + self.action_run_since_last_commit_rollback = True + finally: + self.last_status = "finished" + def execute( self, query: str, @@ -207,7 +241,37 @@ def execute( http == Http.RANDOM and self.rng.choice([True, False]) ) or http == Http.YES if explainable and self.rng.choice([True, False]): - query = f"EXPLAIN OPTIMIZED PLAN AS VERBOSE TEXT FOR {query}" + if self.rng.random() < 0.1: + as_json = " AS JSON" if self.rng.choice([True, False]) else "" + query = f"EXPLAIN TIMESTAMP{as_json} FOR {query}" + else: + stage = self.rng.choice( + [ + "RAW PLAN", + "DECORRELATED PLAN", + "LOCALLY OPTIMIZED PLAN", + "OPTIMIZED PLAN", + "OPTIMIZED PLAN", + "OPTIMIZED PLAN", + "PHYSICAL PLAN", + ] + ) + modifiers = "" + if stage == "OPTIMIZED PLAN" and self.rng.random() < 0.3: + mods = self.rng.sample( + [ + "arity", + "join implementations", + "keys", + "types", + "humanized expressions", + "redacted", + ], + self.rng.randint(1, 3), + ) + modifiers = f" WITH ({', '.join(mods)})" + format = self.rng.choice(["VERBOSE TEXT", "TEXT", "JSON"]) + query = f"EXPLAIN {stage}{modifiers} AS {format} FOR {query}" query += ";" extra_info_str = f" ({extra_info})" if extra_info else "" use_ws = self.use_ws and http != Http.NO diff --git a/misc/python/materialize/parallel_workload/expression.py b/misc/python/materialize/parallel_workload/expression.py index 72daf7e02d252..0d75239f02f16 100644 --- a/misc/python/materialize/parallel_workload/expression.py +++ b/misc/python/materialize/parallel_workload/expression.py @@ -45,6 +45,7 @@ from materialize.parallel_workload.column import ( Column, KafkaColumn, + LoadGeneratorColumn, MySqlColumn, PostgresColumn, SqlServerColumn, @@ -340,11 +341,33 @@ def __init__(self, text: str, params: list, unsupported: ExprKind = ExprKind.ALL # ] +# Edge-case literals injected at low probability as leaves of read/filter +# expressions (ExprKind.ALL only), to surface NaN/Infinity handling in +# arithmetic, comparisons, and aggregation. The generator's `Float` is float4, +# so these MUST be float4: a wider literal (float8/int8) injected where a float4 +# is expected breaks function overload resolution (e.g. round(numeric, bigint)), +# producing spurious failures rather than findings. +# +# Scope is deliberately narrow: NaN/Infinity are the genuinely-new coverage +# (random_value never produces them). Numeric/int boundary and overflow values +# are NOT injected here, because the existing LARGE record size already +# generates overflow-inducing magnitudes and mis-widthed literals only break +# overload resolution. Extreme-year date/timestamp literals are also omitted: +# a 6-digit year deterministically trips the known SS-345 date-parser bug in +# every context, spamming a known-unfixed issue rather than finding new ones. +EDGE_VALUES: dict[type, list[str]] = { + Float: ["'NaN'::float4", "'Infinity'::float4", "'-Infinity'::float4"], +} + + def expression( data_type: type[DataType], columns: list[Column] | ( list[MySqlColumn] - | (list[PostgresColumn] | (list[SqlServerColumn] | list[KafkaColumn])) + | ( + list[PostgresColumn] + | (list[SqlServerColumn] | (list[KafkaColumn] | list[LoadGeneratorColumn])) + ) ), rng: random.Random, kind: ExprKind = ExprKind.ALL, @@ -371,6 +394,12 @@ def expression( if col.data_type == data_type: return str(col) + # Only in read/filter contexts: overflow/eval errors are tolerated there, + # but not in write (INSERT value) or materialized-view-body contexts. + edges = EDGE_VALUES.get(data_type) + if edges and kind == ExprKind.ALL and rng.random() < 0.2: + return rng.choice(edges) + record_size = rng.choice( [RecordSize.TINY, RecordSize.SMALL, RecordSize.MEDIUM, RecordSize.LARGE] ) diff --git a/misc/python/materialize/parallel_workload/negative_accumulation_errors.py b/misc/python/materialize/parallel_workload/negative_accumulation_errors.py index 715ce5ba2704e..6f60be9596676 100644 --- a/misc/python/materialize/parallel_workload/negative_accumulation_errors.py +++ b/misc/python/materialize/parallel_workload/negative_accumulation_errors.py @@ -31,8 +31,18 @@ "Negative multiplicities in TopK", # Reduce "Net-zero records with non-zero accumulation in ReduceAccumulable", + # Client-facing variant of the above (reduce.rs:1514, EvalError), distinct + # from the internal ReduceAccumulable log text. Seen in repeat_row (#8106). + "with non-zero accumulation in accumulable aggregate", "Non-positive multiplicity in DistinctBy", + # Covers the `ReduceInaccumulable`, `ReduceInaccumulable DISTINCT` and + # `ReduceMinsMaxes` sites, which surface their internal log text verbatim. "Non-positive accumulation", + # The hierarchical min/max stage words its client-facing error differently + # from its "Non-positive accumulation in MinsMaxesHierarchical" log line + # (reduce.rs `build_bucketed_stage`), so the entry above does not cover it. + # Seen in repeat_row (build 17664). + "saw non-positive accumulation", "Invalid negative unsigned aggregation in ReduceAccumulable", "saw negative accumulation", # Peek handling @@ -43,6 +53,9 @@ "S3 oneshot sink encountered negative multiplicities", # Constant folding "Negative multiplicity in constant result", + # Constant folding a DISTINCT/INTERSECT/reduce over a repeat_row collection + # with negative diffs. Seen in repeat_row (builds 17205, 17214). + "constant folding encountered reduce on collection with non-positive multiplicities", # Scalar subquery guard "negative number of rows produced in subquery", ] diff --git a/misc/python/materialize/parallel_workload/parallel_workload.py b/misc/python/materialize/parallel_workload/parallel_workload.py index fb24101ca83af..4a2740fc70842 100644 --- a/misc/python/materialize/parallel_workload/parallel_workload.py +++ b/misc/python/materialize/parallel_workload/parallel_workload.py @@ -9,6 +9,7 @@ import argparse import datetime +import faulthandler import gc import os import random @@ -24,12 +25,17 @@ from materialize.parallel_workload.action import ( Action, ActionList, + ApplyReplacementMaterializedViewAction, BackupRestoreAction, CancelAction, + CreateReplacementMaterializedViewAction, DropClusterAction, DropDatabaseAction, + DropRoleAction, DropSchemaAction, + ExplainAnalyzeAction, KillAction, + SealedCollectionCheckAction, StatisticsAction, ZeroDowntimeDeployAction, action_lists, @@ -45,6 +51,7 @@ MAX_KAFKA_SINKS, MAX_KAFKA_SOURCES, MAX_MYSQL_SOURCES, + MAX_NETWORK_POLICIES, MAX_POSTGRES_SOURCES, MAX_ROLES, MAX_SCHEMAS, @@ -67,6 +74,23 @@ REPORT_TIME = 10 +def run_final_sealed_check( + rng: random.Random, database: Database, cur: psycopg.Cursor +) -> None: + """Guaranteed final oracle pass, before the objects are dropped. + + The finalize task seals a wrongly finalized shard only ~5s after the + triggering DROP, so damage from the run's last actions is not visible + immediately, and a run whose random scheduling never picked + SealedCollectionCheckAction would otherwise end without any pass. + """ + exe = Executor(rng, cur, None, database) + sealed_check = SealedCollectionCheckAction(rng, None) + if sealed_check.applicable(exe): + time.sleep(10) + sealed_check.run(exe) + + def run( host: str, ports: dict[str, int], @@ -123,6 +147,9 @@ def run( system_exe.execute( f"ALTER SYSTEM SET max_roles = {MAX_ROLES * 1000 + num_threads}" ) + system_exe.execute( + f"ALTER SYSTEM SET max_network_policies = {MAX_NETWORK_POLICIES + num_threads}" + ) system_exe.execute( f"ALTER SYSTEM SET max_clusters = {MAX_CLUSTERS * 40 + num_threads}" ) @@ -163,7 +190,7 @@ def run( system_exe.execute("DROP CLUSTER quickstart CASCADE") replica_names = [f"r{replica_id}" for replica_id in range(0, replicas)] replica_string = ",".join( - f"{replica_name} (SIZE 'scale=1,workers=4')" + f"{replica_name} (SIZE 'scale=1,workers=1')" for replica_name in replica_names ) system_exe.execute( @@ -228,10 +255,17 @@ def run( ) workers.append(worker) + # Daemon threads, so that a worker still stuck in a long-running query + # cannot block interpreter shutdown. The clean paths below join + # explicitly, and a WorkerFailedException propagates out of `run` past + # those joins. Without this a single wedged worker keeps the process + # alive until the CI step timeout kills it, which reports the timeout + # instead of the failure that actually happened. thread = threading.Thread( name=thread_name, target=worker.run, args=(host, ports["materialized"], ports["http"], "materialize", database), + daemon=True, ) thread.start() threads.append(thread) @@ -252,6 +286,7 @@ def run( name="cancel", target=worker.run, args=(host, ports["mz_system"], ports["http"], "mz_system", database), + daemon=True, ) thread.start() threads.append(thread) @@ -272,6 +307,7 @@ def run( name="kill", target=worker.run, args=(host, ports["materialized"], ports["http"], "materialize", database), + daemon=True, ) thread.start() threads.append(thread) @@ -299,6 +335,7 @@ def run( name="zero-downtime-deploy", target=worker.run, args=(host, ports["materialized"], ports["http"], "materialize", database), + daemon=True, ) thread.start() threads.append(thread) @@ -319,6 +356,7 @@ def run( name="kill", target=worker.run, args=(host, ports["materialized"], ports["http"], "materialize", database), + daemon=True, ) thread.start() threads.append(thread) @@ -343,6 +381,7 @@ def run( name="statistics", target=worker.run, args=(host, ports["mz_system"], ports["http"], "mz_system", database), + daemon=True, ) thread.start() threads.append(thread) @@ -379,13 +418,79 @@ def run( if all([not thread.is_alive() for thread in threads]): break else: - for worker, thread in zip(workers, threads): - if thread.is_alive(): - print( - f"{thread.name} still running ({worker.exe.mz_service}): {worker.exe.last_log} ({worker.exe.last_status})" - ) + alive = [(w, t) for w, t in zip(workers, threads) if t.is_alive()] + for worker, thread in alive: + if worker.exe is None: + # Still in the initial connect retry loop. + print(f"{thread.name} still connecting") + continue + print( + f"{thread.name} still running ({worker.exe.mz_service}): {worker.exe.last_log} ({worker.exe.last_status})" + ) + + # Workers are daemon threads, so a wedged one cannot keep the process + # alive and has to be caught here or it silently costs the run its + # remaining workload. A worker waiting on the server is a server-side + # hang (DB-118) and is tolerated, as is one still in its connect retry + # loop (exe is None), which can only be hanging inside a connect call + # to the server. One that never waits on the server is stuck in the + # workload's own code, e.g. deadlocked on two object locks taken in + # opposite orders, which no timeout ever clears. Nothing legitimate + # pauses without a round trip for the 300s above (the longest such + # pause is BackupRestoreAction's 240s sleep), but a single sample can + # catch a worker between round trips, so only call a worker wedged if + # it stays off the server for a whole confirmation window. + # + # A tolerated server-side hang also explains every other worker: the + # hung statement holds its actions' object locks (an `ALTER SINK ... + # SET FROM` waits unboundedly for the sink's frontier to catch up, + # database-issues#9820, while holding both the sink's and the new base + # object's lock), and a worker queued on such a lock can no more reach + # the server than a deadlocked one can. Telling the two apart from + # outside is not possible, so a workload-side deadlock is only + # diagnosed when no worker is on the server at all. Waiters and + # deadlocks both resolve when the statement does, and calling a waiter + # deadlocked costs the run its final checks. + def not_on_server() -> set[str]: + return { + t.name + for w, t in alive + if t.is_alive() and w.exe is not None and w.exe.last_status != "running" + } + + def on_server() -> set[str]: + return { + t.name + for w, t in alive + if t.is_alive() and w.exe is not None and w.exe.last_status == "running" + } + + wedged = not_on_server() + confirm_until = time.time() + 60 + while wedged and time.time() < confirm_until: + time.sleep(1) + wedged &= not_on_server() + if on_server(): + wedged = set() + break merge_num_queries(num_queries, workers) - print_stats(num_queries, workers, num_threads, scenario) + print_stats(num_queries, workers, num_threads, complexity, scenario) + + # A wedged worker must not mask damage, so run the final oracle pass + # on a fresh connection before exiting. Skipped for 0dt deploys, + # where connecting to the fenced-out environmentd can hang forever. + if scenario != Scenario.ZeroDowntimeDeploy: + try: + check_conn = psycopg.connect( + host=host, port=ports["materialized"], user="materialize" + ) + except Exception as e: + print(f"Skipping final sealed-collection check: {e}") + else: + check_conn.autocommit = True + with check_conn.cursor() as cur: + run_final_sealed_check(rng, database, cur) + check_conn.close() if num_threads >= 50: # Under high load some queries can't finish quickly, especially UPDATE/DELETE @@ -395,9 +500,13 @@ def run( # environmentd will be stuck forever, the promoted environmentd can # take > 10 minutes to become responsive as well os._exit(0) - # TODO: Reenable when https://linear.app/materializeinc/issue/DB-118 is fixed - # print("Threads have not stopped within 5 minutes, exiting hard") - # os._exit(1) + if wedged: + faulthandler.dump_traceback() + print( + "^^^ +++ Threads have not stopped within 5 minutes and are not" + f" waiting on the server, exiting hard: {', '.join(sorted(wedged))}" + ) + os._exit(1) os._exit(0) try: @@ -422,6 +531,8 @@ def run( # increment can still be in flight. database.read_then_write_counter.validate(exe) + run_final_sealed_check(rng, database, cur) + # Dropping the database also releases the long running connections # used by database objects. database.drop(exe) @@ -446,7 +557,7 @@ def run( conn.close() merge_num_queries(num_queries, workers) - print_stats(num_queries, workers, num_threads, scenario) + print_stats(num_queries, workers, num_threads, complexity, scenario) def merge_num_queries( @@ -474,6 +585,7 @@ def print_stats( num_queries: defaultdict[ActionList, Counter[type[Action]]], workers: list[Worker], num_threads: int, + complexity: Complexity, scenario: Scenario, ) -> None: ignored_errors: defaultdict[str, Counter[type[Action]]] = defaultdict(Counter) @@ -527,7 +639,11 @@ def print_stats( # (ReplaceMaterializedView, whose replacements pile up unfinalized in the # workload). "unknown cluster 'dont_exist'" comes from FlipFlags setting the # default cluster to a nonexistent one to hunt for panics, so any action - # needing a cluster fails with it while that flag is live. Tracked + # needing a cluster fails with it while that flag is live. "cannot be + # dropped because some objects depend on it" is the same rejection for + # DROP ROLE: AlterOwnerAction reassigns object ownership to random roles, + # so a role usually owns something and a short run can see DropRoleAction + # never land a dependency-free role. Tracked # separately so such actions don't trip the broken-action assertion below. noise = { "must be owner of", @@ -537,14 +653,26 @@ def print_stats( "because it already has a replacement", "is sealed and thus cannot be replaced", "unknown cluster 'dont_exist'", + "cannot be dropped because some objects depend on it", } + if complexity in (Complexity.DDL, Complexity.DDLOnly): + # CreateTableAction and CreateViewAction make half their objects TEMP, + # and those land in the shared table and view lists every worker samples + # from. A temporary schema belongs to the session that created it, so + # any other worker resolving "db"."mz_temp"."name" fails with "unknown + # schema 'mz_temp'". That is roughly half the attempts of an action + # picking a random table or view, which an action as low-weight as + # RevokePrivilegesAction can spend a whole run on. Concurrent schema + # drops and renames land here too. Only DDL complexity creates temp + # objects or touches schemas, so elsewhere this stays a genuine signal. + noise.add("unknown schema") if scenario == Scenario.Rename: # Concurrent renames invalidate the qualified names an action captured # earlier (a stored SELECT it re-renders, a target it resolved before # the rename), so a rare action can fail on every attempt with a # name-resolution race. Expected only in this scenario. In Regression, # where nothing renames, these would be a genuine signal. - noise |= {"does not exist", "unknown schema", "unknown catalog item"} + noise |= {"does not exist", "unknown catalog item"} num_errored_real: Counter[type[Action]] = Counter() for worker in workers: num_successes.update(worker.num_successes) @@ -558,11 +686,42 @@ def print_stats( for action_list in action_lists for action_class in action_list.action_classes } - # These use RESTRICT and their targets practically always contain - # objects (schemas and databases their items, clusters their sources, - # sinks, and indexes), so they exercise the rejection path and are not - # expected to ever succeed. - action_classes -= {DropClusterAction, DropDatabaseAction, DropSchemaAction} + # A given churny seed can legitimately see zero successes for these, so the + # broken-action assertion must not fire on them: + # * DropCluster/DropDatabase/DropSchema: RESTRICT rejections. Their + # targets practically always contain objects, and the rejection message + # is class-specific, so `noise` above cannot cover it without excusing + # it for every other action too. + # * DropRoleAction: needs a dependency-free role, but AlterOwnerAction + # keeps reassigning object ownership to random roles. + # * ExplainAnalyzeAction: needs a hydrated MV/index on the active cluster, + # which renames/drops keep retiring. + # These succeed in normal runs, so they aren't broken. + # + # NOTE: listed by name rather than matched on a Drop* name prefix. The other + # Drop* actions pick an existing object and are expected to succeed + # constantly, and a broken one (wrong quoting or qualification) fails with + # `unknown catalog item`, which the base errors_to_ignore tolerates and + # `noise` does not list. That is exactly what the assertion is meant to + # catch, so a prefix would silently exempt the failure mode it exists for. + excluded: set[type[Action]] = { + DropClusterAction, + DropDatabaseAction, + DropRoleAction, + DropSchemaAction, + ExplainAnalyzeAction, + } + if scenario == Scenario.Rename: + # CreateReplacementMaterializedViewAction re-renders the view's SELECT + # with the object names captured at creation. Renames invalidate them, + # so CREATE REPLACEMENT fails with a tolerated "does not exist"/"unknown + # schema", and a churny rename seed can legitimately never land a clean + # attempt. With no replacement ever created, + # ApplyReplacementMaterializedViewAction then has nothing to apply + # either. Both succeed in the other scenarios, so they stay checked + # there. + excluded.add(CreateReplacementMaterializedViewAction) + excluded.add(ApplyReplacementMaterializedViewAction) never_succeeded = [] for action_class in sorted(action_classes, key=lambda cls: cls.__name__): successes = num_successes[action_class] @@ -596,7 +755,7 @@ def print_stats( always_erroring = [ action_class.__name__ for action_class, skips, errored in never_succeeded - if num_errored_real[action_class] > 0 + if num_errored_real[action_class] > 0 and action_class not in excluded ] assert ( not always_erroring @@ -636,7 +795,12 @@ def parse_common_args(parser: argparse.ArgumentParser) -> None: parser.add_argument( "--azurite", action="store_true", help="Use Azurite as blob store instead of S3" ) - parser.add_argument("--replicas", type=int, default=2, help="use multiple replicas") + parser.add_argument( + "--replicas", + type=int, + default=1, + help="default replica number for quickstart cluster", + ) def main() -> int: diff --git a/misc/python/materialize/parallel_workload/settings.py b/misc/python/materialize/parallel_workload/settings.py index 79196897d521a..4efdd32a3876e 100644 --- a/misc/python/materialize/parallel_workload/settings.py +++ b/misc/python/materialize/parallel_workload/settings.py @@ -58,4 +58,28 @@ def _missing_(cls, value): # it on outside that scenario is harmless: no Parallel Workload codegen # emits `repeat_row` unless the scenario is active. "enable_repeat_row": "true", + # TODO: Reenable once the frontend-peek path stops acquiring a read hold + # whose `since` has compacted past the chosen `as_of`. The soft-assert + # assert_read_holds_correct (frontend_peek.rs:1759) fires for peeks in + # multi-statement transactions ("... read hold at .. is not enough for + # as_of .."), panicking the coordinator. Peeks still work via the classic + # coordinator path with this off. + # See https://linear.app/materializeinc/issue/SQL-520 + "enable_frontend_peek_sequencing": "false", + # 64 MiB, down from the 1 GiB default. A peek's result is materialized in + # memory on the replica serving it before this bound errors it + # (`to_error_if_exceeds`, compute_state.rs), and measured on a workload-like + # relation a 300 MiB result costs ~1 GiB of replica RSS plus ~0.5 GiB in + # environmentd. The ceiling is per in-flight statement and every worker + # holds a pg and a WebSocket session, so at the default a handful of + # concurrent wide reads can claim tens of GiB on the default cluster, where + # every peek lands. All containers of a run share one cgroup budget (24 GiB + # on the CI agent, environmentd plus every replica process plus Kafka, + # Postgres, MySQL, SQL Server), so that ends in the kernel OOM-killer + # picking a victim and taking unrelated processes down with it. Runs do + # reach the cap ("result exceeds max size of" shows up in the ignored-error + # statistics), which is what makes the tail dangerous. Query shapes are + # unaffected: an oversized result errors instead of being buffered, and that + # error is already tolerated for every action. + "max_result_size": "67108864", } diff --git a/src/adapter/src/catalog/apply.rs b/src/adapter/src/catalog/apply.rs index 0b70c7d72c847..8f964df74c402 100644 --- a/src/adapter/src/catalog/apply.rs +++ b/src/adapter/src/catalog/apply.rs @@ -281,7 +281,7 @@ impl CatalogState { } if update_system_config { - self.system_configuration.dyncfg_updates(); + self.system_configuration.sync_dyncfgs(); } Ok((builtin_table_updates, catalog_updates)) diff --git a/src/adapter/src/catalog/open.rs b/src/adapter/src/catalog/open.rs index c946b5ef5c87a..57e99d313dda8 100644 --- a/src/adapter/src/catalog/open.rs +++ b/src/adapter/src/catalog/open.rs @@ -316,6 +316,17 @@ impl Catalog { .await; builtin_table_updates.extend(builtin_table_update); + // Both the `system_parameter_defaults` set above and the durable + // `SystemConfiguration` values applied via `pre_item_updates` live in + // the `SystemVars` value map by now. Mirror their effective values into + // the dyncfg `ConfigSet`, so that startup-only reads observe configured + // values rather than compile-time defaults, `ENABLE_EXPRESSION_CACHE` + // just below being one of them. `apply_updates` only + // syncs the `ConfigSet` when a `SystemConfiguration` update is present, + // and a deployment configured purely via `system_parameter_default` has + // none, so sync explicitly here. + state.system_config().sync_dyncfgs(); + // Ensure mz_system has a password if configured to have one. // It's important we do this after the `pre_item_updates` so that // the mz_system role exists in the catalog. diff --git a/src/adapter/src/client.rs b/src/adapter/src/client.rs index 705b6056195ed..b6c6298969c47 100644 --- a/src/adapter/src/client.rs +++ b/src/adapter/src/client.rs @@ -18,7 +18,7 @@ use std::time::{Duration, Instant}; use anyhow::bail; use chrono::{DateTime, Utc}; use derivative::Derivative; -use futures::{Stream, StreamExt}; +use futures::{FutureExt, Stream, StreamExt}; use itertools::Itertools; use mz_adapter_types::connection::{ConnectionId, ConnectionIdType}; use mz_auth::password::Password; @@ -60,7 +60,8 @@ use crate::command::{ use crate::config::{ScopedParameters, ScopedParametersScope, SystemParameterFrontend}; use crate::coord::{Coordinator, ExecuteContextGuard}; use crate::error::AdapterError; -use crate::metrics::{self, Metrics}; +use crate::metrics::Metrics; +use crate::peek_client::{ExecutionLogging, TakeOver}; use crate::session::{ EndTransactionAction, PreparedStatement, Session, SessionConfig, StateRevision, TransactionId, }; @@ -609,10 +610,15 @@ Issue a SQL query to get started. Need help? } #[instrument(level = "debug")] - pub(crate) fn send(&self, cmd: Command) { + pub(crate) fn try_send(&self, cmd: Command) -> bool { self.inner_cmd_tx .send((OpenTelemetryContext::obtain(), cmd)) - .expect("coordinator unexpectedly gone"); + .is_ok() + } + + #[instrument(level = "debug")] + pub(crate) fn send(&self, cmd: Command) { + assert!(self.try_send(cmd), "coordinator unexpectedly gone"); } } @@ -788,80 +794,94 @@ impl SessionClient { outer_ctx_extra: Option, ) -> Result<(ExecuteResponse, Instant), AdapterError> { let execute_started = Instant::now(); + let cancel_future = cancel_future.map(|_| ()).shared(); + + // Owning the end-of-execution obligation in this frame is what lets + // cancellation report an error: the inner future can be dropped without + // the obligation going with it. See `ExecutionLogging`. + let mut logging = ExecutionLogging::adopt(outer_ctx_extra, &self.peek_client); - let mut outer_ctx_extra = outer_ctx_extra; + let result = self + .execute_attempts(portal_name, &mut logging, cancel_future) + .await; + + logging.retire(&result); + + result.map(|response| (response, execute_started)) + } + /// Runs the execution paths in order of preference: frontend peek + /// sequencing, frontend read-then-write sequencing, and finally the + /// coordinator via `Command::Execute`. + async fn execute_attempts( + &mut self, + portal_name: String, + logging: &mut ExecutionLogging, + cancel_future: impl Future + Send + Clone, + ) -> Result { // Unroll SQL `EXECUTE (...)` so the inner statement - // flows through `try_frontend_peek` below, rather than being + // flows through `try_frontend_peek` / + // `try_frontend_read_then_write` below, rather than being // re-dispatched via `Command::Execute` from the coordinator's // `Plan::Execute` handler. Without this, a prepared statement // would route differently from the same statement issued // directly. // - // On a successful unroll, `unroll_sql_execute` also begins - // EXECUTE-level statement logging on the outer portal, so that + // On a successful unroll, `unroll_sql_execute` also takes over + // statement logging on the outer portal, so that // `mz_statement_execution_history` records `EXECUTE foo (...)` - // rather than the inner SQL, and installs the resulting - // `ExecuteContextGuard` into `outer_ctx_extra`. - let portal_name = self - .unroll_sql_execute(portal_name, &mut outer_ctx_extra) - .await?; + // rather than the inner SQL. + let portal_name = self.unroll_sql_execute(portal_name, logging).await?; // Attempt peek sequencing in the session task. // If unsupported, fall back to the Coordinator path. // TODO(peek-seq): wire up cancel_future - let peek_result = self - .try_frontend_peek(&portal_name, &mut outer_ctx_extra) - .await?; + let peek_result = self.try_frontend_peek(&portal_name, logging).await?; if let Some(resp) = peek_result { debug!("frontend peek succeeded"); - // Frontend peek handled the execution and retired outer_ctx_extra if it existed. - // No additional work needed here. - return Ok((resp, execute_started)); - } else { - debug!("frontend peek did not happen, falling back to `Command::Execute`"); - // If we bailed out, outer_ctx_extra is still present (if it was originally). - // `Command::Execute` will handle it. - // (This is not true if we bailed out _after_ the frontend peek sequencing has already - // begun its own statement logging. That case would be a bug.) + return Ok(resp); } - - let response = self - .send_with_cancel( - |tx, session| Command::Execute { - portal_name, - session, - tx, - outer_ctx_extra, - }, - cancel_future, - ) - .await?; - Ok((response, execute_started)) + debug!("frontend peek did not happen, falling back to `Command::Execute`"); + + // No frontend path took the statement over, so the coordinator retires + // whatever entry we hold, or begins its own if we hold none. + let outer_ctx_extra = logging.release(); + self.send_with_cancel( + |tx, session| Command::Execute { + portal_name, + session, + tx, + outer_ctx_extra, + }, + cancel_future, + ) + .await } /// If the named portal binds a SQL `EXECUTE `, resolve the /// prepared statement, install a fresh portal for the inner statement /// (carrying the EXECUTE's actual parameter values), and return that - /// portal's name so the caller can run `try_frontend_peek` against it. + /// portal's name so the caller can run `try_frontend_peek` / + /// `try_frontend_read_then_write` against it. /// /// Only ever unrolls one level: the parser rejects /// `PREPARE foo AS EXECUTE bar` (matching Postgres), so the inner /// statement is guaranteed not to be another `EXECUTE`. A failsafe below /// surfaces an internal error if that invariant is ever violated. /// - /// When the portal does not bind an `EXECUTE` — the common case — - /// returns the original portal name, costing only a portal lookup. + /// When the portal does not bind an `EXECUTE`, the common case, returns the + /// original portal name, costing only a portal lookup. async fn unroll_sql_execute( &mut self, portal_name: String, - outer_ctx_extra: &mut Option, + logging: &mut ExecutionLogging, ) -> Result { let (stmt, params, outer_logging, outer_lifecycle_timestamps) = { let session = self.session.as_ref().expect("SessionClient invariant"); let portal = match session.get_portal_unverified(&portal_name) { Some(p) => p, - // No portal: let `try_frontend_peek` surface the + // No portal: let `try_frontend_peek` / + // `try_frontend_read_then_write` surface the // standard "missing portal" error. None => return Ok(portal_name), }; @@ -892,107 +912,31 @@ impl SessionClient { Coordinator::verify_portal(&catalog, session, &portal_name)?; } - // Bump query_total for the outer EXECUTE itself. The inner - // statement gets its own increment inside `try_frontend_peek_inner` - // (or, on bailout, in the coordinator's `handle_execute`). - { - let session = self.session.as_ref().expect("SessionClient invariant"); - session - .metrics() - .query_total(&[ - metrics::session_type_label_value(session.user()), - metrics::statement_type_label_value(&stmt), - ]) - .inc(); - } - - // Begin EXECUTE-level statement logging up front, so that planning - // errors below produce an `Errored` end-event in - // `mz_statement_execution_history` rather than no entry at all. - // - // We pass the *outer* portal's `logging` and pgwire-bound `params` - // so the recorded entry shows the user-visible `EXECUTE foo (...)`, - // not the inner SQL. The id (if any) moves into `outer_ctx_extra` - // below for `try_frontend_peek` to retire; on planning error we - // explicitly emit an `Errored` end-event below. - let began_outer_logging = outer_ctx_extra.is_none(); - let logging_id: Option = - if began_outer_logging { - let session = self.session.as_mut().expect("SessionClient invariant"); - let result = self - .peek_client - .statement_logging_frontend - .begin_statement_execution( - session, - ¶ms, - &outer_logging, - catalog.system_config(), - outer_lifecycle_timestamps, - ); - if let Some((id, began_execution, mseh_update, prepared_statement)) = result { - self.peek_client.log_began_execution( - began_execution, - mseh_update, - prepared_statement, - ); - Some(id) - } else { - None - } - } else { - None - }; - - let new_portal_name = match self.install_inner_portal_for_execute(&catalog, &stmt, ¶ms) - { - Ok(name) => name, - Err(err) => { - if let Some(id) = logging_id { - self.peek_client.log_ended_execution( - id, - StatementEndedExecutionReason::Errored { - error: err.to_string(), - }, - ); - } - return Err(err); - } - }; - - // Hand off to `outer_ctx_extra` whenever we entered the begin path - // for the outer EXECUTE — even if `begin_statement_execution` - // returned `None` (sampling decided not to sample, or logging is - // disabled for the user). This mirrors the original coord path, - // which always installs a guard via - // `ExecuteContextGuard::new(maybe_uuid, ...)`. Without this, the - // inner portal would be treated as a fresh statement by - // `try_frontend_peek` (or the fallback `Command::Execute` path) - // and re-account its bytes against - // `mz_statement_logging_unsampled_bytes`, double-counting the - // inner SQL. - if began_outer_logging { - // Soft invariant: `try_frontend_peek` takes ownership of - // `outer_ctx_extra` immediately, so this guard's `Drop` is - // unreachable on the normal flow and the dummy channel is - // never used. If a panic does fire `Drop` between here and - // that takeover, the `Aborted` end-event is silently lost - // — an acceptable trade given the panic implies the - // connection is going down anyway. - let (dummy_tx, _dummy_rx) = mpsc::unbounded_channel(); - *outer_ctx_extra = Some(ExecuteContextGuard::new(logging_id, dummy_tx)); - } + // Take over the EXECUTE itself, so that a planning error below produces + // an errored end-event in `mz_statement_execution_history` rather than + // no entry at all. We pass the *outer* portal's `logging` and + // pgwire-bound `params` so the recorded entry shows the user-visible + // `EXECUTE foo (...)`, not the inner SQL. The inner statement gets its + // own `query_total` increment when a frontend path takes it over, or, + // on bailout, in the coordinator's `handle_execute`. + logging.take_over( + &self.peek_client, + self.session.as_mut().expect("SessionClient invariant"), + Some(&stmt), + ¶ms, + &outer_logging, + &catalog, + outer_lifecycle_timestamps, + TakeOver::UnrolledExecute, + ); - Ok(new_portal_name) + self.install_inner_portal_for_execute(&catalog, &stmt, ¶ms) } /// Helper for [`Self::unroll_sql_execute`]: plans the outer /// `Statement::Execute`, verifies the referenced prepared statement, and /// installs a fresh portal carrying the inner statement plus the /// EXECUTE's bound parameter values. Returns the new portal's name. - /// - /// Split out so [`Self::unroll_sql_execute`] can wrap the fallible work - /// in a single error-handling site that emits an `Errored` end-event - /// for the EXECUTE-level statement-logging entry. fn install_inner_portal_for_execute( &mut self, catalog: &Arc, @@ -1028,11 +972,11 @@ impl SessionClient { } }; - // Verify and install the inner portal. Mirrors - // `Coordinator::sequence_execute`. The new portal carries the inner - // prepared statement's `logging`, but `try_frontend_peek` will see - // `outer_ctx_extra=Some(...)` and inherit the EXECUTE-level logging - // instead of starting fresh from this portal. + // Verify and install the inner portal. The new portal carries the + // inner prepared statement's `logging`, but the logging slot already + // holds the EXECUTE-level entry, so a frontend path taking the inner + // statement over inherits that entry instead of starting fresh from + // this portal. let session = self.session.as_mut().expect("SessionClient invariant"); Coordinator::verify_prepared_statement(catalog, session, &execute_plan.name)?; let ps = session @@ -1046,7 +990,7 @@ impl SessionClient { // Failsafe: `PREPARE foo AS EXECUTE bar` is rejected by the parser, // so the resolved inner statement must not be another `EXECUTE`. If // that ever changes, we'd silently skip frontend sequencing for the - // deeper EXECUTEs — surface it as an internal error instead. + // deeper EXECUTEs. Surface it as an internal error instead. if let Some(inner) = inner_stmt.as_ref() { if matches!(inner, Statement::Execute(_)) { return Err(AdapterError::Internal(format!( @@ -1325,7 +1269,7 @@ impl SessionClient { async fn send_with_cancel( &mut self, f: F, - cancel_future: impl Future + Send, + cancel_future: impl Future + Send, ) -> Result where F: FnOnce(oneshot::Sender>, Session) -> Command, @@ -1430,7 +1374,7 @@ impl SessionClient { *client_session = Some(res.session); return res.result; }, - _err = &mut cancel_future, if !cancelled => { + _ = &mut cancel_future, if !cancelled => { cancelled = true; inner_client.send(Command::PrivilegedCancelRequest { conn_id: conn_id.clone(), @@ -1476,18 +1420,15 @@ impl SessionClient { /// /// Returns `Ok(Some(response))` if we handled the peek, or `Ok(None)` to fall back to the /// Coordinator's sequencing. If it returns an error, it should be returned to the user. - /// - /// `outer_ctx_extra` is Some when we are executing as part of an outer statement, e.g., a FETCH - /// triggering the execution of the underlying query. pub(crate) async fn try_frontend_peek( &mut self, portal_name: &str, - outer_ctx_extra: &mut Option, + logging: &mut ExecutionLogging, ) -> Result, AdapterError> { if self.enable_frontend_peek_sequencing { let session = self.session.as_mut().expect("SessionClient invariant"); self.peek_client - .try_frontend_peek(portal_name, session, outer_ctx_extra) + .try_frontend_peek(portal_name, session, logging) .await } else { Ok(None) diff --git a/src/adapter/src/command.rs b/src/adapter/src/command.rs index 3174ff359e99f..1665e1035491c 100644 --- a/src/adapter/src/command.rs +++ b/src/adapter/src/command.rs @@ -130,7 +130,11 @@ pub enum Command { portal_name: String, session: Session, tx: oneshot::Sender>, - outer_ctx_extra: Option, + /// The end-of-execution obligation of the statement this execution + /// serves, if any. `Coordinator::handle_execute` arms it again on + /// receipt. `None` means no log entry exists yet, so the coordinator + /// begins one. + outer_ctx_extra: Option, }, /// Attempts to commit or abort the session's transaction. Guarantees that the Coordinator's diff --git a/src/adapter/src/config/frontend.rs b/src/adapter/src/config/frontend.rs index a65a637a9f179..53d42f08ca482 100644 --- a/src/adapter/src/config/frontend.rs +++ b/src/adapter/src/config/frontend.rs @@ -471,7 +471,7 @@ async fn ld_client( let mut backoff = Duration::from_secs(5); let timeout = Duration::from_secs(10); - // TODO(materialize#32030): fix retry logic + // TODO(materialize#38055): fix retry logic loop { match ld_client.wait_for_initialization(timeout).await { Some(true) => break, diff --git a/src/adapter/src/coord.rs b/src/adapter/src/coord.rs index 0796e4b4b7b56..6356d9d3ed0b9 100644 --- a/src/adapter/src/coord.rs +++ b/src/adapter/src/coord.rs @@ -4279,7 +4279,8 @@ impl Coordinator { ctx_extra: ExecuteContextExtra, ) { if let Some(uuid) = ctx_extra.retire() { - self.end_statement_execution(uuid, reason); + let ended_at = self.now(); + self.end_statement_execution(uuid, reason, ended_at); } } diff --git a/src/adapter/src/coord/command_handler.rs b/src/adapter/src/coord/command_handler.rs index 41bbde3c9b7bf..45f582cf42cee 100644 --- a/src/adapter/src/coord/command_handler.rs +++ b/src/adapter/src/coord/command_handler.rs @@ -87,7 +87,7 @@ use crate::webhook::{ }; use crate::{AppendWebhookError, ExecuteContext, catalog, metrics}; -use super::ExecuteContextGuard; +use super::{ExecuteContextExtra, ExecuteContextGuard}; /// The login status of a role, used by authentication handlers to check role /// existence and login permission before proceeding to credential verification. @@ -1107,8 +1107,15 @@ impl Coordinator { // then `outer_context` should be `Some`. // This instructs the coordinator that the // outer execute should be considered finished once the inner one is. - outer_context: Option, + outer_context: Option, ) { + // Arm the obligation before anything can return: from here on, every + // exit either retires it or hands it to a guard whose `Drop` does. + // Nothing could have dropped it on the way here either, the command + // travels a channel that only this loop drains. + let outer_context = outer_context + .map(|extra| ExecuteContextGuard::new(extra.retire(), self.internal_cmd_tx.clone())); + // A new statement is starting, so discard any cancellation that was signaled while no // statement was running. Such a cancellation targeted an earlier statement and must not // cancel the new one. (Like in PostgreSQL, a cancel request that arrives when nothing is diff --git a/src/adapter/src/coord/sequencer.rs b/src/adapter/src/coord/sequencer.rs index 1c1f19303d774..c2c22fc996c83 100644 --- a/src/adapter/src/coord/sequencer.rs +++ b/src/adapter/src/coord/sequencer.rs @@ -624,13 +624,19 @@ impl Coordinator { match self.sequence_execute(ctx.session_mut(), plan) { Ok(portal_name) => { let (tx, _, session, extra, response_barriers) = ctx.into_parts(); + // The obligation travels as data and + // `handle_execute` arms it again. It cannot be + // dropped in between: the command goes to the + // channel this loop drains, and a spawned barrier + // task is only dropped when the process is going + // down, at which point nothing records anything. let command = Message::Command( OpenTelemetryContext::obtain(), Command::Execute { portal_name, session, tx: tx.take(), - outer_ctx_extra: Some(extra), + outer_ctx_extra: Some(extra.defuse()), }, ); if response_barriers.is_empty() { diff --git a/src/adapter/src/coord/sequencer/inner/copy_from.rs b/src/adapter/src/coord/sequencer/inner/copy_from.rs index e89212807a4cc..0c7c5a8d3e892 100644 --- a/src/adapter/src/coord/sequencer/inner/copy_from.rs +++ b/src/adapter/src/coord/sequencer/inner/copy_from.rs @@ -362,11 +362,9 @@ impl Coordinator { session: &session_meta, catalog_state, }; - let mut optimizer = optimize::view::Optimizer::new_with_prep_no_limit( - optimizer_config.clone(), - None, - prep, - ); + let mut optimizer = + optimize::view::Optimizer::new_with_prep(optimizer_config.clone(), None, prep) + .without_fold_constants_limit(); let hir = mz_sql::plan::plan_copy_from( &pcx, diff --git a/src/adapter/src/coord/statement_logging.rs b/src/adapter/src/coord/statement_logging.rs index 3cf13bbfaa175..0b1c35993237c 100644 --- a/src/adapter/src/coord/statement_logging.rs +++ b/src/adapter/src/coord/statement_logging.rs @@ -170,6 +170,7 @@ impl Coordinator { self.end_statement_execution( StatementLoggingId(ended_record.id), ended_record.reason, + ended_record.ended_at, ); } FrontendStatementLoggingEvent::SetCluster { @@ -394,17 +395,22 @@ impl Coordinator { /// Record the end of statement execution for a statement whose beginning /// was logged. Ends are idempotent: the first end wins and later ends for /// the same statement are ignored. + /// + /// `ended_at` is when execution finished, which is not the same as when this + /// runs. A statement sequenced off the coordinator loop finishes in its + /// session task and reports the end as a message, so taking the clock here + /// would charge it for however long that message sat in the queue. pub(crate) fn end_statement_execution( &mut self, id: StatementLoggingId, reason: StatementEndedExecutionReason, + ended_at: EpochMillis, ) { let StatementLoggingId(uuid) = id; - let now = self.now(); let ended_record = StatementEndedExecutionRecord { id: uuid, reason, - ended_at: now, + ended_at, }; let Some(began_record) = self.statement_logging.executions_begun.remove(&uuid) else { @@ -439,7 +445,7 @@ impl Coordinator { self.record_statement_lifecycle_event( &id, &StatementLifecycleEvent::ExecutionFinished, - now, + ended_at, ); } diff --git a/src/adapter/src/frontend_peek.rs b/src/adapter/src/frontend_peek.rs index 016365b994cc6..fc05dab424522 100644 --- a/src/adapter/src/frontend_peek.rs +++ b/src/adapter/src/frontend_peek.rs @@ -45,27 +45,24 @@ use tracing_opentelemetry::OpenTelemetrySpanExt; use crate::catalog::Catalog; use crate::command::Command; +use crate::coord; use crate::coord::peek::{FastPathPlan, PeekPlan}; use crate::coord::sequencer::{eval_copy_to_uri, statistics_oracle}; use crate::coord::timeline::timedomain_for; use crate::coord::timestamp_selection::TimestampDetermination; -use crate::coord::{ - Coordinator, CopyToContext, ExecuteContextGuard, ExplainContext, ExplainPlanContext, - TargetCluster, -}; +use crate::coord::{Coordinator, CopyToContext, ExplainContext, ExplainPlanContext, TargetCluster}; use crate::explain::insights::PlanInsightsContext; use crate::explain::optimizer_trace::OptimizerTrace; use crate::optimize::Optimize; use crate::optimize::dataflows::{ComputeInstanceSnapshot, DataflowBuilder}; -use crate::peek_client::StatementLoggingGuard; +use crate::peek_client::{ExecutionLogging, TakeOver}; use crate::session::{Session, TransactionOps, TransactionStatus}; +use crate::statement_logging::StatementLifecycleEvent; use crate::statement_logging::WatchSetCreation; -use crate::statement_logging::{StatementEndedExecutionReason, StatementLifecycleEvent}; use crate::{ AdapterError, AdapterNotice, CollectionIdBundle, ExecuteResponse, PeekClient, ReadHolds, TimelineContext, TimestampContext, TimestampProvider, optimize, }; -use crate::{coord, metrics}; impl PeekClient { /// Attempt to sequence a peek from the session task. @@ -73,13 +70,14 @@ impl PeekClient { /// Returns `Ok(Some(response))` if we handled the peek, or `Ok(None)` to fall back to the /// Coordinator's sequencing. If it returns an error, it should be returned to the user. /// - /// `outer_ctx_extra` is Some when we are executing as part of an outer statement, e.g., a FETCH - /// triggering the execution of the underlying query. + /// `logging` holds the end-of-execution obligation for this statement. The + /// caller retires it, this function only takes the statement over and, at + /// the dispatch sites that hand execution off, defuses the slot. pub(crate) async fn try_frontend_peek( &mut self, portal_name: &str, session: &mut Session, - outer_ctx_extra: &mut Option, + logging: &mut ExecutionLogging, ) -> Result, AdapterError> { // # From handle_execute @@ -98,26 +96,12 @@ impl PeekClient { let catalog = self.catalog_snapshot("try_frontend_peek").await; - // Extract things from the portal. - let (stmt, params, logging, lifecycle_timestamps) = { - if let Err(err) = Coordinator::verify_portal(&*catalog, session, portal_name) { - // An inherited outer statement (e.g. EXECUTE) ends here and we - // own its end, so we log it. If we discarded the id instead, - // the statement would stay "running" forever in - // mz_statement_execution_history. - if let Some(id) = outer_ctx_extra - .take() - .and_then(|guard| guard.defuse().retire()) - { - self.log_ended_execution( - id, - StatementEndedExecutionReason::Errored { - error: err.to_string(), - }, - ); - } - return Err(err); - } + // Extract things from the portal. A failed verification does not begin + // an entry, mirroring the coordinator: the portal is what statement + // logging draws its record from. An adopted outer entry still gets its + // `Errored` end from the retirement site in `SessionClient::execute`. + let (stmt, params, logging_info, lifecycle_timestamps) = { + Coordinator::verify_portal(&*catalog, session, portal_name)?; let portal = session .get_portal_unverified(portal_name) // The portal is a session-level thing, so it couldn't have concurrently disappeared @@ -125,9 +109,9 @@ impl PeekClient { .expect("called verify_portal above"); let params = portal.parameters.clone(); let stmt = portal.stmt.clone(); - let logging = Arc::clone(&portal.logging); + let logging_info = Arc::clone(&portal.logging); let lifecycle_timestamps = portal.lifecycle_timestamps.clone(); - (stmt, params, logging, lifecycle_timestamps) + (stmt, params, logging_info, lifecycle_timestamps) }; // Before planning, check if this is a statement type we can handle. @@ -201,70 +185,39 @@ impl PeekClient { } } - // Set up statement logging, and log the beginning of execution. - // (But only if we're not executing in the context of another statement.) - let mut logging_guard = self.begin_statement_logging( + // Past this point the coordinator never sees this statement, so every + // exit has to produce an outcome for it. The `Ok(None)` bailouts are all + // above. + logging.take_over( + self, session, + stmt.as_deref(), ¶ms, - &logging, + &logging_info, &catalog, lifecycle_timestamps, - outer_ctx_extra, + TakeOver::StatementToRun, ); - let result = self - .try_frontend_peek_inner(session, catalog, stmt, params, &mut logging_guard) - .await; - - // If we still own end-of-execution logging, retire it with the - // execution's outcome. We don't own it when a dispatch site in - // `try_frontend_peek_inner` handed it off: for streaming responses the - // end is logged asynchronously, by the coordinator for registered - // peeks and by the protocol layer for subscribes. - if logging_guard.id().is_some() { - let reason = match &result { - // Bailout case, which should not happen. - Ok(None) => { - soft_panic_or_log!( - "Bailed out from `try_frontend_peek_inner` after we already logged the beginning of statement execution." - ); - // The old peek sequencing would start its own statement - // logging from scratch; close out this one as errored. - StatementEndedExecutionReason::Errored { - error: "Internal error: bailed out from `try_frontend_peek_inner`" - .to_string(), - } - } - // Streaming responses cannot reach this arm: their dispatch - // sites hand off the guard. The `From` impl panics on them. - Ok(Some(resp)) => resp.into(), - Err(e) => StatementEndedExecutionReason::Errored { - error: e.to_string(), - }, - }; - - logging_guard.retire(reason); - } - - result + self.try_frontend_peek_inner(session, catalog, stmt, params, logging) + .await } /// This is encapsulated in an inner function so that the outer function can still do statement /// logging after the `?` returns of the inner function. /// - /// `logging_guard` owns end-of-execution logging for this statement. /// Dispatch sites that hand the statement to the coordinator for /// asynchronous completion (registered peeks, subscribes) `defuse` the - /// guard at the point where the coordinator takes over. Everywhere else - /// the guard stays armed and the caller logs the end from the returned - /// result. + /// logging slot at the point where the coordinator takes over. Everywhere + /// else the slot stays armed and `SessionClient::execute` logs the end from + /// the returned result. async fn try_frontend_peek_inner( &mut self, session: &mut Session, catalog: Arc, stmt: Option>>, params: Params, - logging_guard: &mut StatementLoggingGuard, + logging: &mut ExecutionLogging, ) -> Result, AdapterError> { let stmt = match stmt { Some(stmt) => stmt, @@ -274,14 +227,6 @@ impl PeekClient { } }; - session - .metrics() - .query_total(&[ - metrics::session_type_label_value(session.user()), - metrics::statement_type_label_value(&stmt), - ]) - .inc(); - // # From handle_execute_inner let conn_catalog = catalog.for_session(session); @@ -377,15 +322,17 @@ impl PeekClient { } _ => { // This shouldn't happen because we already checked for this at the AST - // level before calling `try_frontend_peek_inner`. + // level before calling `try_frontend_peek_inner`. The logging takeover + // has already happened, so falling back to the coordinator would count + // the statement twice. Report the inconsistency instead. soft_panic_or_log!( "unexpected EXPLAIN FILTER PUSHDOWN plan kind in frontend peek sequencing: {:?}", explainee ); - debug!( - "Bailing out from try_frontend_peek_inner, because EXPLAIN FILTER PUSHDOWN is not for a SELECT query or is EXPLAIN BROKEN" - ); - return Ok(None); + return Err(AdapterError::Internal( + "unexpected EXPLAIN FILTER PUSHDOWN plan kind in frontend peek sequencing" + .into(), + )); } } } @@ -445,15 +392,16 @@ impl PeekClient { Plan::Subscribe(subscribe) => (QueryPlan::Subscribe(subscribe), ExplainContext::None), _ => { // This shouldn't happen because we already checked for this at the AST - // level before calling `try_frontend_peek_inner`. + // level before calling `try_frontend_peek_inner`. The logging takeover has + // already happened, so falling back to the coordinator would count the + // statement twice. Report the inconsistency instead. soft_panic_or_log!( "Unexpected plan kind in frontend peek sequencing: {:?}", plan ); - debug!( - "Bailing out from try_frontend_peek_inner, because the Plan is not a SELECT, side-effecting SELECT, EXPLAIN SELECT, EXPLAIN FILTER PUSHDOWN, or COPY TO S3" - ); - return Ok(None); + return Err(AdapterError::Internal( + "unexpected plan kind in frontend peek sequencing".into(), + )); } }; @@ -496,7 +444,7 @@ impl PeekClient { }; // Log cluster selection - if let Some(logging_id) = logging_guard.id() { + if let Some(logging_id) = logging.id() { self.log_set_cluster(logging_id, target_cluster_id, target_cluster_name.clone()); } @@ -1185,7 +1133,7 @@ impl PeekClient { }; // Log optimization finished - if let Some(logging_id) = logging_guard.id() { + if let Some(logging_id) = logging.id() { self.log_lifecycle_event(logging_id, StatementLifecycleEvent::OptimizationFinished); } @@ -1293,7 +1241,7 @@ impl PeekClient { // # Now back to peek_finish - let watch_set = logging_guard.id().map(|logging_id| { + let watch_set = logging.id().map(|logging_id| { WatchSetCreation::new( logging_id, catalog.state(), @@ -1314,7 +1262,7 @@ impl PeekClient { let response = match peek_plan { PeekPlan::FastPath(fast_path_plan) => { - if let Some(logging_id) = logging_guard.id() { + if let Some(logging_id) = logging.id() { // TODO(peek-seq): Actually, we should log it also for // FastPathPlan::Constant. The only reason we are not doing so at the // moment is to match the old peek sequencing, so that statement logging @@ -1362,12 +1310,12 @@ impl PeekClient { session.conn_id().clone(), source_ids, watch_set, - logging_guard, + logging, ) .await? } PeekPlan::SlowPath(dataflow_plan) => { - if let Some(logging_id) = logging_guard.id() { + if let Some(logging_id) = logging.id() { self.log_set_transient_index_id(logging_id, dataflow_plan.id); } @@ -1392,7 +1340,7 @@ impl PeekClient { // coordinator logs nothing (see // `implement_slow_path_peek`), so the guard stays armed // and the caller logs the error. - logging_guard.defuse(); + logging.defuse(); response } }; @@ -1454,7 +1402,7 @@ impl PeekClient { session_uuid: session.uuid(), read_holds, plan: subscribe_plan, - statement_logging_id: logging_guard.id(), + statement_logging_id: logging.id(), tx, }) .await?; @@ -1464,7 +1412,7 @@ impl PeekClient { // coordinator logs nothing (see the `ExecuteSubscribe` // handler), so the guard stays armed and the caller logs the // error. - logging_guard.defuse(); + logging.defuse(); Ok(Some(response)) } Execution::CopyToS3 { @@ -1511,7 +1459,7 @@ impl PeekClient { .await?; // Preflight succeeded, now execute the actual COPY TO dataflow - let watch_set = logging_guard.id().map(|logging_id| { + let watch_set = logging.id().map(|logging_id| { WatchSetCreation::new( logging_id, catalog.state(), diff --git a/src/adapter/src/metrics.rs b/src/adapter/src/metrics.rs index 30abf69c635e2..b5168a4916b98 100644 --- a/src/adapter/src/metrics.rs +++ b/src/adapter/src/metrics.rs @@ -308,6 +308,7 @@ impl Metrics { row_set_finishing_seconds: self.row_set_finishing_seconds(), session_startup_table_writes_seconds: self.session_startup_table_writes_seconds.clone(), query_total: self.query_total.clone(), + subscribe_outputs: self.subscribe_outputs.clone(), determine_timestamp: self.determine_timestamp.clone(), timestamp_difference_for_strict_serializable_ms: self .timestamp_difference_for_strict_serializable_ms @@ -329,6 +330,7 @@ pub struct SessionMetrics { row_set_finishing_seconds: Histogram, session_startup_table_writes_seconds: Histogram, query_total: IntCounterVec, + subscribe_outputs: IntCounterVec, determine_timestamp: IntCounterVec, timestamp_difference_for_strict_serializable_ms: HistogramVec, timestamp_difference_for_bounded_staleness_ms: HistogramVec, @@ -351,6 +353,10 @@ impl SessionMetrics { self.query_total.with_label_values(label_values) } + pub(crate) fn subscribe_outputs(&self, label_values: &[&str]) -> GenericCounter { + self.subscribe_outputs.with_label_values(label_values) + } + pub(crate) fn determine_timestamp(&self, label_values: &[&str]) -> GenericCounter { self.determine_timestamp.with_label_values(label_values) } diff --git a/src/adapter/src/optimize.rs b/src/adapter/src/optimize.rs index 1c842d06c0336..6138c1c6feeb0 100644 --- a/src/adapter/src/optimize.rs +++ b/src/adapter/src/optimize.rs @@ -90,9 +90,8 @@ use crate::TimestampContext; /// A type for a [`DataflowDescription`] backed by `Mir~` plans. Used internally /// by the optimizer implementations. type MirDataflowDescription = DataflowDescription; -/// A type for a [`DataflowDescription`] backed by `Lir~` plans. Used internally -/// by the optimizer implementations. -type LirDataflowDescription = DataflowDescription; +/// A type for a [`DataflowDescription`] backed by `Lir~` plans. +pub type LirDataflowDescription = DataflowDescription; // Core API // -------- diff --git a/src/adapter/src/optimize/subscribe.rs b/src/adapter/src/optimize/subscribe.rs index 96f7839a253c4..19a082a1610ed 100644 --- a/src/adapter/src/optimize/subscribe.rs +++ b/src/adapter/src/optimize/subscribe.rs @@ -17,8 +17,9 @@ use differential_dataflow::lattice::Lattice; use mz_compute_types::ComputeInstanceId; use mz_compute_types::plan::LirRelationExpr; use mz_compute_types::sinks::{ComputeSinkConnection, ComputeSinkDesc, SubscribeSinkConnection}; +use mz_expr::{ColumnOrder, MirRelationExpr}; use mz_ore::soft_assert_or_log; -use mz_repr::{GlobalId, Timestamp}; +use mz_repr::{GlobalId, RelationDesc, Timestamp}; use mz_sql::optimizer_metrics::OptimizerMetrics; use mz_sql::plan::{HirToMirConfig, SubscribeFrom, SubscribePlan}; use mz_transform::TransformCtx; @@ -114,6 +115,118 @@ impl Optimizer { pub fn sink_id(&self) -> GlobalId { self.sink_id } + + /// The single subscribe optimization pipeline. Every subscribe, whatever it + /// reads from, goes through here, so a prep or metainfo step added here + /// applies to all of them. + fn optimize_inner( + &mut self, + source: SubscribeSource, + output: Vec, + ) -> Result, OptimizerError> { + let time = Instant::now(); + + let mut df_builder = { + let compute = self.compute_instance.clone(); + DataflowBuilder::new(&*self.catalog, compute).with_config(&self.config) + }; + let mut df_desc = MirDataflowDescription::new(self.debug_name.clone()); + let mut df_meta = DataflowMetainfo::default(); + + let (from, from_desc) = match source { + SubscribeSource::Id(from_id) => { + let from_desc = self + .catalog + .get_entry(&from_id) + .relation_desc() + .expect("subscribes can only be run on items with descs") + .into_owned(); + + df_builder.import_into_dataflow(&from_id, &mut df_desc, &self.config.features)?; + + (from_id, from_desc) + } + SubscribeSource::Query { expr, from_desc } => { + // MIR ⇒ MIR optimization (local) + let mut transform_ctx = TransformCtx::local( + &self.config.features, + &self.typecheck_ctx, + &mut df_meta, + Some(&mut self.metrics), + Some(self.view_id), + ); + let expr = optimize_mir_local(expr, &mut transform_ctx)?; + + df_builder.import_view_into_dataflow( + &self.view_id, + &expr, + &mut df_desc, + &self.config.features, + )?; + + (self.view_id, from_desc) + } + }; + df_builder.maybe_reoptimize_imported_views(&mut df_desc, &self.config)?; + + // Make SinkDesc + let sink_description = ComputeSinkDesc { + from, + from_desc, + connection: ComputeSinkConnection::Subscribe(SubscribeSinkConnection { output }), + with_snapshot: self.with_snapshot, + up_to: self.up_to.map(Antichain::from_elem).unwrap_or_default(), + // No `FORCE NOT NULL` for subscribes + non_null_assertions: vec![], + // No `REFRESH` for subscribes + refresh_schedule: None, + }; + df_desc.export_sink(self.sink_id, sink_description); + + // Prepare expressions in the assembled dataflow. + let style = ExprPrepMaintained; + df_desc.visit_children( + |r| style.prep_relation_expr(r), + |s| style.prep_scalar_expr(s), + )?; + + // Construct TransformCtx for global optimization. + let mut transform_ctx = TransformCtx::global( + &df_builder, + &mz_transform::EmptyStatisticsOracle, // TODO: wire proper stats + &self.config.features, + &self.typecheck_ctx, + &mut df_meta, + Some(&mut self.metrics), + ); + // Run global optimization. + mz_transform::optimize_dataflow(&mut df_desc, &mut transform_ctx, false)?; + + if self.config.mode == OptimizeMode::Explain { + // Collect the list of indexes used by the dataflow at this point. + trace_plan!(at: "global", &df_meta.used_indexes(&df_desc)); + } + + self.duration += time.elapsed(); + + // Return the (sealed) plan at the end of this optimization step. + Ok(GlobalMirPlan { + df_desc, + df_meta, + phantom: PhantomData::, + }) + } +} + +/// What a subscribe reads from, with any HIR ⇒ MIR lowering already done. +enum SubscribeSource { + /// An existing collection, imported by id. + Id(GlobalId), + /// A query, imported into the dataflow as a view under `view_id`. + Query { + expr: MirRelationExpr, + from_desc: RelationDesc, + }, } /// The (sealed intermediate) result after: @@ -171,122 +284,32 @@ impl Optimize for Optimizer { type To = GlobalMirPlan; fn optimize(&mut self, plan: SubscribePlan) -> Result { - let output = plan.output; - let plan = plan.from; - let time = Instant::now(); + let output = plan.output.row_order().to_vec(); - let mut df_builder = { - let compute = self.compute_instance.clone(); - DataflowBuilder::new(&*self.catalog, compute).with_config(&self.config) - }; - let mut df_desc = MirDataflowDescription::new(self.debug_name.clone()); - let mut df_meta = DataflowMetainfo::default(); - - match plan { - SubscribeFrom::Id(from_id) => { - let from = self.catalog.get_entry(&from_id); - let from_desc = from - .relation_desc() - .expect("subscribes can only be run on items with descs") - .into_owned(); - - df_builder.import_into_dataflow(&from_id, &mut df_desc, &self.config.features)?; - df_builder.maybe_reoptimize_imported_views(&mut df_desc, &self.config)?; - - // Make SinkDesc - let sink_description = ComputeSinkDesc { - from: from_id, - from_desc, - connection: ComputeSinkConnection::Subscribe(SubscribeSinkConnection { - output: output.row_order().to_vec(), - }), - with_snapshot: self.with_snapshot, - up_to: self.up_to.map(Antichain::from_elem).unwrap_or_default(), - // No `FORCE NOT NULL` for subscribes - non_null_assertions: vec![], - // No `REFRESH` for subscribes - refresh_schedule: None, - }; - df_desc.export_sink(self.sink_id, sink_description); - } + match plan.from { + SubscribeFrom::Id(from_id) => self.optimize_inner(SubscribeSource::Id(from_id), output), SubscribeFrom::Query { expr, desc } => { // TODO: Change the `expr` type to be `HirRelationExpr` and run // HIR ⇒ MIR lowering and decorrelation here. This would allow // us implement something like `EXPLAIN RAW PLAN FOR SUBSCRIBE.` // // let typ = expr.top_level_typ(); - // let expr = expr.lower(&self.config)?; - - // MIR ⇒ MIR optimization (local) - let mut transform_ctx = TransformCtx::local( - &self.config.features, - &self.typecheck_ctx, - &mut df_meta, - Some(&mut self.metrics), - Some(self.view_id), - ); + // Lowering counts towards the reported optimization time, so we + // time it here and `optimize_inner` times the rest. + let time = Instant::now(); let expr = expr.lower(HirToMirConfig::from(&self.config), None)?; - let expr = optimize_mir_local(expr, &mut transform_ctx)?; - - df_builder.import_view_into_dataflow( - &self.view_id, - &expr, - &mut df_desc, - &self.config.features, - )?; - df_builder.maybe_reoptimize_imported_views(&mut df_desc, &self.config)?; - - // Make SinkDesc - let sink_description = ComputeSinkDesc { - from: self.view_id, - from_desc: desc.clone(), - connection: ComputeSinkConnection::Subscribe(SubscribeSinkConnection { - output: output.row_order().to_vec(), - }), - with_snapshot: self.with_snapshot, - up_to: self.up_to.map(Antichain::from_elem).unwrap_or_default(), - // No `FORCE NOT NULL` for subscribes - non_null_assertions: vec![], - // No `REFRESH` for subscribes - refresh_schedule: None, - }; - df_desc.export_sink(self.sink_id, sink_description); + self.duration += time.elapsed(); + + self.optimize_inner( + SubscribeSource::Query { + expr, + from_desc: desc, + }, + output, + ) } - }; - - // Prepare expressions in the assembled dataflow. - let style = ExprPrepMaintained; - df_desc.visit_children( - |r| style.prep_relation_expr(r), - |s| style.prep_scalar_expr(s), - )?; - - // Construct TransformCtx for global optimization. - let mut transform_ctx = TransformCtx::global( - &df_builder, - &mz_transform::EmptyStatisticsOracle, // TODO: wire proper stats - &self.config.features, - &self.typecheck_ctx, - &mut df_meta, - Some(&mut self.metrics), - ); - // Run global optimization. - mz_transform::optimize_dataflow(&mut df_desc, &mut transform_ctx, false)?; - - if self.config.mode == OptimizeMode::Explain { - // Collect the list of indexes used by the dataflow at this point. - trace_plan!(at: "global", &df_meta.used_indexes(&df_desc)); } - - self.duration += time.elapsed(); - - // Return the (sealed) plan at the end of this optimization step. - Ok(GlobalMirPlan { - df_desc, - df_meta, - phantom: PhantomData::, - }) } } diff --git a/src/adapter/src/optimize/view.rs b/src/adapter/src/optimize/view.rs index 96476cfc44b83..16dcecb8bcc13 100644 --- a/src/adapter/src/optimize/view.rs +++ b/src/adapter/src/optimize/view.rs @@ -53,23 +53,19 @@ pub struct Optimizer { impl Optimizer { /// Creates an optimizer instance that does not perform any expression - /// preparation. Additionally, this instance calls constant folding with a size limit. + /// preparation. pub fn new(config: OptimizerConfig, metrics: Option) -> Self { - Self { - typecheck_ctx: empty_typechecking_context(), - config, - metrics, - expr_prep_style: ExprPrepNoop, - fold_constants_limit: true, - } + Self::new_with_prep(config, metrics, ExprPrepNoop) } } impl Optimizer { /// Creates an optimizer instance that takes an [`ExprPrep`] to handle - /// unmaterializable functions. Additionally, this instance calls constant - /// folding without a size limit. - pub fn new_with_prep_no_limit( + /// unmaterializable functions. + /// + /// Constant folding runs with the usual size limit, see + /// [`Self::without_fold_constants_limit`]. + pub fn new_with_prep( config: OptimizerConfig, metrics: Option, expr_prep_style: S, @@ -79,9 +75,16 @@ impl Optimizer { config, metrics, expr_prep_style, - fold_constants_limit: false, + fold_constants_limit: true, } } + + /// Folds constants of any size, instead of giving up above the configured + /// limit. + pub fn without_fold_constants_limit(mut self) -> Self { + self.fold_constants_limit = false; + self + } } impl Optimize for Optimizer { diff --git a/src/adapter/src/peek_client.rs b/src/adapter/src/peek_client.rs index d489ca9a4eb09..62321a9c5d469 100644 --- a/src/adapter/src/peek_client.rs +++ b/src/adapter/src/peek_client.rs @@ -18,13 +18,17 @@ use mz_compute_client::protocol::command::PeekTarget; use mz_compute_types::ComputeInstanceId; use mz_expr::row::RowCollection; use mz_ore::cast::CastFrom; +use mz_ore::soft_panic_or_log; use mz_persist_client::PersistClient; use mz_repr::GlobalId; use mz_repr::Timestamp; use mz_repr::global_id::TransientIdGen; use mz_repr::{RelationDesc, Row}; +use mz_sql::ast::{Raw, Statement}; use mz_sql::optimizer_metrics::OptimizerMetrics; use mz_sql::plan::Params; +use mz_sql::session::metadata::SessionMetadata; +use mz_sql_parser::ast::{CopyRelation, CopyStatement, SubscribeStatement}; use mz_storage_types::sources::Timeline; use mz_timestamp_oracle::TimestampOracle; use prometheus::Histogram; @@ -35,15 +39,15 @@ use tokio::sync::oneshot; use uuid::Uuid; use crate::catalog::Catalog; -use crate::command::{CatalogSnapshot, Command}; +use crate::command::{CatalogSnapshot, Command, ExecuteResponse}; use crate::coord::peek::FastPathPlan; -use crate::coord::{Coordinator, ExecuteContextGuard}; +use crate::coord::{Coordinator, ExecuteContextExtra, ExecuteContextGuard}; use crate::session::{LifecycleTimestamps, Session}; use crate::statement_logging::{ FrontendStatementLoggingEvent, PreparedStatementEvent, PreparedStatementLoggingInfo, StatementLoggingFrontend, StatementLoggingId, WatchSetCreation, }; -use crate::{AdapterError, Client, CollectionIdBundle, ReadHolds, statement_logging}; +use crate::{AdapterError, Client, CollectionIdBundle, ReadHolds, metrics, statement_logging}; /// Storage collections trait alias we need to consult for since/frontiers. pub type StorageCollectionsHandle = @@ -271,11 +275,10 @@ impl PeekClient { /// peek target. For slow-path peeks (to be implemented later), we'll need to additionally call /// into the Controller to acquire a hold on the peek target after we create the dataflow. /// - /// `logging_guard` owns end-of-execution logging for this statement. For a - /// constant peek it stays armed and the caller logs the end from the - /// returned result. For a `PeekExisting`/`PeekPersist` peek, successful - /// registration with the coordinator hands ownership of the end to the - /// coordinator and the guard is defused here. That holds even when the + /// For a constant peek the logging slot stays armed and the caller logs the + /// end from the returned result. For a `PeekExisting`/`PeekPersist` peek, + /// successful registration with the coordinator hands ownership of the end + /// to the coordinator and the slot is defused here. That holds even when the /// subsequent `client.peek()` fails to issue. pub(crate) async fn implement_fast_path_peek_plan( &mut self, @@ -294,7 +297,7 @@ impl PeekClient { conn_id: mz_adapter_types::connection::ConnectionId, depends_on: std::collections::BTreeSet, watch_set: Option, - logging_guard: &mut StatementLoggingGuard, + logging: &mut ExecutionLogging, ) -> Result { // If the dataflow optimizes to a constant expression, we can immediately return the result. if let FastPathPlan::Constant(rows_res, _) = fast_path { @@ -431,7 +434,7 @@ impl PeekClient { // cancellation, concurrent teardown (e.g. a DROP CLUSTER), or the // unregistration below. We defuse the guard so the frontend doesn't // also log the end. - logging_guard.defuse(); + logging.defuse(); // Test-only synchronization point: parks a peek between registration // and issue, so a test can land a concurrent DROP CLUSTER in this @@ -494,51 +497,32 @@ impl PeekClient { }) } - /// Set up statement logging for a frontend-sequenced operation. + /// Begins a new statement execution log entry, sampling permitting. /// - /// If `outer_ctx_extra` is `None`, begins a new statement execution log - /// entry. If `outer_ctx_extra` is `Some` (e.g. EXECUTE/FETCH), reuses and - /// retires the existing logging context. - /// - /// Returns a [`StatementLoggingGuard`]. Callers must either - /// [`retire`](StatementLoggingGuard::retire) the guard on the execution's - /// terminal outcome, or [`defuse`](StatementLoggingGuard::defuse) it at - /// the point where end-of-execution logging is handed off to another - /// component. Dropping the guard without retiring it emits an `Aborted` - /// end-execution event. - pub(crate) fn begin_statement_logging( + /// Only [`ExecutionLogging::take_over`] may call this: an entry that exists + /// without the session task owning its end would stay unfinished forever. + fn begin_statement_logging( &self, session: &mut Session, params: &Params, logging: &Arc>, catalog: &Catalog, lifecycle_timestamps: Option, - outer_ctx_extra: &mut Option, ) -> StatementLoggingGuard { - let id = if outer_ctx_extra.is_none() { - // This is a new statement, so begin statement logging. - let result = self.statement_logging_frontend.begin_statement_execution( - session, - params, - logging, - catalog.system_config(), - lifecycle_timestamps, - ); + let result = self.statement_logging_frontend.begin_statement_execution( + session, + params, + logging, + catalog.system_config(), + lifecycle_timestamps, + ); - if let Some((logging_id, began_execution, mseh_update, prepared_statement)) = result { + let id = result.map( + |(logging_id, began_execution, mseh_update, prepared_statement)| { self.log_began_execution(began_execution, mseh_update, prepared_statement); - Some(logging_id) - } else { - None - } - } else { - // We're executing in the context of another statement (e.g. FETCH), - // so take ownership of the outer context and inherit its logging id - // (if any). The end of execution will be logged by the caller. - outer_ctx_extra - .take() - .and_then(|guard| guard.defuse().retire()) - }; + logging_id + }, + ); StatementLoggingGuard { id, @@ -616,38 +600,16 @@ impl PeekClient { FrontendStatementLoggingEvent::Lifecycle { id, event, when }, )); } - - /// Emit a `FrontendStatementLoggingEvent::EndedExecution` for the given - /// logging id. Used by the few callers that hold a bare - /// [`StatementLoggingId`] rather than a [`StatementLoggingGuard`], e.g. - /// error paths of the EXECUTE unrolling where the id is not wrapped in a - /// guard. - pub(crate) fn log_ended_execution( - &self, - id: StatementLoggingId, - reason: statement_logging::StatementEndedExecutionReason, - ) { - let ended_at = (self.statement_logging_frontend.now)(); - let record = statement_logging::StatementEndedExecutionRecord { - id: id.0, - reason, - ended_at, - }; - self.coordinator_client - .send(Command::FrontendStatementLogging( - FrontendStatementLoggingEvent::EndedExecution(record), - )); - } } /// RAII guard owning a frontend statement-logging lifecycle. /// -/// Created by [`PeekClient::begin_statement_logging`]. Unless logging -/// responsibility is handed off via [`defuse`](StatementLoggingGuard::defuse), -/// the guard ensures that every statement for which `BeganExecution` was logged -/// also receives a corresponding `EndedExecution`, even on early-return, panic, -/// or mid-flight drop of the enclosing future: if the guard is dropped without -/// being defused, it emits `StatementEndedExecutionReason::Aborted`. +/// Unless logging responsibility is handed off via +/// [`defuse`](StatementLoggingGuard::defuse), the guard ensures that every +/// statement for which `BeganExecution` was logged also receives a +/// corresponding `EndedExecution`, even on early-return, panic, or mid-flight +/// drop of the enclosing future: if the guard is dropped without being defused, +/// it emits `StatementEndedExecutionReason::Aborted`. /// /// When the guard is `defuse`d, some other component (e.g. the coordinator, for /// streaming peek / subscribe responses) takes over and logs `EndedExecution` @@ -657,7 +619,7 @@ impl PeekClient { /// retirement / drop are no-ops. #[must_use = "StatementLoggingGuard must be explicitly retired or handed off; \ otherwise `Drop` will log the statement as Aborted"] -pub(crate) struct StatementLoggingGuard { +struct StatementLoggingGuard { /// `None` if the statement was not sampled for logging. id: Option, coordinator_client: Client, @@ -665,21 +627,37 @@ pub(crate) struct StatementLoggingGuard { } impl StatementLoggingGuard { + /// Arms a guard for the obligation the coordinator armed for `outer`, the + /// statement whose execution the one we are about to run serves. + fn adopt(outer: ExecuteContextGuard, peek_client: &PeekClient) -> Self { + Self { + id: outer.defuse().retire(), + coordinator_client: peek_client.coordinator_client.clone(), + now: peek_client.statement_logging_frontend.now.clone(), + } + } + /// Returns the logging id, if this statement is being logged. - pub(crate) fn id(&self) -> Option { + fn id(&self) -> Option { self.id } /// Retires the guard with an explicit end-execution reason. /// A no-op if the guard was defused or the statement is not sampled. - pub(crate) fn retire(mut self, reason: statement_logging::StatementEndedExecutionReason) { + fn retire(mut self, reason: statement_logging::StatementEndedExecutionReason) { self.emit(reason); } + /// Turns the obligation back into its transferable form, disarming this + /// guard. + fn release(mut self) -> ExecuteContextExtra { + ExecuteContextExtra::new(self.id.take()) + } + /// Hands off logging responsibility without emitting an end-execution /// event. Call this at the point where another component takes over /// end-of-execution logging. Afterwards the guard is inert. - pub(crate) fn defuse(&mut self) { + fn defuse(&mut self) { self.id = None; } @@ -693,8 +671,12 @@ impl StatementLoggingGuard { reason, ended_at, }; - self.coordinator_client - .send(Command::FrontendStatementLogging( + // A guard can outlive the coordinator during shutdown. Failing to send + // costs us one end event, panicking in `Drop` would cost the whole + // connection. + let _ = self + .coordinator_client + .try_send(Command::FrontendStatementLogging( FrontendStatementLoggingEvent::EndedExecution(record), )); } @@ -708,6 +690,217 @@ impl Drop for StatementLoggingGuard { } } +/// The session task's slot for the end-of-execution obligation of the statement +/// it is running. One slot is held for the whole `SessionClient::execute` call +/// and there is exactly one retirement site. +/// +/// An empty slot means no log entry exists for this execution, so a fallback to +/// the coordinator lets it begin its own. An occupied slot means the session +/// task owes an end event: [`Self::retire`] pays it, [`Self::release`] transfers +/// it to the coordinator. [`Self::id`] returning `None` covers both "not +/// sampled" and "the end is logged elsewhere", which want identical treatment +/// everywhere. +pub(crate) struct ExecutionLogging { + guard: Option, + /// Whether the coordinator must not run this statement, because the session + /// task has already counted it in the metrics `Coordinator::handle_execute` + /// maintains. Handing it over afterwards would count it twice. + coordinator_must_not_run: bool, +} + +/// Which statement the session task is taking the log entry over for. +pub(crate) enum TakeOver { + /// The statement that will run here, so the coordinator must not run it. + StatementToRun, + /// A SQL `EXECUTE` that unrolls into an inner statement, for a session task + /// that will go on to run that inner statement. + /// + /// The entry stays armed, so a failure to unroll is recorded against the + /// `EXECUTE`. Only the `EXECUTE` itself is counted here. The inner + /// statement is a statement in its own right and is counted wherever it + /// ends up running, exactly as the coordinator does when it re-dispatches a + /// `Plan::Execute`, which is why the slot stays releasable. + UnrolledExecute, +} + +impl ExecutionLogging { + /// Adopts the end-of-execution obligation of an outer statement (a FETCH or + /// an EXECUTE running its inner statement), or starts out empty when there + /// is no outer statement. + pub(crate) fn adopt(outer: Option, peek_client: &PeekClient) -> Self { + Self { + guard: outer.map(|outer| StatementLoggingGuard::adopt(outer, peek_client)), + coordinator_must_not_run: false, + } + } + + /// Returns the logging id, if an end event is owed and the statement is + /// being logged. + pub(crate) fn id(&self) -> Option { + self.guard.as_ref().and_then(|guard| guard.id()) + } + + /// Records that the session task, not the coordinator, is executing `stmt`. + /// Bumps the counters `Coordinator::handle_execute` would have bumped and + /// makes sure a log entry exists, inheriting the adopted one if there is + /// one. `stmt` is `None` for an empty portal, which is logged but not + /// counted. + /// + /// For [`TakeOver::StatementToRun`], every exit after this call must produce + /// an outcome for the statement: the coordinator will not see it. + pub(crate) fn take_over( + &mut self, + peek_client: &PeekClient, + session: &mut Session, + stmt: Option<&Statement>, + params: &Params, + logging: &Arc>, + catalog: &Catalog, + lifecycle_timestamps: Option, + taking_over: TakeOver, + ) -> Option { + self.begin_or_inherit( + peek_client, + session, + params, + logging, + catalog, + lifecycle_timestamps, + ); + count_statement(session, stmt); + if matches!(taking_over, TakeOver::StatementToRun) { + self.coordinator_must_not_run = true; + } + self.id() + } + + /// Makes sure a log entry exists for this execution, keeping an adopted one + /// if there is one. + fn begin_or_inherit( + &mut self, + peek_client: &PeekClient, + session: &mut Session, + params: &Params, + logging: &Arc>, + catalog: &Catalog, + lifecycle_timestamps: Option, + ) { + if self.guard.is_none() { + self.guard = Some(peek_client.begin_statement_logging( + session, + params, + logging, + catalog, + lifecycle_timestamps, + )); + } + } + + /// Hands the obligation to the coordinator, which retires it once the + /// statement it dispatches finishes. `None` tells the coordinator to begin + /// its own entry, `Some` with no id inside tells it that an entry already + /// exists or that sampling declined. + #[must_use] + pub(crate) fn release(&mut self) -> Option { + if self.coordinator_must_not_run { + soft_panic_or_log!( + "statement handed to the coordinator after the session task took it over: \ + its per-statement metrics are counted twice" + ); + } + self.guard.take().map(|guard| guard.release()) + } + + /// Emits the end event for `result`, if we still owe one. + pub(crate) fn retire(self, result: &Result) { + let Some(guard) = self.guard else { + return; + }; + // A defused or released slot owes no end event. Bail before mapping + // `result` to an end reason, which soft-panics for the responses whose + // end is logged elsewhere. + if guard.id().is_none() { + return; + } + guard.retire(end_reason(result)); + } + + /// Leaves the slot inert, for dispatch sites that hand the end of execution + /// to the coordinator or the protocol layer (registered peeks, subscribes). + pub(crate) fn defuse(&mut self) { + if let Some(guard) = self.guard.as_mut() { + guard.defuse(); + } + } +} + +/// Maps an execution outcome to the reason to record for it. +/// +/// The responses whose end is logged elsewhere are filtered out first: their +/// `StatementEndedExecutionReason` conversion panics, and this runs for every +/// statement the session task executes, so that panic would take down the +/// connection. +fn end_reason( + result: &Result, +) -> statement_logging::StatementEndedExecutionReason { + if let Ok(response) = result { + if terminates_elsewhere(response) { + soft_panic_or_log!( + "frontend-sequenced statement still owed an end event while returning {:?}", + crate::command::ExecuteResponseKind::from(response) + ); + return statement_logging::StatementEndedExecutionReason::Aborted; + } + } + result.into() +} + +/// Bumps the per-statement counters `Coordinator::handle_execute` maintains. +/// +/// `stmt` is `None` for an empty portal, which is logged but not counted. The +/// coordinator skips it the same way, and matching that is the point of this +/// function. +fn count_statement(session: &Session, stmt: Option<&Statement>) { + let Some(stmt) = stmt else { + return; + }; + let session_type = metrics::session_type_label_value(session.user()); + session + .metrics() + .query_total(&[session_type, metrics::statement_type_label_value(stmt)]) + .inc(); + if let Statement::Subscribe(SubscribeStatement { output, .. }) + | Statement::Copy(CopyStatement { + relation: CopyRelation::Subscribe(SubscribeStatement { output, .. }), + .. + }) = stmt + { + session + .metrics() + .subscribe_outputs(&[session_type, metrics::subscribe_output_label_value(output)]) + .inc(); + } +} + +/// Whether someone else logs the end of execution for `response`: the +/// coordinator for a registered peek, the protocol layer for a subscribe, a +/// FETCH or a COPY FROM. The dispatch sites that produce these defuse the slot, +/// so an armed slot alongside one of them means a dispatch site did not. +fn terminates_elsewhere(response: &ExecuteResponse) -> bool { + match response { + ExecuteResponse::SendingRowsStreaming { .. } + | ExecuteResponse::Subscribing { .. } + | ExecuteResponse::Fetch { .. } + | ExecuteResponse::CopyFrom { .. } => true, + // COPY TO STDOUT of an immediate result terminates here. Anything else + // it can wrap does not. + ExecuteResponse::CopyTo { resp, .. } => { + !matches!(**resp, ExecuteResponse::SendingRowsImmediate { .. }) + } + _ => false, + } +} + /// Errors arising during collection lookup in peek client operations. #[derive(Error, Debug)] pub enum CollectionLookupError { diff --git a/src/dyncfg-launchdarkly/src/lib.rs b/src/dyncfg-launchdarkly/src/lib.rs index 261a5cf695211..0894d305b210e 100644 --- a/src/dyncfg-launchdarkly/src/lib.rs +++ b/src/dyncfg-launchdarkly/src/lib.rs @@ -82,7 +82,7 @@ where let max_backoff = Duration::from_secs(60); let mut backoff = Duration::from_secs(5); - // TODO(materialize#32030): fix retry logic + // TODO(materialize#38055): fix retry logic loop { match client.wait_for_initialization(config_sync_timeout).await { Some(true) => break, diff --git a/src/environmentd/tests/pgwire.rs b/src/environmentd/tests/pgwire.rs index fd6fcbb9886d5..5ad74336c2408 100644 --- a/src/environmentd/tests/pgwire.rs +++ b/src/environmentd/tests/pgwire.rs @@ -16,7 +16,7 @@ use std::net::TcpStream; use std::path::Path; use std::time::Duration; -use bytes::BytesMut; +use bytes::{BufMut, BytesMut}; use fallible_iterator::FallibleIterator; use mz_adapter::session::DEFAULT_DATABASE_NAME; use mz_environmentd::test_util::{self, PostgresErrorExt}; @@ -1125,3 +1125,152 @@ fn test_many_columns() { } } } + +/// Writes a frontend message with type byte `typ`, filling in its length prefix. +/// +/// `postgres_protocol`'s builders cap repeated groups at `i16::MAX` entries, so +/// a message with more parameters than that must be assembled by hand. +fn put_frontend_message(buf: &mut BytesMut, typ: u8, body: F) { + buf.put_u8(typ); + let base = buf.len(); + buf.put_u32(0); + body(buf); + let len = u32::try_from(buf.len() - base).unwrap(); + buf[base..base + 4].copy_from_slice(&len.to_be_bytes()); +} + +/// Reads backend messages up to and including the `ReadyForQuery` that answers +/// `Sync`. +fn read_until_ready(stream: &mut TcpStream) -> Vec { + use postgres_protocol::message::backend::Message; + + let mut buf = BytesMut::new(); + let mut chunk = [0; 1 << 13]; + let mut msgs = Vec::new(); + loop { + if let Some(msg) = Message::parse(&mut buf).unwrap() { + let ready = matches!(msg, Message::ReadyForQuery(_)); + msgs.push(msg); + if ready { + return msgs; + } + } else { + let n = stream.read(&mut chunk).unwrap(); + assert!(n > 0, "connection closed after {} messages", msgs.len()); + buf.extend_from_slice(&chunk[..n]); + } + } +} + +// The counts that precede the repeated groups of Parse and Bind (parameter +// types, format codes, parameter values) are decoded as unsigned, like +// PostgreSQL does, so a client may send more than `i16::MAX` parameters. +// Reading such a count as signed made it negative, which dropped the group and +// left the decoder misaligned for the rest of the message. See SQL-491. +#[mz_ore::test] +#[allow(clippy::disallowed_methods)] +fn test_many_bind_params() { + use postgres_protocol::message::backend::Message; + use postgres_protocol::message::frontend; + + // One past `i16::MAX`. The wire format allows 65535, but a bind that large + // runs into the unrelated `MAX_REQUEST_SIZE` limit. + const PARAMS: usize = 32768; + const INT4_OID: u32 = 23; + // Bound to the last parameter, so the inserted row matches only if the + // decoder stayed aligned to the end of the parameter array. + const LAST_VALUE: i32 = 4242; + + let server = test_util::TestHarness::default().start_blocking(); + let mut client = server.connect(postgres::NoTls).unwrap(); + client + .batch_execute("CREATE TABLE many_params (a int)") + .unwrap(); + + let mut stream = TcpStream::connect(server.sql_local_addr()).unwrap(); + stream + .set_read_timeout(Some(Duration::from_secs(120))) + .unwrap(); + let mut buf = BytesMut::new(); + frontend::startup_message( + vec![ + ("user", "materialize"), + ("database", DEFAULT_DATABASE_NAME), + ("options", "--welcome_message=off"), + ], + &mut buf, + ) + .unwrap(); + stream.write_all(&buf).unwrap(); + read_until_ready(&mut stream); + + // The shape from the report: pgjdbc's reWriteBatchedInserts rewrites a batch + // into one insert whose placeholder count exceeds `i16::MAX`. + let placeholders = (1..=PARAMS) + .map(|i| format!("${i}")) + .collect::>() + .join(","); + let sql = format!("INSERT INTO many_params VALUES (coalesce({placeholders}))"); + + // The statement and portal are unnamed throughout, hence the bare `0`s. + buf.clear(); + put_frontend_message(&mut buf, b'P', |buf| { + buf.put_u8(0); + buf.put_slice(sql.as_bytes()); + buf.put_u8(0); + buf.put_u16(u16::try_from(PARAMS).unwrap()); + for _ in 0..PARAMS { + buf.put_u32(INT4_OID); + } + }); + put_frontend_message(&mut buf, b'D', |buf| { + buf.put_u8(b'S'); // the statement, to get a ParameterDescription back + buf.put_u8(0); + }); + put_frontend_message(&mut buf, b'B', |buf| { + buf.put_u8(0); + buf.put_u8(0); + buf.put_u16(u16::try_from(PARAMS).unwrap()); + for _ in 0..PARAMS { + buf.put_u16(1); // binary + } + buf.put_u16(u16::try_from(PARAMS).unwrap()); + for _ in 1..PARAMS { + buf.put_i32(-1); // NULL + } + buf.put_i32(4); + buf.put_i32(LAST_VALUE); + buf.put_u16(0); // results in text + }); + put_frontend_message(&mut buf, b'E', |buf| { + buf.put_u8(0); + buf.put_i32(0); // no row limit + }); + frontend::sync(&mut buf); + stream.write_all(&buf).unwrap(); + + let mut described_params = None; + let mut tags = Vec::new(); + for msg in read_until_ready(&mut stream) { + match msg { + Message::ErrorResponse(body) => { + let fields: Vec<_> = body + .fields() + .map(|f| Ok(String::from_utf8_lossy(f.value_bytes()).into_owned())) + .collect() + .unwrap(); + panic!("unexpected error response: {fields:?}"); + } + Message::ParameterDescription(body) => { + described_params = Some(body.parameters().count().unwrap()); + } + Message::CommandComplete(body) => tags.push(body.tag().unwrap().to_string()), + _ => (), + } + } + assert_eq!(described_params, Some(PARAMS)); + assert_eq!(tags, ["INSERT 0 1"]); + + let row = client.query_one("SELECT a FROM many_params", &[]).unwrap(); + assert_eq!(row.get::<_, i32>(0), LAST_VALUE); +} diff --git a/src/environmentd/tests/server.rs b/src/environmentd/tests/server.rs index 9ea6298a9d86b..7e3c0dd58f669 100644 --- a/src/environmentd/tests/server.rs +++ b/src/environmentd/tests/server.rs @@ -41,19 +41,15 @@ use mz_frontegg_auth::{ }; use mz_frontegg_mock::{FronteggMockServer, models::ApiToken, models::UserConfig}; use mz_ore::cast::CastFrom; -use mz_ore::cast::CastLossy; -use mz_ore::cast::TryCastFrom; -use mz_ore::collections::CollectionExt; use mz_ore::error::ErrorExt; use mz_ore::metrics::MetricsRegistry; -use mz_ore::now::{NowFn, SYSTEM_TIME, to_datetime}; +use mz_ore::now::{NowFn, SYSTEM_TIME}; use mz_ore::retry::Retry; use mz_ore::{assert_contains, task::RuntimeExt}; -use mz_ore::{assert_err, assert_none, assert_ok, task}; +use mz_ore::{assert_err, assert_ok, task}; use mz_pgrepr::UInt8; use mz_repr::UNKNOWN_COLUMN_NAME; use mz_sql::session::user::{ANALYTICS_USER, HTTP_DEFAULT_USER, SYSTEM_USER}; -use mz_sql_parser::ast::display::AstDisplay; use openssl::ssl::{SslConnectorBuilder, SslVerifyMode}; use openssl::x509::X509; use postgres::config::SslMode; @@ -172,978 +168,6 @@ fn test_persistence() { ); } -/// A wrapper around `TestServerWithRuntime` that runs statement logging checks when dropped. -/// -/// This guard ensures that all statements have finished executing (have non-NULL `finished_at` -/// and `finished_status` in `mz_internal.mz_recent_activity_log`) before the test completes. -struct TestServerWithStatementLoggingChecks { - server: test_util::TestServerWithRuntime, -} - -impl TestServerWithStatementLoggingChecks { - /// Connect to the __internal__ SQL port of the running `environmentd` server. - pub fn connect_internal(&self, tls: T) -> Result - where - T: postgres::tls::MakeTlsConnect + Send + 'static, - T::TlsConnect: Send, - T::Stream: Send, - >::Future: Send, - { - self.server.connect_internal(tls) - } - - /// Returns the metrics registry for the test server. - pub fn metrics_registry(&self) -> &MetricsRegistry { - self.server.metrics_registry() - } -} - -/// Helper to get statement logging record counts from the metrics registry. -/// Returns (sampled_true_count, sampled_false_count). -#[allow(clippy::disallowed_methods)] -fn get_statement_logging_record_counts( - server: &TestServerWithStatementLoggingChecks, -) -> (u64, u64) { - let metrics = server.metrics_registry().gather(); - let record_count_metric = metrics - .into_iter() - .find(|m| m.name() == "mz_statement_logging_record_count") - .expect("mz_statement_logging_record_count metric should exist"); - - let metric_entries = record_count_metric.get_metric(); - let sampled_true = metric_entries - .iter() - .find(|m| { - m.get_label() - .iter() - .any(|l| l.name() == "sample" && l.value() == "true") - }) - .map(|m| u64::cast_lossy(m.get_counter().value())) - .unwrap_or(0); - let sampled_false = metric_entries - .iter() - .find(|m| { - m.get_label() - .iter() - .any(|l| l.name() == "sample" && l.value() == "false") - }) - .map(|m| u64::cast_lossy(m.get_counter().value())) - .unwrap_or(0); - - (sampled_true, sampled_false) -} - -impl Drop for TestServerWithStatementLoggingChecks { - #[allow(clippy::disallowed_methods)] - fn drop(&mut self) { - // Don't run checks if we're already panicking, as this could mask the original error. - if std::thread::panicking() { - return; - } - - let mut mz_client = self - .server - .connect_internal(postgres::NoTls) - .expect("Failed to connect to internal SQL port for statement logging check"); - - // Disable RBAC checks so we can query mz_internal tables. - // (We don't need to restore this afterwards, since no more tests run in the same system.) - mz_client - .batch_execute("ALTER SYSTEM SET enable_rbac_checks = false") - .expect("Failed to disable RBAC checks"); - - // The statement log has a 5-second buffer flush interval, so allow sufficient time. - Retry::default() - .max_duration(Duration::from_secs(30)) - .retry(|_| { - let result = mz_client.query_one( - "SELECT count(*) - FROM mz_internal.mz_recent_activity_log - WHERE - (finished_at IS NULL OR finished_status IS NULL) - AND sql NOT LIKE '%__FILTER-OUT-THIS-QUERY__%' - AND finished_status IS DISTINCT FROM 'aborted'", - &[], - ); - - match result { - Ok(row) => { - let count: i64 = row.get(0); - if count == 0 { - Ok(()) - } else { - Err(format!("{} statements have not finished", count)) - } - } - Err(e) => Err(format!("Query failed: {}", e)), - } - }) - .expect("All statements should have finished executing"); - } -} - -fn setup_statement_logging_core( - max_sample_rate: f64, - sample_rate: f64, - target_data_rate: &str, - test_harness: test_util::TestHarness, -) -> (TestServerWithStatementLoggingChecks, postgres::Client) { - let server = test_harness - .with_system_parameter_default( - "statement_logging_max_sample_rate".to_string(), - max_sample_rate.to_string(), - ) - .with_system_parameter_default( - "statement_logging_default_sample_rate".to_string(), - sample_rate.to_string(), - ) - .with_system_parameter_default( - "statement_logging_max_data_credit".to_string(), - "".to_string(), - ) - .with_system_parameter_default( - "statement_logging_target_data_rate".to_string(), - target_data_rate.to_string(), - ) - .with_system_parameter_default( - "statement_logging_use_reproducible_rng".to_string(), - "true".to_string(), - ) - .start_blocking(); - let client = server.connect(postgres::NoTls).unwrap(); - let server = TestServerWithStatementLoggingChecks { server }; - (server, client) -} - -fn setup_statement_logging( - max_sample_rate: f64, - sample_rate: f64, - target_data_rate: &str, -) -> (TestServerWithStatementLoggingChecks, postgres::Client) { - setup_statement_logging_core( - max_sample_rate, - sample_rate, - target_data_rate, - test_util::TestHarness::default(), - ) -} - -// Test that we log various kinds of statement whose execution terminates in the coordinator. -#[mz_ore::test] -#[allow(clippy::disallowed_methods)] -fn test_statement_logging_immediate() { - let (server, mut client) = setup_statement_logging(1.0, 1.0, ""); - - let mut mz_client = server.connect_internal(postgres::NoTls).unwrap(); - mz_client - .batch_execute("ALTER SYSTEM SET enable_statement_lifecycle_logging = false") - .unwrap(); - mz_client - .batch_execute("ALTER SYSTEM SET statement_logging_max_sample_rate = 1") - .unwrap(); - mz_client - .batch_execute("ALTER SYSTEM SET statement_logging_default_sample_rate = 1") - .unwrap(); - mz_client - .batch_execute("ALTER SYSTEM SET enable_load_generator_counter = true") - .unwrap(); - - let successful_immediates: &[&str] = &[ - "CREATE VIEW v AS SELECT 1;", - "CREATE DEFAULT INDEX i ON v;", - "CREATE TABLE t (x bigint);", - "INSERT INTO t VALUES (1), (2), (3)", - "UPDATE t SET x=x+1", - "DELETE FROM t;", - "CREATE SECRET s AS 'hunter2';", - "DROP SECRET s;", - "", - "CREATE SOURCE s FROM LOAD GENERATOR COUNTER", - "PREPARE foo AS SELECT * FROM t", - "EXECUTE foo", - "BEGIN", - "DECLARE c CURSOR FOR SELECT * FROM t", - "FETCH FORWARD ALL FROM c", - "COMMIT", - "BEGIN", - "ROLLBACK", - "SET application_name='my_application'", - "SHOW ALL", - "SHOW application_name", - ]; - let constants: &[&str] = &["1", "2", "3", "hunter2", "my_application"]; - - for &statement in successful_immediates { - client.execute(statement, &[]).unwrap(); - - // Enforce a small delay to avoid duplicate `began_at` times, which would make the ordering - // of logged statements non-deterministic when we retrieve them below. - thread::sleep(Duration::from_millis(10)); - } - - let mut client = server.connect_internal(postgres::NoTls).unwrap(); - let seh_query = " - SELECT - mseh.sample_rate, - mseh.began_at, - mseh.finished_at, - mseh.finished_status, - mst.sql, - mpsh.prepared_at, - mst.redacted_sql - FROM mz_internal.mz_statement_execution_history AS mseh - LEFT JOIN - mz_internal.mz_prepared_statement_history AS mpsh - ON mseh.prepared_statement_id = mpsh.id - JOIN - (SELECT DISTINCT sql, sql_hash, redacted_sql FROM mz_internal.mz_sql_text) mst - ON mpsh.sql_hash = mst.sql_hash - WHERE - mst.sql !~~ '%mz_statement_execution_history%' AND - mseh.finished_at IS NOT NULL - ORDER BY mseh.began_at"; - - // Statement logging happens async, retry until we get the expected number of logged - // statements. - let mut sl = Vec::new(); - for _ in 0..10 { - thread::sleep(Duration::from_secs(1)); - sl = client.query(seh_query, &[]).unwrap(); - if sl.len() >= successful_immediates.len() { - break; - } - } - - assert_eq!(sl.len(), successful_immediates.len()); - - #[derive(Debug)] - struct Record { - sample_rate: f64, - began_at: DateTime, - finished_at: DateTime, - finished_status: String, - sql: String, - prepared_at: DateTime, - redacted_sql: String, - } - for (r, stmt) in std::iter::zip(sl.iter(), successful_immediates) { - let r = Record { - sample_rate: r.get(0), - began_at: r.get(1), - finished_at: r.get(2), - finished_status: r.get(3), - sql: r.get(4), - prepared_at: r.get(5), - redacted_sql: r.get(6), - }; - assert_eq!(r.sample_rate, 1.0); - - let expected_sql = if r.sql.contains("SECRET") - || r.sql.contains("INSERT") - || r.sql.contains("UPDATE") - || r.sql.contains("EXECUTE") - { - mz_sql::parse::parse(&r.sql) - .unwrap() - .into_element() - .ast - .to_ast_string_redacted() - } else { - stmt.chars().filter(|&ch| ch != ';').collect::() - }; - assert_eq!(r.sql, expected_sql); - assert_eq!(r.finished_status, "success"); - assert!(r.prepared_at <= r.began_at); - assert!(r.began_at <= r.finished_at); - // NB[btv] -- It would be a bit nicer if we could separately mock - // both the start and end time, but the `NowFn` mechanism doesn't - // appear to give us any way to do that. Instead, let's just check - // that none of these statements took longer than 5s wall-clock time. - assert!(r.finished_at - r.began_at <= chrono::Duration::try_seconds(5).unwrap()); - if !r.sql.is_empty() { - let expected_redacted = mz_sql::parse::parse(&r.sql) - .unwrap() - .into_element() - .ast - .to_ast_string_redacted(); - assert_eq!(r.redacted_sql, expected_redacted); - for constant in constants { - assert!(!r.redacted_sql.contains(constant)); - } - } - } -} - -#[mz_ore::test] -#[allow(clippy::disallowed_methods)] -fn test_statement_logging_basic() { - let (server, mut client) = setup_statement_logging(1.0, 1.0, ""); - client.execute("SELECT 1", &[]).unwrap(); - // We test that queries of this view execute on a cluster. - // If we ever change the threshold for constant folding such that - // this gets to run on environmentd, change this query. - client - .execute( - "CREATE VIEW v AS SELECT * FROM generate_series(1, 10001)", - &[], - ) - .unwrap(); - client.execute("SELECT * FROM v", &[]).unwrap(); - client.execute("CREATE DEFAULT INDEX i ON v", &[]).unwrap(); - client.execute("SELECT * FROM v", &[]).unwrap(); - let _ = client.execute("SELECT 1/0", &[]); - client.execute("CREATE TABLE t (x int)", &[]).unwrap(); - client.execute("SELECT * FROM t", &[]).unwrap(); - - #[derive(Debug)] - struct Record { - sample_rate: f64, - began_at: DateTime, - finished_at: DateTime, - finished_status: String, - error_message: Option, - prepared_at: DateTime, - execution_strategy: Option, - result_size: Option, - rows_returned: Option, - execution_timestamp: Option, - } - - let mut client = server.connect_internal(postgres::NoTls).unwrap(); - - let result = Retry::default() - .max_duration(Duration::from_secs(30)) - .retry(|_| { - let sl_results = client - .query( - "SELECT - mseh.sample_rate, - mseh.began_at, - mseh.finished_at, - mseh.finished_status, - mseh.error_message, - mpsh.prepared_at, - mseh.execution_strategy, - mseh.result_size, - mseh.rows_returned, - mseh.execution_timestamp -FROM - mz_internal.mz_statement_execution_history AS mseh - LEFT JOIN - mz_internal.mz_prepared_statement_history AS mpsh - ON mseh.prepared_statement_id = mpsh.id - JOIN - (SELECT DISTINCT sql, sql_hash, redacted_sql FROM mz_internal.mz_sql_text) AS mst - ON mpsh.sql_hash = mst.sql_hash -WHERE (mst.sql ~~ 'SELECT%' -AND mst.sql !~~ '%unique string to prevent this query showing up in results after retries%' -AND mst.sql !~~ '%pg_catalog.pg_type%' --this gets executed behind the scenes by tokio-postgres -OR mst.sql ~~ 'CREATE TABLE%') -AND mseh.finished_at IS NOT NULL -ORDER BY mseh.began_at", - &[], - ) - .unwrap(); - - if sl_results.len() == 6 { - Ok(sl_results) - } else { - Err(sl_results.len()) - } - }); - let sl_results = match result { - Ok(rows) => rows - .into_iter() - .map(|r| Record { - sample_rate: r.get(0), - began_at: r.get(1), - finished_at: r.get(2), - finished_status: r.get(3), - error_message: r.get(4), - prepared_at: r.get(5), - execution_strategy: r.get(6), - result_size: r.get(7), - rows_returned: r.get(8), - execution_timestamp: r.get::<_, Option>(9).map(|UInt8(val)| val), - }) - .collect::>(), - Err(rows) => { - panic!("number of results never became correct: {rows}"); - } - }; - // The two queries on generate_series(1,10001) execute at the maximum timestamp - assert_eq!( - sl_results - .iter() - .filter(|r| r.execution_timestamp == Some(u64::MAX)) - .count(), - 2 - ); - // The two queries that can be satisfied by envd (SELECT 1 and SELECT 1/0) have no execution timestamp - assert_eq!( - sl_results - .iter() - .filter(|r| r.execution_timestamp.is_none()) - .count(), - 2 - ); - // All other queries have an execution timestamp, in particular, including `CREATE TABLE`. - assert_eq!(sl_results.len(), 6); - for r in &sl_results { - assert_eq!(r.sample_rate, 1.0); - assert!(r.prepared_at <= r.began_at); - assert!(r.began_at <= r.finished_at); - // It would be nice to be able to control - // execution timestamp via a `NowFn`, but - // that is hard to get right and interferes with our logic - // about when to flush to persist. So instead, just check that they're sane. - if let Some(ts) = r.execution_timestamp { - if ts != u64::MAX { - let ts = to_datetime(ts); - assert!((ts - r.prepared_at).abs() < chrono::Duration::try_seconds(5).unwrap()) - } - } - } - assert!(sl_results[0].result_size.unwrap_or(0) > 0); - assert_eq!(sl_results[0].rows_returned, Some(1)); - assert_eq!(sl_results[0].finished_status, "success"); - assert_eq!( - sl_results[0].execution_strategy.as_ref().unwrap(), - "constant" - ); - assert!(sl_results[1].result_size.unwrap_or(0) > 0); - assert_eq!(sl_results[1].rows_returned, Some(10001)); - assert_eq!(sl_results[1].finished_status, "success"); - assert_eq!( - sl_results[1].execution_strategy.as_ref().unwrap(), - "standard" - ); - assert!(sl_results[2].result_size.unwrap_or(0) > 0); - assert_eq!(sl_results[2].rows_returned, Some(10001)); - assert_eq!(sl_results[2].finished_status, "success"); - assert_eq!( - sl_results[2].execution_strategy.as_ref().unwrap(), - "fast-path" - ); - assert_eq!(sl_results[3].finished_status, "error"); - assert!( - sl_results[3] - .error_message - .as_ref() - .unwrap() - .contains("division by zero") - ); - assert_none!(sl_results[3].result_size); - assert_none!(sl_results[3].rows_returned); - - // Verify metrics show all statements were sampled (100% sample rate means no unsampled). - let (sampled_true, sampled_false) = get_statement_logging_record_counts(&server); - assert!( - sampled_true > 0, - "some statements should be sampled with 100% rate" - ); - assert_eq!( - sampled_false, 0, - "no statements should be unsampled with 100% rate" - ); - - // Verify statement_logging_actual_bytes metric is being tracked. - // With 100% sample rate, actual_bytes should equal unsampled_bytes. - let metrics = server.metrics_registry().gather(); - let actual_bytes = metrics - .iter() - .find(|m| m.name() == "mz_statement_logging_actual_bytes") - .expect("mz_statement_logging_actual_bytes metric should exist") - .get_metric()[0] - .get_counter() - .value(); - let unsampled_bytes = metrics - .iter() - .find(|m| m.name() == "mz_statement_logging_unsampled_bytes") - .expect("mz_statement_logging_unsampled_bytes metric should exist") - .get_metric()[0] - .get_counter() - .value(); - assert!( - actual_bytes > 0.0, - "actual_bytes should be > 0 with 100% sample rate" - ); - assert_eq!( - actual_bytes, unsampled_bytes, - "with 100% sample rate, actual_bytes should equal unsampled_bytes" - ); -} - -#[allow(clippy::disallowed_methods)] -fn run_throttling_test(use_prepared_statement: bool) { - // The `target_data_rate` should be - // - high enough so that the `SELECT 1` queries get throttled (even with high CPU load due to - // other tests running in parallel), - // - but low enough that the `SELECT 2` query after the sleep doesn't get throttled. - let (server, mut client) = setup_statement_logging(1.0, 1.0, "200"); - thread::sleep(Duration::from_secs(2)); - - if use_prepared_statement { - let statement = client.prepare("SELECT 1").unwrap(); - for _ in 0..100 { - client.execute(&statement, &[]).unwrap(); - } - } else { - for _ in 0..100 { - client.execute("SELECT 1", &[]).unwrap(); - } - } - - thread::sleep(Duration::from_secs(4)); - client.execute("SELECT 2", &[]).unwrap(); - let mut client = server.connect_internal(postgres::NoTls).unwrap(); - let logs = Retry::default() - .max_duration(Duration::from_secs(60)) - .retry(|_| { - let sl_results = client - .query( - "SELECT - sql, - throttled_count -FROM mz_internal.mz_statement_execution_history mseh -JOIN mz_internal.mz_prepared_statement_history mpsh -ON mseh.prepared_statement_id = mpsh.id -JOIN (SELECT DISTINCT sql, sql_hash, redacted_sql FROM mz_internal.mz_sql_text) mst -ON mpsh.sql_hash = mst.sql_hash -WHERE sql IN ('SELECT 1', 'SELECT 2')", - &[], - ) - .unwrap(); - - if sl_results.iter().any(|stmt| { - let sql: String = stmt.get(0); - sql == "SELECT 2" - }) { - Ok(sl_results) - } else { - Err(()) - } - }) - .expect("Never saw last statement (`SELECT 2`)"); - let throttled_count = logs - .iter() - .map(|log| { - let UInt8(throttled_count) = log.get(1); - throttled_count - }) - .sum::(); - assert!( - throttled_count > 0, - "at least some statements should have been throttled" - ); - - assert_eq!(logs.len() + usize::cast_from(throttled_count), 101); -} - -#[mz_ore::test] -fn test_statement_logging_throttling() { - run_throttling_test(false); -} - -#[mz_ore::test] -fn test_statement_logging_prepared_statement_throttling() { - run_throttling_test(true); -} - -/// throttling the first execution of a prepared statement must not make -// every subsequent execution of that statement invisible in -// `mz_recent_activity_log`. -#[mz_ore::test] -#[allow(clippy::disallowed_methods)] -fn test_statement_logging_throttled_prepared_statement_stays_visible() { - const APPLICATION_NAME: &str = "throttled_prepared_repro"; - const NUM_SUBSEQUENT_STATEMENT_EXECUTIONS: i64 = 5; - - let (server, mut client) = setup_statement_logging(1.0, 1.0, "1"); - let mut mz_client = server.connect_internal(postgres::NoTls).unwrap(); - - // cap the credit at a single byte such that the first execution of our - // prepared statement is guaranteed to be throttled. - mz_client - .batch_execute("ALTER SYSTEM SET statement_logging_target_data_rate = 1") - .unwrap(); - mz_client - .batch_execute("ALTER SYSTEM SET statement_logging_max_data_credit = 1") - .unwrap(); - - client - .batch_execute(&format!("SET application_name = '{APPLICATION_NAME}'")) - .unwrap(); - - // First execution of the prepared statement. This one gets throttled, and so - // is never logged. - let stmt = client.prepare("SELECT $1::text").unwrap(); - assert_eq!( - client - .query_one(&stmt, &[&"first"]) - .unwrap() - .get::<_, String>(0), - "first" - ); - - // Lift the throttle so subsequent executions are logged. Sleep afterwards so - // the token bucket has time to refill before the next execution - mz_client - .batch_execute("ALTER SYSTEM SET statement_logging_target_data_rate = 1000000000") - .unwrap(); - mz_client - .batch_execute("ALTER SYSTEM SET statement_logging_max_data_credit = 1000000000") - .unwrap(); - thread::sleep(Duration::from_secs(2)); - - // Subsequent executions of the same prepared statement. - for _ in 0..NUM_SUBSEQUENT_STATEMENT_EXECUTIONS { - assert_eq!( - client - .query_one(&stmt, &[&"again"]) - .unwrap() - .get::<_, String>(0), - "again" - ); - } - - // Statement logging flushes asynchronously; wait until the - // subsequent executions have landed in `mz_statement_execution_history`. - Retry::default() - .max_duration(Duration::from_secs(30)) - .retry(|_| { - let mseh_count: i64 = mz_client - .query_one( - &format!( - "SELECT count(*) - FROM mz_internal.mz_statement_execution_history - WHERE application_name = '{APPLICATION_NAME}'" - ), - &[], - ) - .unwrap() - .get(0); - if mseh_count >= NUM_SUBSEQUENT_STATEMENT_EXECUTIONS { - Ok(()) - } else { - Err(format!( - "only {mseh_count} of {NUM_SUBSEQUENT_STATEMENT_EXECUTIONS} executions logged so far" - )) - } - }) - .expect("subsequent executions never appeared in mz_statement_execution_history"); - - // The subsequent executions in mz_statement_execution_history should - // have a matching entry in mz_prepared_statement_history - let orphan_count: i64 = mz_client - .query_one( - &format!( - "SELECT count(*) - FROM mz_internal.mz_statement_execution_history mseh - LEFT JOIN mz_internal.mz_prepared_statement_history mpsh - ON mseh.prepared_statement_id = mpsh.id - WHERE mseh.application_name = '{APPLICATION_NAME}' - AND mpsh.id IS NULL" - ), - &[], - ) - .unwrap() - .get(0); - assert_eq!( - orphan_count, 0, - "{orphan_count} mz_statement_execution_history rows reference a \ - prepared_statement_id with no matching mz_prepared_statement_history row" - ); -} - -#[mz_ore::test] -#[allow(clippy::disallowed_methods)] -fn test_statement_logging_subscribes() { - let (server, mut client) = setup_statement_logging(1.0, 1.0, ""); - let cancel_token = client.cancel_token(); - - // This should finish - client - .execute( - "SUBSCRIBE TO (SELECT * FROM generate_series(1, 10001))", - &[], - ) - .unwrap(); - - let handle = thread::spawn(move || { - client.execute("CREATE TABLE t (x int)", &[]).unwrap(); - // This should not finish until it's canceled. - let _ = client.execute("SUBSCRIBE TO (SELECT * FROM t)", &[]); - }); - - while !handle.is_finished() { - thread::sleep(Duration::from_secs(1)); - cancel_token.cancel_query(postgres::NoTls).unwrap(); - } - handle.join().unwrap(); - - let mut client = server.connect_internal(postgres::NoTls).unwrap(); - let seh_query = " - SELECT - mseh.sample_rate, - mseh.began_at, - mseh.finished_at, - mseh.finished_status, - mpsh.prepared_at, - mseh.execution_strategy - FROM mz_internal.mz_statement_execution_history AS mseh - LEFT JOIN - mz_internal.mz_prepared_statement_history AS mpsh - ON mseh.prepared_statement_id = mpsh.id - JOIN - mz_internal.mz_sql_text AS mst - ON mpsh.sql_hash = mst.sql_hash - WHERE - mst.sql ~~ 'SUBSCRIBE%' AND - mseh.finished_at IS NOT NULL - ORDER BY mseh.began_at"; - - // Statement logging happens async, retry until we get the expected number of logged - // statements. - let mut sl = Vec::new(); - for _ in 0..10 { - thread::sleep(Duration::from_secs(1)); - sl = client.query(seh_query, &[]).unwrap(); - if sl.len() >= 2 { - break; - } - } - - assert_eq!(sl.len(), 2); - - struct Record { - sample_rate: f64, - began_at: DateTime, - finished_at: DateTime, - finished_status: String, - prepared_at: DateTime, - execution_strategy: Option, - } - - let sl_subscribes = sl - .into_iter() - .map(|r| Record { - sample_rate: r.get(0), - began_at: r.get(1), - finished_at: r.get(2), - finished_status: r.get(3), - prepared_at: r.get(4), - execution_strategy: r.get(5), - }) - .collect::>(); - for r in &sl_subscribes { - assert_eq!(r.sample_rate, 1.0); - assert!(r.prepared_at <= r.began_at); - assert!(r.began_at <= r.finished_at); - assert_none!(r.execution_strategy); - } - assert_eq!(sl_subscribes[0].finished_status, "success"); - assert_eq!(sl_subscribes[1].finished_status, "canceled"); -} - -/// Test that we are sampling approximately 50% of statements. -/// Relies on two assumptions: -/// (1) that the effective sampling rate for the session is 50%, -/// (2) that we are using the deterministic testing RNG. -#[allow(clippy::disallowed_methods)] -fn test_statement_logging_sampling_inner( - server: TestServerWithStatementLoggingChecks, - mut client: postgres::Client, -) { - for i in 0..50 { - client.execute(&format!("SELECT {i}"), &[]).unwrap(); - - // Enforce a small delay to avoid duplicate `began_at` times, which would make the ordering - // of logged statements non-deterministic when we retrieve them below. - thread::sleep(Duration::from_millis(10)); - } - - // 23 randomly sampled out of 50 with 50% sampling. Seems legit! - let expected_sqls = [ - 2, 4, 5, 6, 9, 15, 17, 18, 19, 20, 21, 23, 24, 25, 31, 32, 33, 36, 37, 42, 46, - ] - .into_iter() - .map(|i| format!("SELECT {i}")) - .collect::>(); - - let mut internal_client = server.connect_internal(postgres::NoTls).unwrap(); - let seh_query = " - SELECT mst.sql - FROM mz_internal.mz_statement_execution_history AS mseh - JOIN - mz_internal.mz_prepared_statement_history AS mpsh - ON mseh.prepared_statement_id = mpsh.id - JOIN - mz_internal.mz_sql_text AS mst - ON mpsh.sql_hash = mst.sql_hash - WHERE mst.sql ~~ 'SELECT%' AND mst.sql !~~ '%mz_statement_execution_history%' - ORDER BY mseh.began_at ASC"; - - // Statement logging happens async, retry until we get the expected number of logged - // statements. - let mut sl = Vec::new(); - for _ in 0..10 { - thread::sleep(Duration::from_secs(1)); - sl = internal_client.query(seh_query, &[]).unwrap(); - if sl.len() >= expected_sqls.len() { - break; - } - } - - let sqls: Vec = sl.into_iter().map(|r| r.get(0)).collect(); - assert_eq!(sqls, expected_sqls); - - // Verify the statement_logging_record_count metric correctly tracks sampled vs unsampled. - // With 50% sampling and deterministic RNG, exactly 21 of 50 statements should be sampled. - let (sampled_true, sampled_false) = get_statement_logging_record_counts(&server); - assert_eq!( - sampled_true, 21, - "expected 21 statements to be sampled with 50% rate and deterministic RNG" - ); - assert_eq!( - sampled_false, 29, - "expected 29 statements to not be sampled with 50% rate and deterministic RNG" - ); -} - -#[mz_ore::test] -fn test_statement_logging_sampling() { - let (server, client) = setup_statement_logging(1.0, 0.5, ""); - test_statement_logging_sampling_inner(server, client); -} - -/// Test that we are not allowed to set `statement_logging_sample_rate` -/// arbitrarily high, but that it is constrained by `statement_logging_max_sample_rate`. -#[mz_ore::test] -fn test_statement_logging_sampling_constrained() { - let (server, client) = setup_statement_logging(0.5, 1.0, ""); - test_statement_logging_sampling_inner(server, client); -} - -/// Test that the `mz_statement_logging_unsampled_bytes` metric tracks the total bytes -/// of SQL text that would have been logged if statement logging were fully enabled. -/// We set `sample_rate=0.0` so no statements are actually sampled/logged, but the -/// unsampled_bytes metric still gets incremented for every executed statement. -#[mz_ore::test] -#[allow(clippy::disallowed_methods)] -fn test_statement_logging_unsampled_metrics() { - // Use sample_rate=0.0 so statements are not sampled, but unsampled_bytes metric is still tracked. - let (server, mut client) = setup_statement_logging(1.0, 0.0, ""); - - let batch_queries = [ - "SELECT 'Hello, world!';SELECT 1;;", - "SELECT 'Hello, world again!'", - ]; - let batch_total: usize = batch_queries - .iter() - .map(|s| s.as_bytes().iter().filter(|&&ch| ch != b';').count()) - .sum(); - let single_queries = ["SELECT 'foo'", "SELECT 'bar';;;"]; - let single_total: usize = single_queries - .iter() - .map(|s| s.as_bytes().iter().filter(|&&ch| ch != b';').count()) - .sum(); - let prepared_queries = ["SELECT 'baz';;;", "SELECT 'quux';"]; - let prepared_total: usize = prepared_queries - .iter() - .map(|s| s.as_bytes().iter().filter(|&&ch| ch != b';').count()) - .sum(); - - let named_prepared_inner = "SELECT 42"; - let named_prepared_outer = format!("PREPARE p AS {named_prepared_inner};EXECUTE p;"); - let named_prepared_outer_len = named_prepared_outer - .as_bytes() - .iter() - .filter(|&&ch| ch != b';') - .count(); - - for q in batch_queries { - client.batch_execute(q).unwrap(); - } - - for q in single_queries { - client.execute(q, &[]).unwrap(); - } - - for q in prepared_queries { - let s = client.prepare(q).unwrap(); - client.execute(&s, &[]).unwrap(); - } - - client.batch_execute(&named_prepared_outer).unwrap(); - - // This should NOT be logged, since we never actually execute it. - client.prepare("SELECT 'Hello, not counted!'").unwrap(); - - let expected_total = batch_total + single_total + prepared_total + named_prepared_outer_len; - let metric_value = server - .metrics_registry() - .gather() - .into_iter() - .find(|m| m.name() == "mz_statement_logging_unsampled_bytes") - .unwrap() - .take_metric()[0] - .get_counter() - .value(); - let metric_value = usize::cast_from(u64::try_cast_from(metric_value).unwrap()); - assert_eq!(expected_total, metric_value); - - // Also verify that statement_logging_record_count shows all statements as not sampled - // (since we're using 0% sample rate). - let (sampled_true, _sampled_false) = get_statement_logging_record_counts(&server); - assert_eq!( - sampled_true, 0, - "no statements should be sampled with 0% sample rate" - ); -} - -#[mz_ore::test] -#[allow(clippy::disallowed_methods)] -fn test_enable_internal_statement_logging() { - let (server, mut client) = setup_statement_logging_core( - 1.0, - 1.0, - "", - test_util::TestHarness::default().with_system_parameter_default( - "enable_internal_statement_logging".to_string(), - "true".to_string(), - ), - ); - - client.execute("SELECT 1", &[]).unwrap(); - - let mut client = server.connect_internal(postgres::NoTls).unwrap(); - let num_mz_system_statements = Retry::default() - .max_duration(Duration::from_secs(30)) - .retry(|_| { - let sl_results = client - .query( - "SELECT - count(*) -FROM mz_internal.mz_prepared_statement_history mpsh -JOIN mz_internal.mz_session_history USING (session_id) -WHERE authenticated_user='mz_system'", - &[], - ) - .unwrap(); - - let count: i64 = sl_results[0].get(0); - - if count > 0 { Ok(count) } else { Err(()) } - }) - .expect("at least some statements from mz_system should have been logged"); - - assert!( - num_mz_system_statements > 0, - "statements executed by mz_system should have been logged" - ); -} - // Test the POST and WS server endpoints. #[mz_ore::test] #[allow(clippy::disallowed_methods)] @@ -2246,28 +1270,6 @@ fn test_ws_passes_options() { } } -#[mz_ore::test] -// Test that statement logging -// doesn't cause a crash with subscribes over web sockets, -// which was previously happening (in staging) due to us -// dropping the `ExecuteContext` on the floor in that case. -fn test_statement_logging_ws_subscribe_no_crash() { - let (server, _client) = setup_statement_logging(1.0, 1.0, ""); - - // Create our WebSocket. - let ws_url = server.server.ws_addr(); - let (mut ws, _resp) = tungstenite::connect(ws_url).unwrap(); - test_util::auth_with_ws(&mut ws, Default::default()).unwrap(); - - let query = "SUBSCRIBE (SELECT 1)"; - let json = format!("{{\"query\":\"{query}\"}}"); - let json: serde_json::Value = serde_json::from_str(&json).unwrap(); - ws.send(Message::Text(json.to_string().into())).unwrap(); - - // Give the server time to crash, if it's going to. - std::thread::sleep(Duration::from_secs(1)) -} - #[mz_ore::test] #[cfg_attr(miri, ignore)] // too slow fn test_ws_notifies_for_bad_options() { diff --git a/src/environmentd/tests/statement_logging.rs b/src/environmentd/tests/statement_logging.rs new file mode 100644 index 0000000000000..859ef1dd493a1 --- /dev/null +++ b/src/environmentd/tests/statement_logging.rs @@ -0,0 +1,1091 @@ +// Copyright Materialize, Inc. and contributors. All rights reserved. +// +// Use of this software is governed by the Business Source License +// included in the LICENSE file. +// +// As of the Change Date specified in that file, in accordance with +// the Business Source License, use of this software will be governed +// by the Apache License, Version 2.0. + +//! Integration tests for statement logging. + +use std::thread; +use std::time::Duration; + +use chrono::{DateTime, Utc}; +use mz_environmentd::test_util; +use mz_ore::assert_none; +use mz_ore::cast::{CastFrom, CastLossy, TryCastFrom}; +use mz_ore::collections::CollectionExt; +use mz_ore::metrics::MetricsRegistry; +use mz_ore::now::to_datetime; +use mz_ore::retry::Retry; +use mz_pgrepr::UInt8; +use mz_sql_parser::ast::display::AstDisplay; +use tungstenite::Message; + +/// A wrapper around `TestServerWithRuntime` that runs statement logging checks when dropped. +/// +/// This guard ensures that all statements have finished executing (have non-NULL `finished_at` +/// and `finished_status` in `mz_internal.mz_recent_activity_log`) before the test completes. +struct TestServerWithStatementLoggingChecks { + server: test_util::TestServerWithRuntime, +} + +impl TestServerWithStatementLoggingChecks { + /// Connect to the __internal__ SQL port of the running `environmentd` server. + pub fn connect_internal(&self, tls: T) -> Result + where + T: postgres::tls::MakeTlsConnect + Send + 'static, + T::TlsConnect: Send, + T::Stream: Send, + >::Future: Send, + { + self.server.connect_internal(tls) + } + + /// Returns the metrics registry for the test server. + pub fn metrics_registry(&self) -> &MetricsRegistry { + self.server.metrics_registry() + } +} + +/// Helper to get statement logging record counts from the metrics registry. +/// Returns (sampled_true_count, sampled_false_count). +#[allow(clippy::disallowed_methods)] +fn get_statement_logging_record_counts( + server: &TestServerWithStatementLoggingChecks, +) -> (u64, u64) { + let metrics = server.metrics_registry().gather(); + let record_count_metric = metrics + .into_iter() + .find(|m| m.name() == "mz_statement_logging_record_count") + .expect("mz_statement_logging_record_count metric should exist"); + + let metric_entries = record_count_metric.get_metric(); + let sampled_true = metric_entries + .iter() + .find(|m| { + m.get_label() + .iter() + .any(|l| l.name() == "sample" && l.value() == "true") + }) + .map(|m| u64::cast_lossy(m.get_counter().value())) + .unwrap_or(0); + let sampled_false = metric_entries + .iter() + .find(|m| { + m.get_label() + .iter() + .any(|l| l.name() == "sample" && l.value() == "false") + }) + .map(|m| u64::cast_lossy(m.get_counter().value())) + .unwrap_or(0); + + (sampled_true, sampled_false) +} + +impl Drop for TestServerWithStatementLoggingChecks { + #[allow(clippy::disallowed_methods)] + fn drop(&mut self) { + // Don't run checks if we're already panicking, as this could mask the original error. + if std::thread::panicking() { + return; + } + + let mut mz_client = self + .server + .connect_internal(postgres::NoTls) + .expect("Failed to connect to internal SQL port for statement logging check"); + + // Disable RBAC checks so we can query mz_internal tables. + // (We don't need to restore this afterwards, since no more tests run in the same system.) + mz_client + .batch_execute("ALTER SYSTEM SET enable_rbac_checks = false") + .expect("Failed to disable RBAC checks"); + + // The statement log has a 5-second buffer flush interval, so allow sufficient time. + Retry::default() + .max_duration(Duration::from_secs(30)) + .retry(|_| { + let result = mz_client.query_one( + "SELECT count(*) + FROM mz_internal.mz_recent_activity_log + WHERE + (finished_at IS NULL OR finished_status IS NULL) + AND sql NOT LIKE '%__FILTER-OUT-THIS-QUERY__%' + AND finished_status IS DISTINCT FROM 'aborted'", + &[], + ); + + match result { + Ok(row) => { + let count: i64 = row.get(0); + if count == 0 { + Ok(()) + } else { + Err(format!("{} statements have not finished", count)) + } + } + Err(e) => Err(format!("Query failed: {}", e)), + } + }) + .expect("All statements should have finished executing"); + } +} + +fn setup_statement_logging_core( + max_sample_rate: f64, + sample_rate: f64, + target_data_rate: &str, + test_harness: test_util::TestHarness, +) -> (TestServerWithStatementLoggingChecks, postgres::Client) { + let server = test_harness + .with_system_parameter_default( + "statement_logging_max_sample_rate".to_string(), + max_sample_rate.to_string(), + ) + .with_system_parameter_default( + "statement_logging_default_sample_rate".to_string(), + sample_rate.to_string(), + ) + .with_system_parameter_default( + "statement_logging_max_data_credit".to_string(), + "".to_string(), + ) + .with_system_parameter_default( + "statement_logging_target_data_rate".to_string(), + target_data_rate.to_string(), + ) + .with_system_parameter_default( + "statement_logging_use_reproducible_rng".to_string(), + "true".to_string(), + ) + .start_blocking(); + let client = server.connect(postgres::NoTls).unwrap(); + let server = TestServerWithStatementLoggingChecks { server }; + (server, client) +} + +fn setup_statement_logging( + max_sample_rate: f64, + sample_rate: f64, + target_data_rate: &str, +) -> (TestServerWithStatementLoggingChecks, postgres::Client) { + setup_statement_logging_core( + max_sample_rate, + sample_rate, + target_data_rate, + test_util::TestHarness::default(), + ) +} + +// Test that we log various kinds of statement whose execution terminates in the coordinator. +#[mz_ore::test] +#[allow(clippy::disallowed_methods)] +fn test_statement_logging_immediate() { + let (server, mut client) = setup_statement_logging(1.0, 1.0, ""); + + let mut mz_client = server.connect_internal(postgres::NoTls).unwrap(); + mz_client + .batch_execute("ALTER SYSTEM SET enable_statement_lifecycle_logging = false") + .unwrap(); + mz_client + .batch_execute("ALTER SYSTEM SET statement_logging_max_sample_rate = 1") + .unwrap(); + mz_client + .batch_execute("ALTER SYSTEM SET statement_logging_default_sample_rate = 1") + .unwrap(); + mz_client + .batch_execute("ALTER SYSTEM SET enable_load_generator_counter = true") + .unwrap(); + + let successful_immediates: &[&str] = &[ + "CREATE VIEW v AS SELECT 1;", + "CREATE DEFAULT INDEX i ON v;", + "CREATE TABLE t (x bigint);", + "INSERT INTO t VALUES (1), (2), (3)", + "UPDATE t SET x=x+1", + "DELETE FROM t;", + "CREATE SECRET s AS 'hunter2';", + "DROP SECRET s;", + "", + "CREATE SOURCE s FROM LOAD GENERATOR COUNTER", + "PREPARE foo AS SELECT * FROM t", + "EXECUTE foo", + "BEGIN", + "DECLARE c CURSOR FOR SELECT * FROM t", + "FETCH FORWARD ALL FROM c", + "COMMIT", + "BEGIN", + "ROLLBACK", + "SET application_name='my_application'", + "SHOW ALL", + "SHOW application_name", + ]; + let constants: &[&str] = &["1", "2", "3", "hunter2", "my_application"]; + + for &statement in successful_immediates { + client.execute(statement, &[]).unwrap(); + + // Enforce a small delay to avoid duplicate `began_at` times, which would make the ordering + // of logged statements non-deterministic when we retrieve them below. + thread::sleep(Duration::from_millis(10)); + } + + let mut client = server.connect_internal(postgres::NoTls).unwrap(); + let seh_query = " + SELECT + mseh.sample_rate, + mseh.began_at, + mseh.finished_at, + mseh.finished_status, + mst.sql, + mpsh.prepared_at, + mst.redacted_sql + FROM mz_internal.mz_statement_execution_history AS mseh + LEFT JOIN + mz_internal.mz_prepared_statement_history AS mpsh + ON mseh.prepared_statement_id = mpsh.id + JOIN + (SELECT DISTINCT sql, sql_hash, redacted_sql FROM mz_internal.mz_sql_text) mst + ON mpsh.sql_hash = mst.sql_hash + WHERE + mst.sql !~~ '%mz_statement_execution_history%' AND + mseh.finished_at IS NOT NULL + ORDER BY mseh.began_at"; + + // Statement logging happens async, retry until we get the expected number of logged + // statements. + let mut sl = Vec::new(); + for _ in 0..10 { + thread::sleep(Duration::from_secs(1)); + sl = client.query(seh_query, &[]).unwrap(); + if sl.len() >= successful_immediates.len() { + break; + } + } + + assert_eq!(sl.len(), successful_immediates.len()); + + #[derive(Debug)] + struct Record { + sample_rate: f64, + began_at: DateTime, + finished_at: DateTime, + finished_status: String, + sql: String, + prepared_at: DateTime, + redacted_sql: String, + } + for (r, stmt) in std::iter::zip(sl.iter(), successful_immediates) { + let r = Record { + sample_rate: r.get(0), + began_at: r.get(1), + finished_at: r.get(2), + finished_status: r.get(3), + sql: r.get(4), + prepared_at: r.get(5), + redacted_sql: r.get(6), + }; + assert_eq!(r.sample_rate, 1.0); + + let expected_sql = if r.sql.contains("SECRET") + || r.sql.contains("INSERT") + || r.sql.contains("UPDATE") + || r.sql.contains("EXECUTE") + { + mz_sql::parse::parse(&r.sql) + .unwrap() + .into_element() + .ast + .to_ast_string_redacted() + } else { + stmt.chars().filter(|&ch| ch != ';').collect::() + }; + assert_eq!(r.sql, expected_sql); + assert_eq!(r.finished_status, "success"); + assert!(r.prepared_at <= r.began_at); + assert!(r.began_at <= r.finished_at); + // NB[btv] -- It would be a bit nicer if we could separately mock + // both the start and end time, but the `NowFn` mechanism doesn't + // appear to give us any way to do that. Instead, let's just check + // that none of these statements took longer than 5s wall-clock time. + assert!(r.finished_at - r.began_at <= chrono::Duration::try_seconds(5).unwrap()); + if !r.sql.is_empty() { + let expected_redacted = mz_sql::parse::parse(&r.sql) + .unwrap() + .into_element() + .ast + .to_ast_string_redacted(); + assert_eq!(r.redacted_sql, expected_redacted); + for constant in constants { + assert!(!r.redacted_sql.contains(constant)); + } + } + } +} + +#[mz_ore::test] +#[allow(clippy::disallowed_methods)] +fn test_statement_logging_basic() { + let (server, mut client) = setup_statement_logging(1.0, 1.0, ""); + client.execute("SELECT 1", &[]).unwrap(); + // We test that queries of this view execute on a cluster. + // If we ever change the threshold for constant folding such that + // this gets to run on environmentd, change this query. + client + .execute( + "CREATE VIEW v AS SELECT * FROM generate_series(1, 10001)", + &[], + ) + .unwrap(); + client.execute("SELECT * FROM v", &[]).unwrap(); + client.execute("CREATE DEFAULT INDEX i ON v", &[]).unwrap(); + client.execute("SELECT * FROM v", &[]).unwrap(); + let _ = client.execute("SELECT 1/0", &[]); + client.execute("CREATE TABLE t (x int)", &[]).unwrap(); + client.execute("SELECT * FROM t", &[]).unwrap(); + + #[derive(Debug)] + struct Record { + sample_rate: f64, + began_at: DateTime, + finished_at: DateTime, + finished_status: String, + error_message: Option, + prepared_at: DateTime, + execution_strategy: Option, + result_size: Option, + rows_returned: Option, + execution_timestamp: Option, + } + + let mut client = server.connect_internal(postgres::NoTls).unwrap(); + + let result = Retry::default() + .max_duration(Duration::from_secs(30)) + .retry(|_| { + let sl_results = client + .query( + "SELECT + mseh.sample_rate, + mseh.began_at, + mseh.finished_at, + mseh.finished_status, + mseh.error_message, + mpsh.prepared_at, + mseh.execution_strategy, + mseh.result_size, + mseh.rows_returned, + mseh.execution_timestamp +FROM + mz_internal.mz_statement_execution_history AS mseh + LEFT JOIN + mz_internal.mz_prepared_statement_history AS mpsh + ON mseh.prepared_statement_id = mpsh.id + JOIN + (SELECT DISTINCT sql, sql_hash, redacted_sql FROM mz_internal.mz_sql_text) AS mst + ON mpsh.sql_hash = mst.sql_hash +WHERE (mst.sql ~~ 'SELECT%' +AND mst.sql !~~ '%unique string to prevent this query showing up in results after retries%' +AND mst.sql !~~ '%pg_catalog.pg_type%' --this gets executed behind the scenes by tokio-postgres +OR mst.sql ~~ 'CREATE TABLE%') +AND mseh.finished_at IS NOT NULL +ORDER BY mseh.began_at", + &[], + ) + .unwrap(); + + if sl_results.len() == 6 { + Ok(sl_results) + } else { + Err(sl_results.len()) + } + }); + let sl_results = match result { + Ok(rows) => rows + .into_iter() + .map(|r| Record { + sample_rate: r.get(0), + began_at: r.get(1), + finished_at: r.get(2), + finished_status: r.get(3), + error_message: r.get(4), + prepared_at: r.get(5), + execution_strategy: r.get(6), + result_size: r.get(7), + rows_returned: r.get(8), + execution_timestamp: r.get::<_, Option>(9).map(|UInt8(val)| val), + }) + .collect::>(), + Err(rows) => { + panic!("number of results never became correct: {rows}"); + } + }; + // The two queries on generate_series(1,10001) execute at the maximum timestamp + assert_eq!( + sl_results + .iter() + .filter(|r| r.execution_timestamp == Some(u64::MAX)) + .count(), + 2 + ); + // The two queries that can be satisfied by envd (SELECT 1 and SELECT 1/0) have no execution timestamp + assert_eq!( + sl_results + .iter() + .filter(|r| r.execution_timestamp.is_none()) + .count(), + 2 + ); + // All other queries have an execution timestamp, in particular, including `CREATE TABLE`. + assert_eq!(sl_results.len(), 6); + for r in &sl_results { + assert_eq!(r.sample_rate, 1.0); + assert!(r.prepared_at <= r.began_at); + assert!(r.began_at <= r.finished_at); + // It would be nice to be able to control + // execution timestamp via a `NowFn`, but + // that is hard to get right and interferes with our logic + // about when to flush to persist. So instead, just check that they're sane. + if let Some(ts) = r.execution_timestamp { + if ts != u64::MAX { + let ts = to_datetime(ts); + assert!((ts - r.prepared_at).abs() < chrono::Duration::try_seconds(5).unwrap()) + } + } + } + assert!(sl_results[0].result_size.unwrap_or(0) > 0); + assert_eq!(sl_results[0].rows_returned, Some(1)); + assert_eq!(sl_results[0].finished_status, "success"); + assert_eq!( + sl_results[0].execution_strategy.as_ref().unwrap(), + "constant" + ); + assert!(sl_results[1].result_size.unwrap_or(0) > 0); + assert_eq!(sl_results[1].rows_returned, Some(10001)); + assert_eq!(sl_results[1].finished_status, "success"); + assert_eq!( + sl_results[1].execution_strategy.as_ref().unwrap(), + "standard" + ); + assert!(sl_results[2].result_size.unwrap_or(0) > 0); + assert_eq!(sl_results[2].rows_returned, Some(10001)); + assert_eq!(sl_results[2].finished_status, "success"); + assert_eq!( + sl_results[2].execution_strategy.as_ref().unwrap(), + "fast-path" + ); + assert_eq!(sl_results[3].finished_status, "error"); + assert!( + sl_results[3] + .error_message + .as_ref() + .unwrap() + .contains("division by zero") + ); + assert_none!(sl_results[3].result_size); + assert_none!(sl_results[3].rows_returned); + + // Verify metrics show all statements were sampled (100% sample rate means no unsampled). + let (sampled_true, sampled_false) = get_statement_logging_record_counts(&server); + assert!( + sampled_true > 0, + "some statements should be sampled with 100% rate" + ); + assert_eq!( + sampled_false, 0, + "no statements should be unsampled with 100% rate" + ); + + // Verify statement_logging_actual_bytes metric is being tracked. + // With 100% sample rate, actual_bytes should equal unsampled_bytes. + let metrics = server.metrics_registry().gather(); + let actual_bytes = metrics + .iter() + .find(|m| m.name() == "mz_statement_logging_actual_bytes") + .expect("mz_statement_logging_actual_bytes metric should exist") + .get_metric()[0] + .get_counter() + .value(); + let unsampled_bytes = metrics + .iter() + .find(|m| m.name() == "mz_statement_logging_unsampled_bytes") + .expect("mz_statement_logging_unsampled_bytes metric should exist") + .get_metric()[0] + .get_counter() + .value(); + assert!( + actual_bytes > 0.0, + "actual_bytes should be > 0 with 100% sample rate" + ); + assert_eq!( + actual_bytes, unsampled_bytes, + "with 100% sample rate, actual_bytes should equal unsampled_bytes" + ); +} + +#[allow(clippy::disallowed_methods)] +fn run_throttling_test(use_prepared_statement: bool) { + // The `target_data_rate` should be + // - high enough so that the `SELECT 1` queries get throttled (even with high CPU load due to + // other tests running in parallel), + // - but low enough that the `SELECT 2` query after the sleep doesn't get throttled. + let (server, mut client) = setup_statement_logging(1.0, 1.0, "200"); + thread::sleep(Duration::from_secs(2)); + + if use_prepared_statement { + let statement = client.prepare("SELECT 1").unwrap(); + for _ in 0..100 { + client.execute(&statement, &[]).unwrap(); + } + } else { + for _ in 0..100 { + client.execute("SELECT 1", &[]).unwrap(); + } + } + + thread::sleep(Duration::from_secs(4)); + client.execute("SELECT 2", &[]).unwrap(); + let mut client = server.connect_internal(postgres::NoTls).unwrap(); + let logs = Retry::default() + .max_duration(Duration::from_secs(60)) + .retry(|_| { + let sl_results = client + .query( + "SELECT + sql, + throttled_count +FROM mz_internal.mz_statement_execution_history mseh +JOIN mz_internal.mz_prepared_statement_history mpsh +ON mseh.prepared_statement_id = mpsh.id +JOIN (SELECT DISTINCT sql, sql_hash, redacted_sql FROM mz_internal.mz_sql_text) mst +ON mpsh.sql_hash = mst.sql_hash +WHERE sql IN ('SELECT 1', 'SELECT 2')", + &[], + ) + .unwrap(); + + if sl_results.iter().any(|stmt| { + let sql: String = stmt.get(0); + sql == "SELECT 2" + }) { + Ok(sl_results) + } else { + Err(()) + } + }) + .expect("Never saw last statement (`SELECT 2`)"); + let throttled_count = logs + .iter() + .map(|log| { + let UInt8(throttled_count) = log.get(1); + throttled_count + }) + .sum::(); + assert!( + throttled_count > 0, + "at least some statements should have been throttled" + ); + + assert_eq!(logs.len() + usize::cast_from(throttled_count), 101); +} + +#[mz_ore::test] +fn test_statement_logging_throttling() { + run_throttling_test(false); +} + +#[mz_ore::test] +fn test_statement_logging_prepared_statement_throttling() { + run_throttling_test(true); +} + +/// throttling the first execution of a prepared statement must not make +// every subsequent execution of that statement invisible in +// `mz_recent_activity_log`. +#[mz_ore::test] +#[allow(clippy::disallowed_methods)] +fn test_statement_logging_throttled_prepared_statement_stays_visible() { + const APPLICATION_NAME: &str = "throttled_prepared_repro"; + const NUM_SUBSEQUENT_STATEMENT_EXECUTIONS: i64 = 5; + + let (server, mut client) = setup_statement_logging(1.0, 1.0, "1"); + let mut mz_client = server.connect_internal(postgres::NoTls).unwrap(); + + // cap the credit at a single byte such that the first execution of our + // prepared statement is guaranteed to be throttled. + mz_client + .batch_execute("ALTER SYSTEM SET statement_logging_target_data_rate = 1") + .unwrap(); + mz_client + .batch_execute("ALTER SYSTEM SET statement_logging_max_data_credit = 1") + .unwrap(); + + client + .batch_execute(&format!("SET application_name = '{APPLICATION_NAME}'")) + .unwrap(); + + // First execution of the prepared statement. This one gets throttled, and so + // is never logged. + let stmt = client.prepare("SELECT $1::text").unwrap(); + assert_eq!( + client + .query_one(&stmt, &[&"first"]) + .unwrap() + .get::<_, String>(0), + "first" + ); + + // Lift the throttle so subsequent executions are logged. Sleep afterwards so + // the token bucket has time to refill before the next execution + mz_client + .batch_execute("ALTER SYSTEM SET statement_logging_target_data_rate = 1000000000") + .unwrap(); + mz_client + .batch_execute("ALTER SYSTEM SET statement_logging_max_data_credit = 1000000000") + .unwrap(); + thread::sleep(Duration::from_secs(2)); + + // Subsequent executions of the same prepared statement. + for _ in 0..NUM_SUBSEQUENT_STATEMENT_EXECUTIONS { + assert_eq!( + client + .query_one(&stmt, &[&"again"]) + .unwrap() + .get::<_, String>(0), + "again" + ); + } + + // Statement logging flushes asynchronously; wait until the + // subsequent executions have landed in `mz_statement_execution_history`. + Retry::default() + .max_duration(Duration::from_secs(30)) + .retry(|_| { + let mseh_count: i64 = mz_client + .query_one( + &format!( + "SELECT count(*) + FROM mz_internal.mz_statement_execution_history + WHERE application_name = '{APPLICATION_NAME}'" + ), + &[], + ) + .unwrap() + .get(0); + if mseh_count >= NUM_SUBSEQUENT_STATEMENT_EXECUTIONS { + Ok(()) + } else { + Err(format!( + "only {mseh_count} of {NUM_SUBSEQUENT_STATEMENT_EXECUTIONS} executions logged so far" + )) + } + }) + .expect("subsequent executions never appeared in mz_statement_execution_history"); + + // The subsequent executions in mz_statement_execution_history should + // have a matching entry in mz_prepared_statement_history + let orphan_count: i64 = mz_client + .query_one( + &format!( + "SELECT count(*) + FROM mz_internal.mz_statement_execution_history mseh + LEFT JOIN mz_internal.mz_prepared_statement_history mpsh + ON mseh.prepared_statement_id = mpsh.id + WHERE mseh.application_name = '{APPLICATION_NAME}' + AND mpsh.id IS NULL" + ), + &[], + ) + .unwrap() + .get(0); + assert_eq!( + orphan_count, 0, + "{orphan_count} mz_statement_execution_history rows reference a \ + prepared_statement_id with no matching mz_prepared_statement_history row" + ); +} + +#[mz_ore::test] +#[allow(clippy::disallowed_methods)] +fn test_statement_logging_subscribes() { + let (server, mut client) = setup_statement_logging(1.0, 1.0, ""); + let cancel_token = client.cancel_token(); + + // This should finish + client + .execute( + "SUBSCRIBE TO (SELECT * FROM generate_series(1, 10001))", + &[], + ) + .unwrap(); + + let handle = thread::spawn(move || { + client.execute("CREATE TABLE t (x int)", &[]).unwrap(); + // This should not finish until it's canceled. + let _ = client.execute("SUBSCRIBE TO (SELECT * FROM t)", &[]); + }); + + while !handle.is_finished() { + thread::sleep(Duration::from_secs(1)); + cancel_token.cancel_query(postgres::NoTls).unwrap(); + } + handle.join().unwrap(); + + let mut client = server.connect_internal(postgres::NoTls).unwrap(); + let seh_query = " + SELECT + mseh.sample_rate, + mseh.began_at, + mseh.finished_at, + mseh.finished_status, + mpsh.prepared_at, + mseh.execution_strategy + FROM mz_internal.mz_statement_execution_history AS mseh + LEFT JOIN + mz_internal.mz_prepared_statement_history AS mpsh + ON mseh.prepared_statement_id = mpsh.id + JOIN + mz_internal.mz_sql_text AS mst + ON mpsh.sql_hash = mst.sql_hash + WHERE + mst.sql ~~ 'SUBSCRIBE%' AND + mseh.finished_at IS NOT NULL + ORDER BY mseh.began_at"; + + // Statement logging happens async, retry until we get the expected number of logged + // statements. + let mut sl = Vec::new(); + for _ in 0..10 { + thread::sleep(Duration::from_secs(1)); + sl = client.query(seh_query, &[]).unwrap(); + if sl.len() >= 2 { + break; + } + } + + assert_eq!(sl.len(), 2); + + struct Record { + sample_rate: f64, + began_at: DateTime, + finished_at: DateTime, + finished_status: String, + prepared_at: DateTime, + execution_strategy: Option, + } + + let sl_subscribes = sl + .into_iter() + .map(|r| Record { + sample_rate: r.get(0), + began_at: r.get(1), + finished_at: r.get(2), + finished_status: r.get(3), + prepared_at: r.get(4), + execution_strategy: r.get(5), + }) + .collect::>(); + for r in &sl_subscribes { + assert_eq!(r.sample_rate, 1.0); + assert!(r.prepared_at <= r.began_at); + assert!(r.began_at <= r.finished_at); + assert_none!(r.execution_strategy); + } + assert_eq!(sl_subscribes[0].finished_status, "success"); + assert_eq!(sl_subscribes[1].finished_status, "canceled"); +} + +/// Test that we are sampling approximately 50% of statements. +/// Relies on two assumptions: +/// (1) that the effective sampling rate for the session is 50%, +/// (2) that we are using the deterministic testing RNG. +#[allow(clippy::disallowed_methods)] +fn test_statement_logging_sampling_inner( + server: TestServerWithStatementLoggingChecks, + mut client: postgres::Client, +) { + for i in 0..50 { + client.execute(&format!("SELECT {i}"), &[]).unwrap(); + + // Enforce a small delay to avoid duplicate `began_at` times, which would make the ordering + // of logged statements non-deterministic when we retrieve them below. + thread::sleep(Duration::from_millis(10)); + } + + // 23 randomly sampled out of 50 with 50% sampling. Seems legit! + let expected_sqls = [ + 2, 4, 5, 6, 9, 15, 17, 18, 19, 20, 21, 23, 24, 25, 31, 32, 33, 36, 37, 42, 46, + ] + .into_iter() + .map(|i| format!("SELECT {i}")) + .collect::>(); + + let mut internal_client = server.connect_internal(postgres::NoTls).unwrap(); + let seh_query = " + SELECT mst.sql + FROM mz_internal.mz_statement_execution_history AS mseh + JOIN + mz_internal.mz_prepared_statement_history AS mpsh + ON mseh.prepared_statement_id = mpsh.id + JOIN + mz_internal.mz_sql_text AS mst + ON mpsh.sql_hash = mst.sql_hash + WHERE mst.sql ~~ 'SELECT%' AND mst.sql !~~ '%mz_statement_execution_history%' + ORDER BY mseh.began_at ASC"; + + // Statement logging happens async, retry until we get the expected number of logged + // statements. + let mut sl = Vec::new(); + for _ in 0..10 { + thread::sleep(Duration::from_secs(1)); + sl = internal_client.query(seh_query, &[]).unwrap(); + if sl.len() >= expected_sqls.len() { + break; + } + } + + let sqls: Vec = sl.into_iter().map(|r| r.get(0)).collect(); + assert_eq!(sqls, expected_sqls); + + // Verify the statement_logging_record_count metric correctly tracks sampled vs unsampled. + // With 50% sampling and deterministic RNG, exactly 21 of 50 statements should be sampled. + let (sampled_true, sampled_false) = get_statement_logging_record_counts(&server); + assert_eq!( + sampled_true, 21, + "expected 21 statements to be sampled with 50% rate and deterministic RNG" + ); + assert_eq!( + sampled_false, 29, + "expected 29 statements to not be sampled with 50% rate and deterministic RNG" + ); +} + +#[mz_ore::test] +fn test_statement_logging_sampling() { + let (server, client) = setup_statement_logging(1.0, 0.5, ""); + test_statement_logging_sampling_inner(server, client); +} + +/// Test that we are not allowed to set `statement_logging_sample_rate` +/// arbitrarily high, but that it is constrained by `statement_logging_max_sample_rate`. +#[mz_ore::test] +fn test_statement_logging_sampling_constrained() { + let (server, client) = setup_statement_logging(0.5, 1.0, ""); + test_statement_logging_sampling_inner(server, client); +} + +/// Test that the `mz_statement_logging_unsampled_bytes` metric tracks the total bytes +/// of SQL text that would have been logged if statement logging were fully enabled. +/// We set `sample_rate=0.0` so no statements are actually sampled/logged, but the +/// unsampled_bytes metric still gets incremented for every executed statement. +#[mz_ore::test] +#[allow(clippy::disallowed_methods)] +fn test_statement_logging_unsampled_metrics() { + // Use sample_rate=0.0 so statements are not sampled, but unsampled_bytes metric is still tracked. + let (server, mut client) = setup_statement_logging(1.0, 0.0, ""); + + let batch_queries = [ + "SELECT 'Hello, world!';SELECT 1;;", + "SELECT 'Hello, world again!'", + ]; + let batch_total: usize = batch_queries + .iter() + .map(|s| s.as_bytes().iter().filter(|&&ch| ch != b';').count()) + .sum(); + let single_queries = ["SELECT 'foo'", "SELECT 'bar';;;"]; + let single_total: usize = single_queries + .iter() + .map(|s| s.as_bytes().iter().filter(|&&ch| ch != b';').count()) + .sum(); + let prepared_queries = ["SELECT 'baz';;;", "SELECT 'quux';"]; + let prepared_total: usize = prepared_queries + .iter() + .map(|s| s.as_bytes().iter().filter(|&&ch| ch != b';').count()) + .sum(); + + let named_prepared_inner = "SELECT 42"; + let named_prepared_outer = format!("PREPARE p AS {named_prepared_inner};EXECUTE p;"); + let named_prepared_outer_len = named_prepared_outer + .as_bytes() + .iter() + .filter(|&&ch| ch != b';') + .count(); + + for q in batch_queries { + client.batch_execute(q).unwrap(); + } + + for q in single_queries { + client.execute(q, &[]).unwrap(); + } + + for q in prepared_queries { + let s = client.prepare(q).unwrap(); + client.execute(&s, &[]).unwrap(); + } + + client.batch_execute(&named_prepared_outer).unwrap(); + + // This should NOT be logged, since we never actually execute it. + client.prepare("SELECT 'Hello, not counted!'").unwrap(); + + let expected_total = batch_total + single_total + prepared_total + named_prepared_outer_len; + let metric_value = server + .metrics_registry() + .gather() + .into_iter() + .find(|m| m.name() == "mz_statement_logging_unsampled_bytes") + .unwrap() + .take_metric()[0] + .get_counter() + .value(); + let metric_value = usize::cast_from(u64::try_cast_from(metric_value).unwrap()); + assert_eq!(expected_total, metric_value); + + // Also verify that statement_logging_record_count shows all statements as not sampled + // (since we're using 0% sample rate). + let (sampled_true, _sampled_false) = get_statement_logging_record_counts(&server); + assert_eq!( + sampled_true, 0, + "no statements should be sampled with 0% sample rate" + ); +} + +#[mz_ore::test] +#[allow(clippy::disallowed_methods)] +fn test_enable_internal_statement_logging() { + let (server, mut client) = setup_statement_logging_core( + 1.0, + 1.0, + "", + test_util::TestHarness::default().with_system_parameter_default( + "enable_internal_statement_logging".to_string(), + "true".to_string(), + ), + ); + + client.execute("SELECT 1", &[]).unwrap(); + + let mut client = server.connect_internal(postgres::NoTls).unwrap(); + let num_mz_system_statements = Retry::default() + .max_duration(Duration::from_secs(30)) + .retry(|_| { + let sl_results = client + .query( + "SELECT + count(*) +FROM mz_internal.mz_prepared_statement_history mpsh +JOIN mz_internal.mz_session_history USING (session_id) +WHERE authenticated_user='mz_system'", + &[], + ) + .unwrap(); + + let count: i64 = sl_results[0].get(0); + + if count > 0 { Ok(count) } else { Err(()) } + }) + .expect("at least some statements from mz_system should have been logged"); + + assert!( + num_mz_system_statements > 0, + "statements executed by mz_system should have been logged" + ); +} + +#[mz_ore::test] +// Test that statement logging +// doesn't cause a crash with subscribes over web sockets, +// which was previously happening (in staging) due to us +// dropping the `ExecuteContext` on the floor in that case. +fn test_statement_logging_ws_subscribe_no_crash() { + let (server, _client) = setup_statement_logging(1.0, 1.0, ""); + + // Create our WebSocket. + let ws_url = server.server.ws_addr(); + let (mut ws, _resp) = tungstenite::connect(ws_url).unwrap(); + test_util::auth_with_ws(&mut ws, Default::default()).unwrap(); + + let query = "SUBSCRIBE (SELECT 1)"; + let json = format!("{{\"query\":\"{query}\"}}"); + let json: serde_json::Value = serde_json::from_str(&json).unwrap(); + ws.send(Message::Text(json.to_string().into())).unwrap(); + + // Give the server time to crash, if it's going to. + std::thread::sleep(Duration::from_secs(1)) +} + +/// `finished_at` must record when execution finished, not when the coordinator +/// got around to the end event. A statement the session task retires itself +/// reports its own end timestamp, so a busy coordinator cannot inflate it. +/// +/// The coordinator is stalled inside `Catalog::transact` while the measured +/// statement runs. That failpoint sits ahead of every catalog mutation, so the +/// catalog revision does not move while it sleeps and the session's cached +/// snapshot stays valid, which is what keeps the measured statement from +/// needing the coordinator at all. +#[mz_ore::test] +#[allow(clippy::disallowed_methods)] +fn test_statement_logging_finished_at_excludes_coordinator_queue() { + let (server, mut client) = setup_statement_logging(1.0, 1.0, ""); + + // Populate this session's catalog snapshot cache, so the measured statement + // below needs nothing from the stalled coordinator. + client.execute("SELECT 1", &[]).unwrap(); + + let mut ddl = server.connect_internal(postgres::NoTls).unwrap(); + fail::cfg("catalog_transact", "sleep(3000)").unwrap(); + let stall = thread::spawn(move || { + let _ = ddl.batch_execute("CREATE TABLE stalls_the_coordinator (x int)"); + }); + // Give the DDL time to reach the failpoint before we measure. + thread::sleep(Duration::from_millis(500)); + + // Logged timestamps are epoch milliseconds and truncate, so floor the lower + // bound and ceil the upper one rather than comparing against the + // sub-millisecond instants `Utc::now` reports. + let began_bound = Utc::now().timestamp_millis(); + client + .execute("SELECT 1 AS finished_at_probe", &[]) + .unwrap(); + let finished_bound = Utc::now().timestamp_millis() + 1; + + stall.join().unwrap(); + fail::remove("catalog_transact"); + + let mut internal = server.connect_internal(postgres::NoTls).unwrap(); + let query = " + SELECT mseh.began_at, mseh.finished_at + FROM mz_internal.mz_statement_execution_history AS mseh + JOIN mz_internal.mz_prepared_statement_history AS mpsh + ON mseh.prepared_statement_id = mpsh.id + JOIN mz_internal.mz_sql_text AS mst ON mpsh.sql_hash = mst.sql_hash + WHERE mst.sql ~~ '%finished_at_probe%' AND mseh.finished_at IS NOT NULL"; + + let mut rows = Vec::new(); + for _ in 0..30 { + rows = internal.query(query, &[]).unwrap(); + if !rows.is_empty() { + break; + } + thread::sleep(Duration::from_secs(1)); + } + assert_eq!(rows.len(), 1, "expected exactly one logged probe statement"); + + let began_at: DateTime = rows[0].get(0); + let finished_at: DateTime = rows[0].get(1); + assert!( + began_at.timestamp_millis() >= began_bound, + "began_at precedes the client-observed start of the statement by {} ms", + began_bound - began_at.timestamp_millis() + ); + assert!( + finished_at.timestamp_millis() <= finished_bound, + "finished_at is {} ms past the client-observed end of the statement, so it \ + charges the statement for time its end event spent queued for the coordinator", + finished_at.timestamp_millis() - finished_bound + ); +} diff --git a/src/ore/src/iter.rs b/src/ore/src/iter.rs index 9ee9020741aed..9ac5c9f98926c 100644 --- a/src/ore/src/iter.rs +++ b/src/ore/src/iter.rs @@ -259,6 +259,7 @@ mod tests { #[crate::test] #[cfg(feature = "differential-dataflow")] + #[cfg_attr(miri, ignore)] // too slow fn test_consolidate_sorted() { proptest!(|(mut data in any::>())| { data.sort(); @@ -270,6 +271,7 @@ mod tests { #[crate::test] #[cfg(feature = "differential-dataflow")] + #[cfg_attr(miri, ignore)] // too slow fn test_consolidate() { proptest!(|(mut data in any::>())| { let streamed: Vec<_> = consolidate_iter(data.iter().copied()).collect(); @@ -279,6 +281,7 @@ mod tests { } #[crate::test] + #[cfg_attr(miri, ignore)] // too slow fn test_merge() { proptest!(|(mut data in vec(vec(0usize..100usize, 0..10), 0..10))| { let mut expected: Vec<_> = data.iter().flatten().copied().collect(); diff --git a/src/ore/src/pool.rs b/src/ore/src/pool.rs index 74185bf24cdda..23a1a5920e898 100644 --- a/src/ore/src/pool.rs +++ b/src/ore/src/pool.rs @@ -94,7 +94,9 @@ use crate::pool::region::{Region, SIZE_CLASSES}; /// (address space costs nothing, and touched pages are bounded by peak /// residency) so that no realistic budget, on any machine size, reaches the /// heap-fallback path. -const CLASS_CAPACITY_BYTES: usize = 1 << 40; +/// +/// NOTE: Seen OoMs with Miri since it actually allocates the capacity. +const CLASS_CAPACITY_BYTES: usize = if cfg!(miri) { 16 << 20 } else { 1 << 40 }; /// A chunk-provided transform between a chunk's body bytes and the stored /// bytes its extent holds. The pool owns scheduling: spill threads, the @@ -2987,6 +2989,7 @@ mod tests { /// the same victims (the pool's only two-chunk lock edge). Contents are /// asserted on every read. #[mz_ore::test] + #[cfg_attr(miri, ignore)] // too slow fn concurrent_admits_exercise_the_steal_path() { const CHUNKS: u64 = 8; let pool = test_pool(usize::MAX); @@ -3069,6 +3072,7 @@ mod tests { /// contents stay correct, contended steals degrade to skips, and the /// accounting identity settles to zero. #[mz_ore::test] + #[cfg_attr(miri, ignore)] // too slow fn concurrent_admits_race_cleanly() { let pool = test_pool(64 << 10); let per_thread = rounds(50, 3); @@ -3364,6 +3368,7 @@ mod tests { } #[mz_ore::test] + #[cfg_attr(miri, ignore)] // too slow fn multithreaded_smoke() { // Budget of one small chunk: four inserting threads keep the pool // over budget, so every insert's enforcement pass selects victims @@ -3399,6 +3404,7 @@ mod tests { } #[mz_ore::test] + #[cfg_attr(miri, ignore)] // too slow fn concurrent_read_enforce_churn() { // Races the three actors that can touch one chunk's slot: readers // copying shared chunks out and verifying them, an enforcer evicting @@ -3480,6 +3486,7 @@ mod tests { /// drops) extent-queue entries, and pruning alone must keep the queue /// proportional to the live resident extents. #[mz_ore::test] + #[cfg_attr(miri, ignore)] // too slow fn extent_queue_stays_bounded_under_cap() { let pool = test_pool(256 << 20); pool.set_rss_target(1 << 40); @@ -3636,6 +3643,7 @@ mod tests { } #[mz_ore::test] + #[cfg_attr(miri, ignore)] // too slow fn spill_threads_end_to_end() { let pool = test_pool(128 << 10); pool.set_spill_threads(2); diff --git a/src/pgwire-common/src/codec.rs b/src/pgwire-common/src/codec.rs index 84fac7a1ea3c2..04a9763a322d7 100644 --- a/src/pgwire-common/src/codec.rs +++ b/src/pgwire-common/src/codec.rs @@ -52,6 +52,7 @@ impl fmt::Display for CodecError { pub trait Pgbuf: BufMut { fn put_string(&mut self, s: &str); fn put_length_i16(&mut self, len: usize) -> Result<(), io::Error>; + fn put_length_u16(&mut self, len: usize) -> Result<(), io::Error>; fn put_format_i8(&mut self, format: Format); fn put_format_i16(&mut self, format: Format); } @@ -70,6 +71,17 @@ impl Pgbuf for B { Ok(()) } + /// Writes a count field as unsigned, so it may exceed 32767. The protocol + /// calls these fields `Int16`, but PostgreSQL clients decode them as + /// unsigned. + fn put_length_u16(&mut self, len: usize) -> Result<(), io::Error> { + let len = u16::try_from(len).map_err(|_| { + io::Error::new(io::ErrorKind::InvalidData, "length does not fit in a u16") + })?; + self.put_u16(len); + Ok(()) + } + fn put_format_i8(&mut self, format: Format) { self.put_i8(format.into()) } @@ -285,17 +297,6 @@ impl<'a> Cursor<'a> { } } - /// Reads the next 16-bit signed integer, advancing the cursor by two - /// bytes. - pub fn read_i16(&mut self) -> Result { - if self.buf.len() < 2 { - return Err(input_err("not enough buffer for an Int16")); - } - let val = NetworkEndian::read_i16(self.buf); - self.advance(2); - Ok(val) - } - /// Reads the next 32-bit signed integer, advancing the cursor by four /// bytes. pub fn read_i32(&mut self) -> Result { @@ -307,6 +308,17 @@ impl<'a> Cursor<'a> { Ok(val) } + /// Reads the next 16-bit unsigned integer, advancing the cursor by two + /// bytes. + pub fn read_u16(&mut self) -> Result { + if self.buf.len() < 2 { + return Err(input_err("not enough buffer for an Int16")); + } + let val = NetworkEndian::read_u16(self.buf); + self.advance(2); + Ok(val) + } + /// Reads the next 32-bit unsigned integer, advancing the cursor by four /// bytes. pub fn read_u32(&mut self) -> Result { @@ -320,7 +332,7 @@ impl<'a> Cursor<'a> { /// Reads the next 16-bit format code, advancing the cursor by two bytes. pub fn read_format(&mut self) -> Result { - Format::try_from(self.read_i16()?) + Format::try_from(self.read_u16()?) } /// Advances the cursor by `n` bytes. diff --git a/src/pgwire-common/src/format.rs b/src/pgwire-common/src/format.rs index 0502c3829c617..ebc72ee20dbef 100644 --- a/src/pgwire-common/src/format.rs +++ b/src/pgwire-common/src/format.rs @@ -26,10 +26,10 @@ pub enum Format { Binary, } -impl TryFrom for Format { +impl TryFrom for Format { type Error = io::Error; - fn try_from(value: i16) -> Result { + fn try_from(value: u16) -> Result { match value { 0 => Ok(Format::Text), 1 => Ok(Format::Binary), diff --git a/src/pgwire/src/codec.rs b/src/pgwire/src/codec.rs index 91311c07986ab..e54583bb0e4fb 100644 --- a/src/pgwire/src/codec.rs +++ b/src/pgwire/src/codec.rs @@ -430,17 +430,17 @@ impl Codec { dst.put_u32(secret_key); } BackendMessage::ParameterDescription(params) => { - if params.len() > usize::try_from(i16::MAX).expect("i16::MAX is positive") { + if params.len() > usize::from(u16::MAX) { return Err(io::Error::new( io::ErrorKind::InvalidData, format!( "{} params in parameter description, which exceeds {}", params.len(), - i16::MAX + u16::MAX ), )); } - dst.put_length_i16(params.len())?; + dst.put_length_u16(params.len())?; for param in params { dst.put_u32(param.oid()); } @@ -792,8 +792,12 @@ fn decode_parse(mut buf: Cursor) -> Result { let name = buf.read_cstr()?; let sql = buf.read_cstr()?; + // NOTE: the protocol calls the counts that precede repeated groups `Int16`, + // but PostgreSQL decodes them as unsigned, so they reach 65535. Reading one + // as signed makes it negative, which drops the group and leaves the cursor + // misaligned for the rest of the message. let mut param_types = vec![]; - for _ in 0..buf.read_i16()? { + for _ in 0..buf.read_u16()? { param_types.push(buf.read_u32()?); } @@ -833,13 +837,15 @@ fn decode_bind(mut buf: Cursor) -> Result { let portal_name = buf.read_cstr()?.to_string(); let statement_name = buf.read_cstr()?.to_string(); + // The three counts below are `Int16` in the protocol but decoded as + // unsigned, for the reason spelled out in `decode_parse`. let mut param_formats = Vec::new(); - for _ in 0..buf.read_i16()? { + for _ in 0..buf.read_u16()? { param_formats.push(buf.read_format()?); } let mut raw_params = Vec::new(); - for _ in 0..buf.read_i16()? { + for _ in 0..buf.read_u16()? { let len = buf.read_i32()?; if len == -1 { raw_params.push(None); // NULL @@ -854,7 +860,7 @@ fn decode_bind(mut buf: Cursor) -> Result { } let mut result_formats = Vec::new(); - for _ in 0..buf.read_i16()? { + for _ in 0..buf.read_u16()? { result_formats.push(buf.read_format()?); } diff --git a/src/repr/src/adt/datetime.rs b/src/repr/src/adt/datetime.rs index 64eb838384536..c1bf1a93b3078 100644 --- a/src/repr/src/adt/datetime.rs +++ b/src/repr/src/adt/datetime.rs @@ -3667,6 +3667,7 @@ mod tests { } #[mz_ore::test] + #[cfg_attr(miri, ignore)] // too slow fn proptest_packed_naive_time_sort_order() { let time = add_arb_duration(NaiveTime::from_hms_opt(0, 0, 0).unwrap()); let strat = proptest::collection::vec(time, 0..128); diff --git a/src/repr/src/adt/interval.rs b/src/repr/src/adt/interval.rs index a61d64b1253e5..eb5a1dbc32c97 100644 --- a/src/repr/src/adt/interval.rs +++ b/src/repr/src/adt/interval.rs @@ -1325,6 +1325,7 @@ mod test { } #[mz_ore::test] + #[cfg_attr(miri, ignore)] // too slow fn proptest_packed_interval_sorts() { fn sort_intervals(mut og: Vec) { let mut packed: Vec<_> = og.iter().copied().map(PackedInterval::from_value).collect(); diff --git a/src/repr/src/adt/timestamp.rs b/src/repr/src/adt/timestamp.rs index 019d657567df2..0a47c25c29a4a 100644 --- a/src/repr/src/adt/timestamp.rs +++ b/src/repr/src/adt/timestamp.rs @@ -1220,6 +1220,7 @@ mod test { } #[mz_ore::test] + #[cfg_attr(miri, ignore)] // too slow fn proptest_packed_naive_date_time_sort_order() { let strat = proptest::collection::vec(arb_naive_date_time(), 0..128); proptest!(|(mut times in strat)| { diff --git a/src/sql-parser/src/ast/defs/name.rs b/src/sql-parser/src/ast/defs/name.rs index f732b52ff8a06..c6559ebdfcc4d 100644 --- a/src/sql-parser/src/ast/defs/name.rs +++ b/src/sql-parser/src/ast/defs/name.rs @@ -333,9 +333,12 @@ impl Ident { // (`list[1]` is a valid one-element list), so a bare `list` // identifier that gets subscripted — `"list"[1]` — would // reparse as a list literal instead of a subscript. (`ARRAY` - // is reserved-in-scalar-expression and so already quoted; - // `MAP[...]` requires `=>`, so `map[1]` is unambiguously a - // subscript.) + // is reserved-in-scalar-expression and so already quoted. + // `MAP` needs no clause here even though `map[…]` is a + // literal too: `MAP` is a context-sensitive keyword, so + // `write_subscript_receiver` parenthesizes the receiver into + // `(map)[1]`, and every grammar that starts a map literal + // requires a `[` right after the keyword.) || kw == LIST // `DEALLOCATE [PREPARE] ` accepts an optional // `PREPARE` keyword before the name, so a bare `prepare` diff --git a/src/sql-parser/src/parser.rs b/src/sql-parser/src/parser.rs index e0c2a39f779eb..34763fdb55465 100644 --- a/src/sql-parser/src/parser.rs +++ b/src/sql-parser/src/parser.rs @@ -5721,6 +5721,15 @@ impl<'a> Parser<'a> { fn parse_option_map( &mut self, ) -> Result>>, ParserError> { + // `MAP` only begins a map literal when a `[` follows it. Committing on + // the keyword alone makes a bare `map` in an option-value position a hard + // error on the missing bracket, rather than falling through to the item + // name it is: `TOPIC CONFIG = "map"` prints as `TOPIC CONFIG = map`, + // since `map` is a legal bare identifier everywhere else, and that output + // then failed to reparse. + if !(self.peek_keyword(MAP) && self.peek_nth_token(1) == Some(Token::LBracket)) { + return Ok(None); + } Ok(if self.parse_keyword(MAP) { self.expect_token(&Token::LBracket)?; let mut map = BTreeMap::new(); diff --git a/src/sql-parser/tests/sqlparser_common.rs b/src/sql-parser/tests/sqlparser_common.rs index 1e67fb0c313fc..27326a43fb58d 100644 --- a/src/sql-parser/tests/sqlparser_common.rs +++ b/src/sql-parser/tests/sqlparser_common.rs @@ -878,6 +878,29 @@ fn test_list_keyword_bare_identifier_subscript_display_roundtrip() { } } +#[mz_ore::test] +#[cfg_attr(miri, ignore)] // unsupported operation: can't call foreign function `rust_psm_stack_pointer` on OS `linux` +fn test_map_keyword_bare_identifier_option_value_display_roundtrip() { + // A `MAP` that no `[` follows is the item name `map`, which is how a quoted + // `"map"` prints. `parse_option_map` used to commit on the keyword alone, so + // the printed form failed to reparse in every generic option-value position. + // Regression for the sql_roundtrip fuzz finding `(TOPIC = "map")`. + for sql in [ + r#"CREATE SINK s FROM t INTO KAFKA CONNECTION c (TOPIC = "map") FORMAT BYTES ENVELOPE DEBEZIUM"#, + r#"CREATE SOURCE s FROM KAFKA CONNECTION c (TOPIC = "map") FORMAT BYTES"#, + r#"CREATE MATERIALIZED VIEW v WITH (PARTITION BY = "map") AS SELECT 1"#, + // A real map-literal option value still takes the map branch. + r#"CREATE SINK s FROM t INTO KAFKA CONNECTION c (TOPIC = 't', TOPIC CONFIG = MAP['a' => 'b']) FORMAT BYTES ENVELOPE DEBEZIUM"#, + // Expression position needs no quoting: a subscripted receiver prints as + // `(map)[1]`, which is why `can_be_printed_bare` has no `MAP` clause. + r#"SELECT "map""#, + r#"SELECT "map"[1]"#, + r#"SELECT map['a' => 1]"#, + ] { + assert_display_roundtrips(sql); + } +} + #[mz_ore::test] #[cfg_attr(miri, ignore)] // unsupported operation: can't call foreign function `rust_psm_stack_pointer` on OS `linux` fn test_table_function_special_name_display_roundtrip() { diff --git a/src/sql-parser/tests/testdata/ddl b/src/sql-parser/tests/testdata/ddl index 62597942f6bf5..6ee7e4a46002d 100644 --- a/src/sql-parser/tests/testdata/ddl +++ b/src/sql-parser/tests/testdata/ddl @@ -3781,3 +3781,13 @@ CREATE SINK snk FROM t INTO KAFKA CONNECTION conn1 (TOPIC 'topic') FORMAT AVRO U CREATE SINK snk FROM t INTO KAFKA CONNECTION conn1 (TOPIC = 'topic') FORMAT AVRO USING AWS GLUE SCHEMA REGISTRY CONNECTION glue_conn (KEY SCHEMA NAME = 'k', VALUE SCHEMA NAME = 'v', KEY COMPATIBILITY LEVEL = 'BACKWARD', VALUE COMPATIBILITY LEVEL = 'FULL') ENVELOPE UPSERT => CreateSink(CreateSinkStatement { name: Some(UnresolvedItemName([Ident("snk")])), in_cluster: None, if_not_exists: false, from: Name(UnresolvedItemName([Ident("t")])), connection: Kafka { connection: Name(UnresolvedItemName([Ident("conn1")])), options: [KafkaSinkConfigOption { name: Topic, value: Some(Value(String("topic"))) }], key: None, headers: None }, format: Some(Bare(Avro(Glue { connection: Name(UnresolvedItemName([Ident("glue_conn")])), with_options: [GlueAvroOption { name: KeySchemaName, value: Some(Value(String("k"))) }, GlueAvroOption { name: ValueSchemaName, value: Some(Value(String("v"))) }, GlueAvroOption { name: KeyCompatibilityLevel, value: Some(Value(String("BACKWARD"))) }, GlueAvroOption { name: ValueCompatibilityLevel, value: Some(Value(String("FULL"))) }], seed: None }))), envelope: Some(Upsert), mode: None, with_options: [] }) + +# An option value that is the bare identifier `map`. `MAP` only starts a map +# literal when a `[` follows, so this is the item name `map`, which is what +# `AstDisplay` prints for a quoted `"map"`. +parse-statement +CREATE SINK foo FROM bar INTO KAFKA CONNECTION baz (TOPIC CONFIG = "map") +---- +CREATE SINK foo FROM bar INTO KAFKA CONNECTION baz (TOPIC CONFIG = map) +=> +CreateSink(CreateSinkStatement { name: Some(UnresolvedItemName([Ident("foo")])), in_cluster: None, if_not_exists: false, from: Name(UnresolvedItemName([Ident("bar")])), connection: Kafka { connection: Name(UnresolvedItemName([Ident("baz")])), options: [KafkaSinkConfigOption { name: TopicConfig, value: Some(UnresolvedItemName(UnresolvedItemName([Ident("map")]))) }], key: None, headers: None }, format: None, envelope: None, mode: None, with_options: [] }) diff --git a/src/sql/src/session/vars.rs b/src/sql/src/session/vars.rs index 485168d6593d8..1cdc5dd702056 100644 --- a/src/sql/src/session/vars.rs +++ b/src/sql/src/session/vars.rs @@ -2053,6 +2053,10 @@ impl SystemVars { *self.expect_value(&SCRAM_ITERATIONS) } + /// Computes an update for every dyncfg from its configured value. + /// + /// This does not touch our own [`ConfigSet`]. Callers that need this + /// process's dyncfgs to reflect the result want [`Self::sync_dyncfgs`]. pub fn dyncfg_updates(&self) -> ConfigUpdates { let mut updates = ConfigUpdates::default(); for entry in self.dyncfgs.entries() { @@ -2080,6 +2084,18 @@ impl SystemVars { }; updates.add_dynamic(entry.name(), val); } + updates + } + + /// Applies the configured dyncfg values to our own [`ConfigSet`], and + /// returns them. + /// + /// Two callers own keeping this process's dyncfgs in step with the + /// catalog: catalog open, and every durable system-config change. Everyone + /// else only forwards the updates elsewhere and wants + /// [`Self::dyncfg_updates`] instead. + pub fn sync_dyncfgs(&self) -> ConfigUpdates { + let updates = self.dyncfg_updates(); updates.apply(&self.dyncfgs); updates } diff --git a/src/timely-util/src/columnar.rs b/src/timely-util/src/columnar.rs index e66c66b07198e..e8d9134ab03f1 100644 --- a/src/timely-util/src/columnar.rs +++ b/src/timely-util/src/columnar.rs @@ -373,6 +373,7 @@ mod tests { /// The ship signal is monotone: once it fires it stays fired, even when /// a single wide record steps far past the 2 MiB boundary in one push. #[mz_ore::test] + #[cfg_attr(miri, ignore)] // too slow fn ship_threshold_monotone() { use columnar::Push; let mut container = as Columnar>::Container::default(); diff --git a/src/timely-util/src/columnar/chunk.rs b/src/timely-util/src/columnar/chunk.rs index 002c0f012c072..4d693af116b86 100644 --- a/src/timely-util/src/columnar/chunk.rs +++ b/src/timely-util/src/columnar/chunk.rs @@ -1459,6 +1459,7 @@ mod tests { /// A body large enough to spill round-trips through the pool with resident /// metadata intact, and the batcher produces spilled sealed output. #[mz_ore::test] + #[cfg_attr(miri, ignore)] // too slow fn spill_round_trip() { set_spill_override(Some(test_pool())); @@ -1490,6 +1491,7 @@ mod tests { /// Merging spilled chains loads bodies call-scoped and consolidates /// correctly, and an untouched survivor keeps its spilled body. #[mz_ore::test] + #[cfg_attr(miri, ignore)] // too slow fn merge_spilled_chains() { set_spill_override(Some(test_pool())); @@ -1604,6 +1606,7 @@ mod tests { /// commit large enough to spill carries the depth into its spilled /// metadata (and thus into the pool hints). #[mz_ore::test] + #[cfg_attr(miri, ignore)] // too slow fn settle_commits_at_accumulated_depth() { set_spill_override(Some(test_pool())); let big: Vec = (0..100_000u64).map(|i| ((i, 0), 0, 1i64)).collect(); @@ -1627,6 +1630,7 @@ mod tests { /// periodic ship window, so mid-window chunk sizes cannot make it grow /// past the target unbounded. #[mz_ore::test] + #[cfg_attr(miri, ignore)] // too slow fn settle_carry_commits_at_target() { // ~1.5 MiB per chunk: inside the dead zone of the periodic window // check (see `at_commit_size`).