tests: More shared code based on the patterns we have

ci/test/lint-main/checks/check-mzcompose-files.sh

@@ -58,7 +58,7 @@ check_default_workflow_references_others() {
       if (( MATCHES_COUNT > 1 )); then
         # mzcompose file contains more than one workflow
 
-        LOOP_DETECTED=$(grep "c.workflow(name" "$file" -c)
+        LOOP_DETECTED=$(grep -c -E "c\.workflow\(name|c\.run_all_workflows\(" "$file")
 
         if (( LOOP_DETECTED < 1 )); then
           echo "$file contains more than one workflow but does not seem to loop over the workflows"

misc/python/materialize/backup_restore.py

@@ -0,0 +1,71 @@
+# Copyright Materialize, Inc. and contributors. All rights reserved.
+#
+# Use of this software is governed by the Business Source License
+# included in the LICENSE file at the root of this repository.
+#
+# 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.
+
+"""Shared backup & restore test logic for CockroachDB and Postgres metadata stores."""
+
+from textwrap import dedent
+
+from materialize.mzcompose.composition import Composition, Service
+
+
+def workflow_default(c: Composition) -> None:
+    """Basic Backup & Restore test: create a table, back up, insert more, restore, verify."""
+
+    # Enable versioning for the Persist bucket
+    c.enable_minio_versioning()
+
+    # Start Materialize, and set up some basic state in it
+    c.up("materialized", Service("testdrive", idle=True))
+    c.testdrive(
+        dedent(
+            """
+                > DROP TABLE IF EXISTS numbers;
+                > CREATE TABLE IF NOT EXISTS numbers (id BIGINT);
+                > INSERT INTO numbers SELECT * from generate_series(1, 1);
+                > INSERT INTO numbers SELECT * from generate_series(1, 10);
+                > INSERT INTO numbers SELECT * from generate_series(1, 100);
+                """
+        )
+    )
+
+    c.backup()
+
+    # Make further updates to Materialize's state
+    for i in range(0, 100):
+        # TODO: This seems to be enough to produce interesting shard state;
+        # ie. if we remove the restore-blob step we can see the restore fail.
+        # Is there any cheaper or more obvious way to do that?
+        c.testdrive(
+            dedent(
+                """
+                    > INSERT INTO numbers SELECT * from generate_series(1, 1);
+                    > INSERT INTO numbers SELECT * from generate_series(1, 10);
+                    > INSERT INTO numbers SELECT * from generate_series(1, 100);
+                    """
+            )
+        )
+
+    # Restore from backup, run persistcli restore-blob and restart Mz
+    c.restore()
+
+    # Confirm that the database is readable / has shard data
+    c.exec_metadata_store_sql(
+        "SELECT shard, min(sequence_number), max(sequence_number) "
+        "FROM consensus.consensus GROUP BY 1 ORDER BY 2 DESC, 3 DESC, 1 ASC LIMIT 32;"
+    )
+
+    # Check that the cluster is up and that it answers queries as of the old state
+    c.testdrive(
+        dedent(
+            """
+                > SELECT count(*) FROM numbers;
+                111
+                """
+        )
+    )

misc/python/materialize/checks/all_checks/alter_connection.py

@@ -14,15 +14,11 @@
 
 from materialize.checks.actions import Testdrive
 from materialize.checks.checks import Check, externally_idempotent
-from materialize.checks.common import KAFKA_SCHEMA_WITH_SINGLE_STRING_FIELD
+from materialize.checks.common import KAFKA_SCHEMA
 from materialize.checks.executors import Executor
 from materialize.mz_version import MzVersion
 
 
-def schema() -> str:
-    return dedent(KAFKA_SCHEMA_WITH_SINGLE_STRING_FIELD)
-
-
 class SshChange(Enum):
     ADD_SSH = 1
     DROP_SSH = 2
@@ -50,7 +46,7 @@ def initialize(self) -> Testdrive:
         i = self.index
 
         return Testdrive(
-            schema()
+            KAFKA_SCHEMA
             + dedent(
                 f"""
                 $ postgres-execute connection=postgres://mz_system:materialize@${{testdrive.materialize-internal-sql-addr}}
@@ -95,7 +91,7 @@ def manipulate(self) -> list[Testdrive]:
         i = self.index
 
         return [
-            Testdrive(schema() + dedent(s))
+            Testdrive(KAFKA_SCHEMA + dedent(s))
             for s in [
                 f"""
                 $ kafka-ingest topic=alter-connection-{i}a format=bytes

misc/python/materialize/checks/all_checks/alter_index.py

@@ -10,18 +10,14 @@
 
 from materialize.checks.actions import Testdrive
 from materialize.checks.checks import Check, externally_idempotent
-from materialize.checks.common import KAFKA_SCHEMA_WITH_SINGLE_STRING_FIELD
-
-
-def schema() -> str:
-    return dedent(KAFKA_SCHEMA_WITH_SINGLE_STRING_FIELD)
+from materialize.checks.common import KAFKA_SCHEMA
 
 
 @externally_idempotent(False)
 class AlterIndex(Check):
     def initialize(self) -> Testdrive:
         return Testdrive(
-            schema()
+            KAFKA_SCHEMA
             + dedent(
                 """
                 > CREATE TABLE alter_index_table (f1 STRING);
@@ -46,7 +42,7 @@ def initialize(self) -> Testdrive:
 
     def manipulate(self) -> list[Testdrive]:
         return [
-            Testdrive(schema() + dedent(s))
+            Testdrive(KAFKA_SCHEMA + dedent(s))
             for s in [
                 """
                 $ postgres-execute connection=postgres://mz_system:materialize@${testdrive.materialize-internal-sql-addr}

misc/python/materialize/checks/all_checks/continual_task.py

@@ -10,15 +10,11 @@
 
 from materialize.checks.actions import Testdrive
 from materialize.checks.checks import Check
-from materialize.checks.common import KAFKA_SCHEMA_WITH_SINGLE_STRING_FIELD
+from materialize.checks.common import KAFKA_SCHEMA
 from materialize.checks.executors import Executor
 from materialize.mz_version import MzVersion
 
 
-def schemas() -> str:
-    return dedent(KAFKA_SCHEMA_WITH_SINGLE_STRING_FIELD)
-
-
 class AuditLogCT(Check):
     """Continual Task for audit logging"""
 
@@ -27,7 +23,7 @@ def _can_run(self, e: Executor) -> bool:
 
     def initialize(self) -> Testdrive:
         return Testdrive(
-            schemas()
+            KAFKA_SCHEMA
             + dedent(
                 """
                 > CREATE TABLE t_input (key INT);
@@ -42,7 +38,7 @@ def initialize(self) -> Testdrive:
 
     def manipulate(self) -> list[Testdrive]:
         return [
-            Testdrive(schemas() + dedent(s))
+            Testdrive(KAFKA_SCHEMA + dedent(s))
             for s in [
                 """
                 > INSERT INTO t_input VALUES (2), (3);
@@ -74,7 +70,7 @@ def _can_run(self, e: Executor) -> bool:
 
     def initialize(self) -> Testdrive:
         return Testdrive(
-            schemas()
+            KAFKA_SCHEMA
             + dedent(
                 """
                 > CREATE TABLE big (key INT);
@@ -94,7 +90,7 @@ def initialize(self) -> Testdrive:
 
     def manipulate(self) -> list[Testdrive]:
         return [
-            Testdrive(schemas() + dedent(s))
+            Testdrive(KAFKA_SCHEMA + dedent(s))
             for s in [
                 """
                 > UPDATE small SET val = 'v' || val;
@@ -140,7 +136,7 @@ def _can_run(self, e: Executor) -> bool:
 
     def initialize(self) -> Testdrive:
         return Testdrive(
-            schemas()
+            KAFKA_SCHEMA
             + dedent(
                 """
                 > CREATE TABLE append_only (key INT, val INT);
@@ -155,7 +151,7 @@ def initialize(self) -> Testdrive:
 
     def manipulate(self) -> list[Testdrive]:
         return [
-            Testdrive(schemas() + dedent(s))
+            Testdrive(KAFKA_SCHEMA + dedent(s))
             for s in [
                 """
                 > INSERT INTO append_only VALUES (1, 3), (2, 4)

misc/python/materialize/checks/all_checks/explain_catalog_item.py

@@ -7,7 +7,6 @@
 # the Business Source License, use of this software will be governed
 # by the Apache License, Version 2.0.
 
-import re
 from textwrap import dedent
 
 from materialize.checks.actions import Testdrive
@@ -135,7 +134,3 @@ def validate(self) -> Testdrive:
         )
 
         return Testdrive(sql)
-
-
-def remove_target_cluster_from_explain(sql: str) -> str:
-    return re.sub(r"\n\s*Target cluster: \w+\n", "", sql)

misc/python/materialize/checks/all_checks/identifiers.py

@@ -14,7 +14,7 @@
 
 from materialize.checks.actions import Testdrive
 from materialize.checks.checks import Check
-from materialize.checks.common import KAFKA_SCHEMA_WITH_SINGLE_STRING_FIELD
+from materialize.checks.common import KAFKA_SCHEMA
 from materialize.mz_version import MzVersion
 from materialize.util import naughty_strings
 
@@ -34,10 +34,6 @@ def sq(ident: str) -> Any:
     return literal(ident)
 
 
-def schemas() -> str:
-    return dedent(KAFKA_SCHEMA_WITH_SINGLE_STRING_FIELD)
-
-
 def cluster() -> str:
     return "> CREATE CLUSTER identifiers SIZE 'scale=1,workers=4'\n"
 
@@ -126,7 +122,7 @@ def initialize(self) -> Testdrive:
             > COMMENT ON COLUMN {dq(self.ident["schema"])}.{dq(self.ident["table"])}.{dq(self.ident["column"])} IS {sq(self.ident["comment_column"])};
             """
 
-        return Testdrive(schemas() + cluster() + dedent(cmds))
+        return Testdrive(KAFKA_SCHEMA + cluster() + dedent(cmds))
 
     def manipulate(self) -> list[Testdrive]:
         cmds = [

misc/python/materialize/checks/all_checks/materialized_views.py

@@ -6,7 +6,6 @@
 # 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.
-import re
 from textwrap import dedent
 
 from materialize.checks.actions import Testdrive
@@ -384,7 +383,3 @@ def validate(self) -> Testdrive:
             """
             )
         )
-
-
-def remove_target_cluster_from_explain(sql: str) -> str:
-    return re.sub(r"\n\s*Target cluster: \w+\n", "", sql)

misc/python/materialize/checks/all_checks/multiple_partitions.py

@@ -10,11 +10,7 @@
 
 from materialize.checks.actions import Testdrive
 from materialize.checks.checks import Check, externally_idempotent
-from materialize.checks.common import KAFKA_SCHEMA_WITH_SINGLE_STRING_FIELD
-
-
-def schemas() -> str:
-    return dedent(KAFKA_SCHEMA_WITH_SINGLE_STRING_FIELD)
+from materialize.checks.common import KAFKA_SCHEMA
 
 
 @externally_idempotent(False)
@@ -23,7 +19,7 @@ class MultiplePartitions(Check):
 
     def initialize(self) -> Testdrive:
         return Testdrive(
-            schemas()
+            KAFKA_SCHEMA
             + dedent(
                 """
                 $ postgres-execute connection=postgres://mz_system:materialize@${testdrive.materialize-internal-sql-addr}
@@ -49,7 +45,7 @@ def initialize(self) -> Testdrive:
 
     def manipulate(self) -> list[Testdrive]:
         return [
-            Testdrive(schemas() + dedent(s))
+            Testdrive(KAFKA_SCHEMA + dedent(s))
             for s in [
                 """
                 # ingest B-key entries

misc/python/materialize/checks/all_checks/mysql_cdc.py

@@ -7,7 +7,6 @@
 # the Business Source License, use of this software will be governed
 # by the Apache License, Version 2.0.
 
-import re
 from random import Random
 from textwrap import dedent
 from typing import Any
@@ -590,7 +589,3 @@ def validate(self) -> Testdrive:
                 """
             )
         )
-
-
-def remove_target_cluster_from_explain(sql: str) -> str:
-    return re.sub(r"\n\s*Target cluster: \w+\n", "", sql)

misc/python/materialize/checks/all_checks/pg_cdc.py

@@ -7,7 +7,6 @@
 # the Business Source License, use of this software will be governed
 # by the Apache License, Version 2.0.
 
-import re
 from random import Random
 from textwrap import dedent
 from typing import Any
@@ -382,7 +381,3 @@ def validate(self) -> Testdrive:
                 """
             )
         )
-
-
-def remove_target_cluster_from_explain(sql: str) -> str:
-    return re.sub(r"\n\s*Target cluster: \w+\n", "", sql)

misc/python/materialize/checks/all_checks/ranges.py

@@ -10,17 +10,13 @@
 
 from materialize.checks.actions import Testdrive
 from materialize.checks.checks import Check
-from materialize.checks.common import KAFKA_SCHEMA_WITH_SINGLE_STRING_FIELD
-
-
-def schema() -> str:
-    return dedent(KAFKA_SCHEMA_WITH_SINGLE_STRING_FIELD)
+from materialize.checks.common import KAFKA_SCHEMA
 
 
 class Range(Check):
     def initialize(self) -> Testdrive:
         return Testdrive(
-            schema()
+            KAFKA_SCHEMA
             + dedent(
                 """
             > CREATE TABLE range_table (