misc.python.materialize.zippy.table_actions

  1# Copyright Materialize, Inc. and contributors. All rights reserved.
  2#
  3# Use of this software is governed by the Business Source License
  4# included in the LICENSE file at the root of this repository.
  5#
  6# As of the Change Date specified in that file, in accordance with
  7# the Business Source License, use of this software will be governed
  8# by the Apache License, Version 2.0.
  9
 10import random
 11from textwrap import dedent
 12
 13from materialize.mzcompose.composition import Composition
 14from materialize.zippy.balancerd_capabilities import BalancerdIsRunning
 15from materialize.zippy.framework import (
 16    Action,
 17    ActionFactory,
 18    Capabilities,
 19    Capability,
 20    State,
 21)
 22from materialize.zippy.mz_capabilities import MzIsRunning
 23from materialize.zippy.table_capabilities import TableExists
 24from materialize.zippy.view_capabilities import ViewExists
 25
 26MAX_ROWS_PER_ACTION = 10000
 27
 28
 29class CreateTableParameterized(ActionFactory):
 30    def __init__(
 31        self, max_tables: int = 10, max_rows_per_action: int = MAX_ROWS_PER_ACTION
 32    ) -> None:
 33        self.max_tables = max_tables
 34        self.max_rows_per_action = max_rows_per_action
 35
 36    @classmethod
 37    def requires(cls) -> set[type[Capability]]:
 38        return {BalancerdIsRunning, MzIsRunning}
 39
 40    def new(self, capabilities: Capabilities) -> list[Action]:
 41        new_table_name = capabilities.get_free_capability_name(
 42            TableExists, self.max_tables
 43        )
 44
 45        if new_table_name:
 46            return [
 47                CreateTable(
 48                    capabilities=capabilities,
 49                    table=TableExists(
 50                        name=new_table_name,
 51                        has_index=random.choice([True, False]),
 52                        max_rows_per_action=self.max_rows_per_action,
 53                    ),
 54                )
 55            ]
 56        else:
 57            return []
 58
 59
 60class CreateTable(Action):
 61    """Creates a table on the Mz instance. 50% of the tables have a default index."""
 62
 63    @classmethod
 64    def requires(cls) -> set[type[Capability]]:
 65        return {BalancerdIsRunning, MzIsRunning}
 66
 67    def __init__(self, table: TableExists, capabilities: Capabilities) -> None:
 68        assert (
 69            table is not None
 70        ), "CreateTable Action can not be referenced directly, it is produced by CreateTableParameterized factory"
 71        self.table = table
 72        super().__init__(capabilities)
 73
 74    def run(self, c: Composition, state: State) -> None:
 75        index = (
 76            f"> CREATE DEFAULT INDEX ON {self.table.name}"
 77            if self.table.has_index
 78            else ""
 79        )
 80        c.testdrive(
 81            dedent(f"""
 82                > CREATE TABLE {self.table.name} (f1 INTEGER);
 83                {index}
 84                > INSERT INTO {self.table.name} VALUES ({self.table.watermarks.max});
 85                """),
 86            mz_service=state.mz_service,
 87        )
 88
 89    def provides(self) -> list[Capability]:
 90        return [self.table]
 91
 92
 93class ValidateTable(Action):
 94    """Validates that a single table contains data that is consistent with the expected min/max watermark."""
 95
 96    @classmethod
 97    def requires(cls) -> set[type[Capability]]:
 98        return {BalancerdIsRunning, MzIsRunning, TableExists}
 99
