diff --git a/misc/python/materialize/parallel_workload/action.py b/misc/python/materialize/parallel_workload/action.py index e6d36e6748c87..6ce4fb8b7bc5c 100644 --- a/misc/python/materialize/parallel_workload/action.py +++ b/misc/python/materialize/parallel_workload/action.py @@ -24,7 +24,7 @@ from psycopg import Connection from psycopg.errors import OperationalError -import materialize.parallel_workload.database +import materialize.parallel_workload.column from materialize.data_ingest.data_type import ( NUMBER_TYPES, Boolean, @@ -221,7 +221,7 @@ def errors_to_ignore(self, exe: Executor) -> list[str]: ) if exe.db.scenario == Scenario.Rename: result.extend(["unknown schema", "ambiguous reference to schema name"]) - if materialize.parallel_workload.database.NAUGHTY_IDENTIFIERS: + if materialize.parallel_workload.column.NAUGHTY_IDENTIFIERS: result.extend(["identifier length exceeds 255 bytes"]) return result diff --git a/misc/python/materialize/parallel_workload/column.py b/misc/python/materialize/parallel_workload/column.py index 8731efdbe1b76..fa63c08067bcf 100644 --- a/misc/python/materialize/parallel_workload/column.py +++ b/misc/python/materialize/parallel_workload/column.py @@ -24,6 +24,11 @@ NAUGHTY_IDENTIFIERS = False +def set_naughty_identifiers(value: bool) -> None: + global NAUGHTY_IDENTIFIERS + NAUGHTY_IDENTIFIERS = value + + def naughtify(name: str) -> str: """Makes a string into a naughty identifier, always returns the same identifier when called with the same input.""" diff --git a/misc/python/materialize/parallel_workload/database.py b/misc/python/materialize/parallel_workload/database.py index fe5154f1fc6d7..46a44696fe65a 100644 --- a/misc/python/materialize/parallel_workload/database.py +++ b/misc/python/materialize/parallel_workload/database.py @@ -49,6 +49,7 @@ SqlServerColumn, WebhookColumn, naughtify, + set_naughty_identifiers, ) from materialize.parallel_workload.executor import Executor from materialize.parallel_workload.expression import ExprKind, expression @@ -872,13 +873,12 @@ def __init__( scenario: Scenario, naughty_identifiers: bool, ): - global NAUGHTY_IDENTIFIERS self.host = host self.ports = ports self.complexity = complexity self.scenario = scenario self.seed = seed - NAUGHTY_IDENTIFIERS = naughty_identifiers + set_naughty_identifiers(naughty_identifiers) self.s3_path = 0 self.dbs = [DB(seed, i) for i in range(rng.randint(1, MAX_INITIAL_DBS))]