From 02f9c086df8805ea5198fdc1ba01744dfb86fb65 Mon Sep 17 00:00:00 2001 From: Dennis Felsing Date: Wed, 8 Jul 2026 21:26:42 +0000 Subject: [PATCH] data-ingest: Fix data generation and executor issues - Text.random_value: the special-values branch (NULL, unicode, numbers as strings) was dead code, always overwritten by the record-size branch - Insert: restart keys on each generate() call so cycling workloads (DeleteDataAtEndOfDay*) insert again after the first cycle instead of degenerating into no-op deletes - DeleteDataAtEndOfDay*: the delete phase now actually runs as a single transaction, as its comment claims - Upsert: honor the keyspace parameter, reject the unimplemented EXISTING - Replace the type blocklists with explicit allowlists and drop types from DATA_TYPES_FOR_MYSQL that cannot work: MzTimestamp and TimestampTz have no MySQL type, Timestamp and Date values exceed MySQL's ranges - Seed the row-value rng so --seed makes runs fully reproducible - Sort WORKLOADS so rng.choice in parallel-workload is deterministic - PrintExecutor: don't connect to Mz, --verbose crashed when the randomly picked materialized2 wasn't up yet - Executor.execute: only retry OperationalErrors against Mz, never redirect an upstream Postgres query to the Mz connection Co-Authored-By: Claude Fable 5 --- .../materialize/data_ingest/data_type.py | 128 +++--------------- .../materialize/data_ingest/definition.py | 35 ++--- .../materialize/data_ingest/executor.py | 15 +- .../materialize/data_ingest/workload.py | 18 ++- test/data-ingest/mzcompose.py | 4 + 5 files changed, 59 insertions(+), 141 deletions(-) diff --git a/misc/python/materialize/data_ingest/data_type.py b/misc/python/materialize/data_ingest/data_type.py index 7e96d514df230..741f09ad6b3ef 100644 --- a/misc/python/materialize/data_ingest/data_type.py +++ b/misc/python/materialize/data_ingest/data_type.py @@ -479,6 +479,7 @@ def random_value( rng.randint(-100, 100), ] ) + return literal(str(result)) if in_query else str(result) # Fails: unterminated dollar-quoted string # chars = string.printable chars = string.ascii_letters + string.digits @@ -1190,120 +1191,23 @@ def name(backend: Backend = Backend.MATERIALIZE) -> str: # Sort to keep determinism for reproducible runs with specific seed DATA_TYPES = sorted(list(all_subclasses(DataType)), key=repr) -# fastavro._schema_common.UnknownType: record -# bytea requires Python bytes type instead of str -DATA_TYPES_FOR_AVRO = sorted( - list( - set(DATA_TYPES) - - { - TextTextMap, - Jsonb, - Bytea, - Boolean, - UUID, - Interval, - IntList, - IntArray, - Time, - Date, - Timestamp, - TimestampTz, - MzTimestamp, - Oid, - Numeric, - Numeric383, - UInt2, - UInt4, - UInt8, - Float, - Double, - Char, - VarChar, - Int4Range, - Int8Range, - NumRange, - DateRange, - TsRange, - TsTzRange, - } - ), - key=repr, -) - -DATA_TYPES_FOR_MYSQL = sorted( - list( - set(DATA_TYPES) - - { - IntList, - IntArray, - UUID, - TextTextMap, - Interval, - Oid, - Jsonb, - Bytea, - Boolean, - Numeric, - Numeric383, - UInt2, - UInt4, - UInt8, - Char, - VarChar, - Int4Range, - Int8Range, - NumRange, - DateRange, - TsRange, - TsTzRange, - } - ), - key=repr, -) - -DATA_TYPES_FOR_SQL_SERVER = sorted( - list( - set(DATA_TYPES) - - { - IntList, - IntArray, - UUID, - TextTextMap, - Interval, - Oid, - Jsonb, - Bytea, - Boolean, - Numeric, - Numeric383, - UInt2, - UInt4, - UInt8, - Date, - Time, - Timestamp, - TimestampTz, - MzTimestamp, - Float, - Double, - Char, - VarChar, - Int4Range, - Int8Range, - NumRange, - DateRange, - TsRange, - TsTzRange, - } - ), - key=repr, -) +# Explicit allowlists so that the actually exercised types are visible at a +# glance. A type belongs in one of these lists only if it maps to a native +# type on the backend, generated values fit the backend's ranges, and values +# roundtrip identically for the cross-backend result comparison. + +# Types the Kafka executor can encode in Avro directly, no records, bytes or +# logical types. +DATA_TYPES_FOR_AVRO = [Int, Long, SmallInt, Text] + +# Timestamp and Date are excluded because random_value generates years outside +# MySQL's supported ranges. MzTimestamp and TimestampTz have no MySQL type. +DATA_TYPES_FOR_MYSQL = [Double, Float, Int, Long, SmallInt, Text, Time] + +DATA_TYPES_FOR_SQL_SERVER = [Int, Long, SmallInt, Text] # MySQL doesn't support keys of unlimited size -DATA_TYPES_FOR_KEY = sorted( - list(set(DATA_TYPES_FOR_AVRO) - {Text, Bytea, IntList, IntArray, Float, Double}), - key=repr, -) +DATA_TYPES_FOR_KEY = [Int, Long, SmallInt] NUMBER_TYPES = [ SmallInt, diff --git a/misc/python/materialize/data_ingest/definition.py b/misc/python/materialize/data_ingest/definition.py index f9a53b6f11415..857ae52e97d14 100644 --- a/misc/python/materialize/data_ingest/definition.py +++ b/misc/python/materialize/data_ingest/definition.py @@ -33,12 +33,6 @@ class Keyspace(Enum): EXISTING = 3 -class Target(Enum): - KAFKA = 1 - POSTGRES = 2 - PRINT = 3 - - class Definition: def generate(self, fields: list[Field]) -> Iterator[RowList]: raise NotImplementedError @@ -48,7 +42,6 @@ class Insert(Definition): def __init__(self, count: Records, record_size: RecordSize): self.count = count.value self.record_size = record_size - self.current_key = 0 def max_key(self) -> int: if self.count < 1: @@ -63,19 +56,17 @@ def generate(self, fields: list[Field]) -> Iterator[RowList]: f'Unexpected count {self.count}, doesn\'t make sense to generate "ALL" values' ) - for i in range(self.count): - if self.current_key >= self.count: - break - + # Keys restart at 0 on every call so that a workload cycling through + # insert and delete phases inserts the same keys again in each cycle. + for key in range(self.count): values = [ ( - field.data_type.numeric_value(self.current_key) + field.data_type.numeric_value(key) if field.is_key else field.data_type.random_value(rng, self.record_size) ) for field in fields ] - self.current_key += 1 yield RowList( [ @@ -90,6 +81,8 @@ def generate(self, fields: list[Field]) -> Iterator[RowList]: class Upsert(Definition): def __init__(self, keyspace: Keyspace, count: Records, record_size: RecordSize): + if keyspace not in (Keyspace.SINGLE_VALUE, Keyspace.LARGE): + raise ValueError(f"Unsupported keyspace {keyspace}") self.keyspace = keyspace self.count = count.value self.record_size = record_size @@ -101,14 +94,14 @@ def generate(self, fields: list[Field]) -> Iterator[RowList]: ) for i in range(self.count): - values = [ - ( - field.data_type.numeric_value(0) - if field.is_key - else field.data_type.random_value(rng, self.record_size) - ) - for field in fields - ] + values = [] + for field in fields: + if not field.is_key: + values.append(field.data_type.random_value(rng, self.record_size)) + elif self.keyspace == Keyspace.SINGLE_VALUE: + values.append(field.data_type.numeric_value(0)) + else: + values.append(field.data_type.random_value(rng, self.record_size)) yield RowList( [ diff --git a/misc/python/materialize/data_ingest/executor.py b/misc/python/materialize/data_ingest/executor.py index f9b025df43b6c..219243b80c813 100644 --- a/misc/python/materialize/data_ingest/executor.py +++ b/misc/python/materialize/data_ingest/executor.py @@ -104,7 +104,14 @@ def execute(self, cur: psycopg.Cursor | pymysql.cursors.Cursor, query: str) -> N else cur.execute(query) ) except OperationalError: - # Can happen after Mz disruptions if we are running queries against Mz + # Can happen after Mz disruptions if we are running queries against + # Mz. Only queries against Mz can be retried this way, a query on + # an upstream connection must not be redirected to Mz. + if ( + not isinstance(cur, psycopg.Cursor) + or cur.connection is not self.mz_conn + ): + raise print("Network error, retrying") time.sleep(0.01) self.reconnect() @@ -136,6 +143,12 @@ def execute_with_retry_on_error( class PrintExecutor(Executor): + def reconnect(self) -> None: + # Only prints transactions, needs no Materialize connection. The + # random service pick in the base implementation could also hit a + # service that is not running yet. + pass + def create(self, logging_exe: Any | None = None) -> None: pass diff --git a/misc/python/materialize/data_ingest/workload.py b/misc/python/materialize/data_ingest/workload.py index f9cf4582744a0..894eb294305f8 100644 --- a/misc/python/materialize/data_ingest/workload.py +++ b/misc/python/materialize/data_ingest/workload.py @@ -162,13 +162,14 @@ def __init__( ) # Delete all records in a single transaction delete_phase = TransactionDef( - [ + size=TransactionSize.HUGE, + operations=[ Delete( number_of_records=Records.ALL, record_size=RecordSize.SMALL, num=insert.max_key(), ) - ] + ], ) self.cycle = [ insert_phase, @@ -195,13 +196,14 @@ def __init__( ) # Delete all records in a single transaction delete_phase = TransactionDef( - [ + size=TransactionSize.HUGE, + operations=[ Delete( number_of_records=Records.ALL, record_size=RecordSize.SMALL, num=insert.max_key(), ) - ] + ], ) self.cycle = [ insert_phase, @@ -235,13 +237,14 @@ def __init__( ) # Delete all records in a single transaction delete_phase = TransactionDef( - [ + size=TransactionSize.HUGE, + operations=[ Delete( number_of_records=Records.ALL, record_size=RecordSize.SMALL, num=insert.max_key(), ) - ] + ], ) self.cycle = [ insert_phase, @@ -266,7 +269,8 @@ def __init__( # ] -WORKLOADS = all_subclasses(Workload) +# Sort to keep determinism for reproducible runs with specific seed +WORKLOADS = sorted(all_subclasses(Workload), key=repr) def execute_workload( diff --git a/test/data-ingest/mzcompose.py b/test/data-ingest/mzcompose.py index 72ca8fea5fa2a..a0fc19f6081f5 100644 --- a/test/data-ingest/mzcompose.py +++ b/test/data-ingest/mzcompose.py @@ -16,6 +16,7 @@ import time from materialize import buildkite +from materialize.data_ingest import definition from materialize.data_ingest.executor import ( KafkaExecutor, MySqlExecutor, @@ -191,6 +192,9 @@ def workflow_default(c: Composition, parser: WorkflowArgumentParser) -> None: for i, workload_class in enumerate(workloads): random.seed(args.seed) + # Row values come from a separate rng, seed it too so that runs + # with the same seed are fully reproducible. + definition.rng.seed(args.seed) print(f"--- Testing workload {workload_class.__name__}") workload = workload_class(args.azurite, c, mz_service, deploy_generation) execute_workload(