100    def __init__(
101        self, capabilities: Capabilities, table: TableExists | None = None
102    ) -> None:
103        if table is not None:
104            self.table = table
105        else:
106            self.table = random.choice(capabilities.get(TableExists))
107
108        self.select_limit = random.choices([True, False], weights=[0.2, 0.8], k=1)[0]
109        super().__init__(capabilities)
110
111    def run(self, c: Composition, state: State) -> None:
112        # Validating via SELECT ... LIMIT is expensive as it requires creating a temporary table
113        # Therefore, only use it in 20% of validations.
114        if self.select_limit:
115            c.testdrive(
116                dedent(f"""
117                    > CREATE TEMPORARY TABLE {self.table.name}_select_limit (f1 INTEGER);
118                    > INSERT INTO {self.table.name}_select_limit SELECT * FROM {self.table.name} LIMIT 999999999;
119                    > SELECT MIN(f1), MAX(f1), COUNT(f1), COUNT(DISTINCT f1) FROM {self.table.name}_select_limit;
120                    {self.table.watermarks.min} {self.table.watermarks.max} {(self.table.watermarks.max-self.table.watermarks.min)+1} {(self.table.watermarks.max-self.table.watermarks.min)+1}
121                    > DROP TABLE {self.table.name}_select_limit
122                    """),
123                mz_service=state.mz_service,
124            )
125        else:
126            c.testdrive(
127                dedent(f"""
128                    > SELECT MIN(f1), MAX(f1), COUNT(f1), COUNT(DISTINCT f1) FROM {self.table.name};
129                    {self.table.watermarks.min} {self.table.watermarks.max} {(self.table.watermarks.max-self.table.watermarks.min)+1} {(self.table.watermarks.max-self.table.watermarks.min)+1}
130                    """),
131                mz_service=state.mz_service,
132            )
133
134
135class DropTable(Action):
136    """Drops a table that no view reads from."""
137
138    @classmethod
139    def requires(cls) -> set[type[Capability]]:
140        return {BalancerdIsRunning, MzIsRunning, TableExists}
141
142    def __init__(self, capabilities: Capabilities) -> None:
143        referenced: set[str] = set()
144        for view in capabilities.get(ViewExists):
145            referenced.update(input.name for input in view.inputs)
146
147        candidates = [
148            t for t in capabilities.get(TableExists) if t.name not in referenced
149        ]
150        self.table: TableExists | None = (
151            random.choice(candidates) if candidates else None
152        )
153        if self.table is not None:
154            capabilities.remove_capability_instance(self.table)
155        super().__init__(capabilities)
156
157    def __str__(self) -> str:
158        return f"{Action.__str__(self)} {self.table.name if self.table else '<none>'}"
159
160    def run(self, c: Composition, state: State) -> None:
161        if self.table is not None:
162            c.testdrive(
163                f"> DROP TABLE {self.table.name};",
164                mz_service=state.mz_service,
165            )
166
167
168class DML(Action):
169    """Performs an INSERT, DELETE or UPDATE against a table."""
170
171    @classmethod
172    def requires(cls) -> set[type[Capability]]:
173        return {BalancerdIsRunning, MzIsRunning, TableExists}
174
175    def __init__(self, capabilities: Capabilities) -> None:
176        self.table = random.choice(capabilities.get(TableExists))
177        self.delta = random.randint(1, self.table.max_rows_per_action)
178        super().__init__(capabilities)
179
180    def __str__(self) -> str:
181        return f"{Action.__str__(self)} {self.table.name}"
182
183
184class Insert(DML):
185    """Inserts rows into a table."""
186
187    def run(self, c: Composition, state: State) -> None:
188        prev_max = self.table.watermarks.max
189        self.table.watermarks.max = prev_max + self.delta
190        c.testdrive(
191            f"> INSERT INTO {self.table.name} SELECT * FROM generate_series({prev_max + 1}, {self.table.watermarks.max});",
192            mz_service=state.mz_service,
193        )
194
195
196class ShiftForward(DML):
197    """Update all rows from a table by incrementing their values by a constant."""
198
199    def run(self, c: Composition, state: State) -> None:
200        self.table.watermarks.shift(self.delta)
201        c.testdrive(
202            f"> UPDATE {self.table.name} SET f1 = f1 + {self.delta};",
203            mz_service=state.mz_service,
204        )
205
206
207class ShiftBackward(DML):
208    """Update all rows from a table by decrementing their values by a constant."""
209
210    def run(self, c: Composition, state: State) -> None:
211        self.table.watermarks.shift(-self.delta)
212        c.testdrive(
213            f"> UPDATE {self.table.name} SET f1 = f1 - {self.delta};",
214            mz_service=state.mz_service,
215        )
216
217
218class DeleteFromHead(DML):
219    """Delete the largest values from a table"""
220
221    def run(self, c: Composition, state: State) -> None:
222        self.table.watermarks.max = max(
223            self.table.watermarks.max - self.delta, self.table.watermarks.min
224        )
225        c.testdrive(
226            f"> DELETE FROM {self.table.name} WHERE f1 > {self.table.watermarks.max};",
227            mz_service=state.mz_service,
228        )
229
230
231class DeleteFromTail(DML):
232    """Delete the smallest values from a table"""
233
234    def run(self, c: Composition, state: State) -> None:
235        self.table.watermarks.min = min(
236            self.table.watermarks.min + self.delta, self.table.watermarks.max
237        )
238        c.testdrive(
239            f"> DELETE FROM {self.table.name} WHERE f1 < {self.table.watermarks.min};",
240            mz_service=state.mz_service,
241        )