Avoid pickling schemas on _every_ DDL statement when restoring a dump (…

…#6362)

`restore()` doesn't care about anything other than compiled SQL or
config ops, so pickling and returning the schema with every DDL query
unit is wasteful, and, if the schema is large, quite slow.

edb/server/compiler/compiler.py

@@ -120,6 +120,7 @@ class CompileContext:
     bootstrap_mode: bool = False
     internal_schema_mode: bool = False
     log_ddl_as_migrations: bool = True
+    dump_restore_mode: bool = False
     notebook: bool = False
 
     def _assert_not_in_migration_block(
@@ -1164,6 +1165,7 @@ def describe_database_restore(
             schema_object_ids=schema_object_ids,
             log_ddl_as_migrations=False,
             protocol_version=protocol_version,
+            dump_restore_mode=True,
         )
 
         ctx.state.start_tx()
@@ -2437,18 +2439,19 @@ def _try_compile(
             unit.drop_db = comp.drop_db
             unit.create_db_template = comp.create_db_template
             unit.ddl_stmt_id = comp.ddl_stmt_id
-            if comp.user_schema is not None:
-                final_user_schema = comp.user_schema
-                unit.user_schema = pickle.dumps(comp.user_schema, -1)
-                unit.extensions, unit.ext_config_settings = (
-                    _extract_extensions(ctx, comp.user_schema)
-                )
-            if comp.cached_reflection is not None:
-                unit.cached_reflection = \
-                    pickle.dumps(comp.cached_reflection, -1)
-            if comp.global_schema is not None:
-                unit.global_schema = pickle.dumps(comp.global_schema, -1)
-                unit.roles = _extract_roles(comp.global_schema)
+            if not ctx.dump_restore_mode:
+                if comp.user_schema is not None:
+                    final_user_schema = comp.user_schema
+                    unit.user_schema = pickle.dumps(comp.user_schema, -1)
+                    unit.extensions, unit.ext_config_settings = (
+                        _extract_extensions(ctx, comp.user_schema)
+                    )
+                if comp.cached_reflection is not None:
+                    unit.cached_reflection = \
+                        pickle.dumps(comp.cached_reflection, -1)
+                if comp.global_schema is not None:
+                    unit.global_schema = pickle.dumps(comp.global_schema, -1)
+                    unit.roles = _extract_roles(comp.global_schema)
 
             unit.config_ops.extend(comp.config_ops)
 
@@ -2460,18 +2463,20 @@ def _try_compile(
                 )
             unit.sql = comp.sql
             unit.cacheable = comp.cacheable
-            if comp.user_schema is not None:
-                final_user_schema = comp.user_schema
-                unit.user_schema = pickle.dumps(comp.user_schema, -1)
-                unit.extensions, unit.ext_config_settings = (
-                    _extract_extensions(ctx, comp.user_schema)
-                )
-            if comp.cached_reflection is not None:
-                unit.cached_reflection = \
-                    pickle.dumps(comp.cached_reflection, -1)
-            if comp.global_schema is not None:
-                unit.global_schema = pickle.dumps(comp.global_schema, -1)
-                unit.roles = _extract_roles(comp.global_schema)
+
+            if not ctx.dump_restore_mode:
+                if comp.user_schema is not None:
+                    final_user_schema = comp.user_schema
+                    unit.user_schema = pickle.dumps(comp.user_schema, -1)
+                    unit.extensions, unit.ext_config_settings = (
+                        _extract_extensions(ctx, comp.user_schema)
+                    )
+                if comp.cached_reflection is not None:
+                    unit.cached_reflection = \
+                        pickle.dumps(comp.cached_reflection, -1)
+                if comp.global_schema is not None:
+                    unit.global_schema = pickle.dumps(comp.global_schema, -1)
+                    unit.roles = _extract_roles(comp.global_schema)
 
             if comp.modaliases is not None:
                 unit.modaliases = comp.modaliases
@@ -2496,15 +2501,17 @@ def _try_compile(
         elif isinstance(comp, dbstate.MigrationControlQuery):
             unit.sql = comp.sql
             unit.cacheable = comp.cacheable
-            if comp.user_schema is not None:
-                final_user_schema = comp.user_schema
-                unit.user_schema = pickle.dumps(comp.user_schema, -1)
-                unit.extensions, unit.ext_config_settings = (
-                    _extract_extensions(ctx, comp.user_schema)
-                )
-            if comp.cached_reflection is not None:
-                unit.cached_reflection = \
-                    pickle.dumps(comp.cached_reflection, -1)
+
+            if not ctx.dump_restore_mode:
+                if comp.user_schema is not None:
+                    final_user_schema = comp.user_schema
+                    unit.user_schema = pickle.dumps(comp.user_schema, -1)
+                    unit.extensions, unit.ext_config_settings = (
+                        _extract_extensions(ctx, comp.user_schema)
+                    )
+                if comp.cached_reflection is not None:
+                    unit.cached_reflection = \
+                        pickle.dumps(comp.cached_reflection, -1)
             unit.ddl_stmt_id = comp.ddl_stmt_id
 
             if comp.modaliases is not None: