Add force_database_error hooks to enable testing failing restores and…

… upgrades (#7287)

Add two new optional fields to the JSON object used to configure
`force_database_error`.

 * `_scopes`: Determines what operations will trigger an
   error. Defaults to `["query"]`. Other checked scopes are `restore`
   and `startup`. The `startup` scope is only checked against
   `INSTANCE` configuration, not database configurations.
 * `_versions`: If set, only fail if the current version is present
   in the array.

edb/server/compiler/__init__.py

@@ -24,6 +24,7 @@
 from .compiler import compile_edgeql_script
 from .compiler import new_compiler, new_compiler_from_pg, new_compiler_context
 from .compiler import compile, compile_schema_storage_in_delta
+from .compiler import maybe_force_database_error
 from .dbstate import QueryUnit, QueryUnitGroup
 from .enums import Capability, Cardinality
 from .enums import InputFormat, OutputFormat
@@ -43,6 +44,7 @@
     'OutputFormat',
     'analyze_explain_output',
     'compile_edgeql_script',
+    'maybe_force_database_error',
     'new_compiler',
     'new_compiler_from_pg',
     'new_compiler_context',

edb/server/compiler/compiler.py

@@ -45,6 +45,7 @@
 
 import immutables
 
+from edb import buildmeta
 from edb import errors
 
 from edb.common.typeutils import not_none
@@ -1306,6 +1307,8 @@ def describe_database_restore(
         # so it's simply ignored here and the I/O process will do it on its own
         units = compile(ctx=ctx, source=ddl_source).units
 
+        _check_force_database_error(ctx, scope='restore')
+
         schema = ctx.state.current_tx().get_schema(
             ctx.compiler_state.std_schema)
 
@@ -3263,22 +3266,26 @@ def _add_fake_property(
     )
 
 
-def _check_force_database_error(
-    ctx: CompileContext,
-    ql: qlast.Base,
+def maybe_force_database_error(
+    val: Optional[str],
+    *,
+    scope: str,
 ) -> None:
-    if isinstance(ql, qlast.ConfigOp):
+    # Check the string directly for false to skip a deserialization
+    if val is None or val == 'false':
         return
-
     try:
-        val = _get_config_val(ctx, 'force_database_error')
-        # Check the string directly for false to skip a deserialization
-        if val is None or val == 'false':
-            return
         err = json.loads(val)
         if not err:
             return
 
+        scopes = err.get('_scopes', ['query'])
+        if scope not in scopes:
+            return
+        versions = err.get('_versions')
+        if versions and buildmeta.get_version_string() not in versions:
+            return
+
         errcls = errors.EdgeDBError.get_error_class_from_name(err['type'])
         if context := err.get('context'):
             filename = context.get('filename')
@@ -3303,6 +3310,19 @@ def _check_force_database_error(
     raise errval
 
 
+def _check_force_database_error(
+    ctx: CompileContext,
+    ql: Optional[qlast.Base]=None,
+    *,
+    scope: str='query',
+) -> None:
+    if isinstance(ql, qlast.ConfigOp):
+        return
+
+    val = _get_config_val(ctx, 'force_database_error')
+    maybe_force_database_error(val, scope=scope)
+
+
 def _get_config_val(
     ctx: CompileContext,
     name: str,

edb/server/tenant.py

@@ -58,6 +58,8 @@
 from . import defines
 from . import metrics
 from . import pgcon
+from . import compiler as edbcompiler
+
 from .ha import adaptive as adaptive_ha
 from .ha import base as ha_base
 from .pgcon import errors as pgcon_errors
@@ -393,6 +395,12 @@ async def init(self) -> None:
         default_sysconfig = await self._load_sys_config("sysconfig_default")
         await self._load_reported_config()
 
+        # To make in-place upgrade failures more testable, check
+        # 'force_database_error' with a 'startup' scope.
+        force_error = self._server.config_lookup(
+            'force_database_error', sys_config)
+        edbcompiler.maybe_force_database_error(force_error, scope='startup')
+
         self._dbindex = dbview.DatabaseIndex(
             self,
             std_schema=self._server.get_std_schema(),

tests/test_server_config.py

@@ -1300,6 +1300,11 @@ async def test_server_proto_orphan_rollback_state(self):
     async def test_server_proto_configure_error(self):
         con1 = self.con
         con2 = await self.connect(database=con1.dbname)
+
+        version_str = await con1.query_single('''
+            select sys::get_version_as_str();
+        ''')
+
         try:
             await con2.execute('''
                 select 1;
@@ -1329,6 +1334,49 @@ async def test_server_proto_configure_error(self):
                             async with tx:
                                 await tx.query('select schema::Object')
 
+            # If we change the '_version' to something else we
+            # should be good
+            err = {
+                'type': 'SchemaError',
+                'message': 'danger',
+                'context': {'start': 42},
+                '_versions': [version_str + '1'],
+            }
+            await con1.execute(f'''
+                configure current database set force_database_error :=
+                  {qlquote.quote_literal(json.dumps(err))};
+            ''')
+            await con1.query('select schema::Object')
+
+            # It should also be fine if we set a '_scopes' other than 'query'
+            err = {
+                'type': 'SchemaError',
+                'message': 'danger',
+                'context': {'start': 42},
+                '_scopes': ['restore'],
+            }
+            await con1.execute(f'''
+                configure current database set force_database_error :=
+                  {qlquote.quote_literal(json.dumps(err))};
+            ''')
+            await con1.query('select schema::Object')
+
+            # But if we make it the current version it should still fail
+            err = {
+                'type': 'SchemaError',
+                'message': 'danger',
+                'context': {'start': 42},
+                '_versions': [version_str],
+            }
+            await con1.execute(f'''
+                configure current database set force_database_error :=
+                  {qlquote.quote_literal(json.dumps(err))};
+            ''')
+            with self.assertRaisesRegex(edgedb.SchemaError, 'danger'):
+                async for tx in con1.retrying_transaction():
+                    async with tx:
+                        await tx.query('select schema::Object')
+
             await con2.execute(f'''
                 configure session set force_database_error := "false";
             ''')