Add an internal-only TEMPLATE BRANCH that directly uses sql templates (…

…#7477)

The test suite used to rely on this for cloning databases, and switching
away caused some big slowdowns. Both because this method is somewhat faster
but more because the new method only allows one branch at a time to avoid
going overbudget on connections.

Progress on #7329.
On my machine this speed up edb test from 13:49 to 11:20.

edb/edgeql-parser/src/keywords.rs

@@ -95,6 +95,7 @@ pub const UNRESERVED_KEYWORDS: phf::Set<&str> = phf_set!(
     "superuser",
     "system",
     "target",
+    "template",
     "ternary",
     "text",
     "then",

edb/edgeql/ast.py

@@ -857,6 +857,7 @@ class BranchType(s_enum.StrEnum):
     EMPTY = 'EMPTY'
     SCHEMA = 'SCHEMA'
     DATA = 'DATA'
+    TEMPLATE = 'TEMPLATE'
 
 
 class DatabaseCommand(ExternalObjectCommand):

edb/edgeql/parser/grammar/ddl.py

@@ -833,6 +833,16 @@ def reduce_create_data_branch(self, *kids):
             branch_type=qlast.BranchType.DATA,
         )
 
+    def reduce_create_template_branch(self, *kids):
+        """%reduce
+            CREATE TEMPLATE BRANCH DatabaseName FROM DatabaseName
+        """
+        self.val = qlast.CreateDatabase(
+            name=kids[3].val,
+            template=kids[5].val,
+            branch_type=qlast.BranchType.TEMPLATE,
+        )
+
 
 #
 # DROP BRANCH

edb/pgsql/delta.py

@@ -7045,11 +7045,18 @@ def apply(
         tenant_id = self._get_tenant_id(context)
         db_name = common.get_database_backend_name(
             str(self.classname), tenant_id=tenant_id)
-        # We always use the base template, even for branches, since we
+        # We use the base template for SCHEMA and DATA branches, since we
         # implement branches ourselves using pg_dump in order to avoid
         # connection restrictions.
+        # For internal-only TEMPLATE branches, we use the source as
+        # the template.
+        template = (
+            self.template
+            if self.template and self.branch_type == ql_ast.BranchType.TEMPLATE
+            else edbdef.EDGEDB_TEMPLATE_DB
+        )
         tpl_name = common.get_database_backend_name(
-            edbdef.EDGEDB_TEMPLATE_DB, tenant_id=tenant_id)
+            template, tenant_id=tenant_id)
         self.pgops.add(
             dbops.CreateDatabase(
                 dbops.Database(

edb/schema/database.py

@@ -90,6 +90,16 @@ def _cmd_tree_from_ast(
         assert isinstance(astnode, qlast.CreateDatabase)
         if astnode.template is not None:
             cmd.template = astnode.template.name
+
+        if (
+            astnode.branch_type == qlast.BranchType.TEMPLATE
+            and not context.testmode
+        ):
+            raise errors.EdgeQLSyntaxError(
+                f'unexpected TEMPLATE',
+                span=astnode.span,
+            )
+
         cmd.branch_type = astnode.branch_type
 
         return cmd

edb/server/protocol/execute.pyx

@@ -259,6 +259,7 @@ async def execute(
             await tenant.on_before_create_db_from_template(
                 query_unit.create_db_template,
                 dbv.dbname,
+                query_unit.create_db_mode,
             )
         if query_unit.drop_db:
             await tenant.on_before_drop_db(

edb/server/tenant.py

@@ -1323,16 +1323,23 @@ async def on_before_drop_db(
         )
 
     async def on_before_create_db_from_template(
-        self, dbname: str, current_dbname: str
+        self, dbname: str, current_dbname: str, mode: str
     ) -> None:
         # Make sure the database exists.
         # TODO: Is it worth producing a nicer error message if it
         # fails on the backside? (Because of a race?)
         self.get_db(dbname=dbname)
 
+        if mode == 'TEMPLATE':
+            await self.ensure_database_not_connected(dbname)
+
     async def on_after_create_db_from_template(
         self, tgt_dbname: str, src_dbname: str, mode: str
     ) -> None:
+        if mode == 'TEMPLATE':
+            self.allow_database_connections(tgt_dbname)
+            return
+
         logger.info('Starting copy from %s to %s', src_dbname, tgt_dbname)
         from edb.pgsql import common
         from . import bootstrap  # noqa: F402

edb/testbase/server.py

@@ -1224,9 +1224,13 @@ async def setup_and_connect(cls):
             base_db_name, _, _ = dbname.rpartition('_')
 
             if cls.get_setup_script():
+                await admin_conn.execute('''
+                    configure session set __internal_testmode := true;
+                ''')
+
                 create_command = (
-                    f'CREATE DATA BRANCH {qlquote.quote_ident(dbname)}'
-                    f' FROM {qlquote.quote_ident(base_db_name)}'
+                    f'CREATE TEMPLATE BRANCH {qlquote.quote_ident(dbname)}'
+                    f' FROM {qlquote.quote_ident(base_db_name)};'
                 )
             else:
                 create_command = (