Support extension upgrades (#7998)

- ExtensionPackageMigration is a new schema object that specifies
  extension package upgrade scripts.
  Syntax is:
```
create extension package asdf migration
from version '2.0' to version '3.0'
```
- DDL: `ALTER MIGRATION asdf TO VERSION '2.0'`
- All extension versions are treated as being >=, though we
  will still prefer to do an exact match if possible.

---------

Co-authored-by: Aljaž Mur Eržen <[email protected]>

edb/buildmeta.py

@@ -60,7 +60,7 @@
 # The merge conflict there is a nice reminder that you probably need
 # to write a patch in edb/pgsql/patches.py, and then you should preserve
 # the old value.
-EDGEDB_CATALOG_VERSION = 2024_11_22_00_00
+EDGEDB_CATALOG_VERSION = 2024_11_26_00_00
 EDGEDB_MAJOR_VERSION = 6
 
 

edb/edgeql/ast.py

@@ -673,7 +673,6 @@ class DDLCommand(DDLOperation):
 
 class NonTransactionalDDLCommand(DDLCommand):
     __abstract_node__ = True
-    __rust_ignore__ = True
 
 
 class AlterAddInherit(DDLOperation):
@@ -900,20 +899,40 @@ class DropExtensionPackage(DropObject, ExtensionPackageCommand):
     pass
 
 
-class ExtensionCommand(UnqualifiedObjectCommand):
+class ExtensionPackageMigrationCommand(GlobalObjectCommand):
+    __abstract_node__ = True
+
+
+class CreateExtensionPackageMigration(
+    CreateObject, ExtensionPackageMigrationCommand
+):
+    from_version: Constant
+    to_version: Constant
+    body: NestedQLBlock
+
+
+class DropExtensionPackageMigration(
+    DropObject, ExtensionPackageMigrationCommand
+):
+    from_version: Constant
+    to_version: Constant
 
+
+class ExtensionCommand(UnqualifiedObjectCommand):
     __abstract_node__ = True
-    version: typing.Optional[Constant] = None
 
 
 class CreateExtension(CreateObject, ExtensionCommand):
-    # HACK: I think there is a bug in our plugin that made us not
-    # understand that this was overridden in ExtensionCommand.
-    pass
+    version: typing.Optional[Constant] = None
+
+
+class AlterExtension(DropObject, ExtensionCommand):
+    version: typing.Optional[Constant] = None
+    to_version: Constant
 
 
 class DropExtension(DropObject, ExtensionCommand):
-    pass
+    version: typing.Optional[Constant] = None
 
 
 class FutureCommand(UnqualifiedObjectCommand):
@@ -925,7 +944,7 @@ class CreateFuture(CreateObject, FutureCommand):
     pass
 
 
-class DropFuture(DropObject, ExtensionCommand):
+class DropFuture(DropObject, FutureCommand):
     pass
 
 

edb/edgeql/codegen.py

@@ -1199,6 +1199,41 @@ def after_name() -> None:
 
         self._visit_DropObject(node, 'EXTENSION PACKAGE', after_name=after_name)
 
+    def visit_CreateExtensionPackageMigration(
+        self,
+        node: qlast.CreateExtensionPackageMigration,
+    ) -> None:
+        self._write_keywords('CREATE EXTENSION PACKAGE')
+        self.write(' ')
+        self.write(ident_to_str(node.name.name))
+        self._write_keywords(' MIGRATION FROM ')
+        self._write_keywords(' VERSION ')
+        self.visit(node.from_version)
+        self._write_keywords(' TO ')
+        self.visit(node.to_version)
+
+        if node.body.text:
+            self.write(' {')
+            self._block_ws(1)
+            self.write(self.indent_text(node.body.text))
+            self._block_ws(-1)
+            self.write('}')
+        elif node.body.commands:
+            self._ddl_visit_body(node.body.commands)
+
+    def visit_DropExtensionPackageMigration(
+        self,
+        node: qlast.DropExtensionPackageMigration,
+    ) -> None:
+        self._write_keywords('DROP EXTENSION PACKAGE')
+        self.write(' ')
+        self.write(ident_to_str(node.name.name))
+        self._write_keywords(' MIGRATION FROM ')
+        self._write_keywords(' VERSION ')
+        self.visit(node.from_version)
+        self._write_keywords(' TO ')
+        self.visit(node.to_version)
+
     def visit_CreateExtension(
         self,
         node: qlast.CreateExtension,
@@ -1215,6 +1250,13 @@ def visit_CreateExtension(
         if node.commands:
             self._ddl_visit_body(node.commands)
 
+    def visit_AlterExtension(self, node: qlast.AlterExtension) -> None:
+        self._write_keywords('ALTER EXTENSION')
+        self.write(' ')
+        self.write(ident_to_str(node.name.name))
+        self._write_keywords(' TO VERSION ')
+        self.visit(node.to_version)
+
     def visit_DropExtension(
         self,
         node: qlast.DropExtension,

edb/edgeql/parser/grammar/ddl.py

@@ -914,6 +914,14 @@ def reduce_CreateExtensionPackageStmt(self, *kids):
     def reduce_DropExtensionPackageStmt(self, *kids):
         pass
 
+    @parsing.inline(0)
+    def reduce_CreateExtensionPackageMigrationStmt(self, *kids):
+        pass
+
+    @parsing.inline(0)
+    def reduce_DropExtensionPackageMigrationStmt(self, *kids):
+        pass
+
 
 #
 # CREATE EXTENSION PACKAGE
@@ -972,6 +980,68 @@ def reduce_DropExtensionPackageStmt(self, *kids):
         )
 
 
+#
+# CREATE EXTENSION PACKAGE MIGRATION
+#
+
+class CreateExtensionPackageMigrationBodyBlock(NestedQLBlock):
+
+    @property
+    def allowed_fields(self) -> typing.FrozenSet[str]:
+        return frozenset(
+            {'early_sql_script', 'late_sql_script'}
+        )
+
+    @property
+    def result(self) -> typing.Any:
+        return ExtensionPackageBody
+
+
+nested_ql_block(
+    'CreateExtensionPackage',
+    production_tpl=CreateExtensionPackageBodyBlock,
+)
+
+
+class CreateExtensionPackageMigrationStmt(Nonterm):
+
+    def reduce_CreateExtensionPackageMigrationStmt(self, *kids):
+        r"""%reduce CREATE EXTENSIONPACKAGE ShortNodeName
+                    MIGRATION FROM
+                    ExtensionVersion TO
+                    ExtensionVersion
+                    OptCreateExtensionPackageCommandsBlock
+        """
+        _, _, name, _, _, from_version, _, to_version, block = kids
+        self.val = qlast.CreateExtensionPackageMigration(
+            name=name.val,
+            from_version=from_version.val,
+            to_version=to_version.val,
+            body=block.val.body,
+            commands=block.val.fields,
+        )
+
+
+#
+# DROP EXTENSION PACKAGE MIGRATION
+#
+class DropExtensionPackageMigrationStmt(Nonterm):
+
+    def reduce_DropExtensionPackageMigrationStmt(self, *kids):
+        r"""%reduce DROP EXTENSIONPACKAGE ShortNodeName
+                    MIGRATION FROM
+                    ExtensionVersion TO
+                    ExtensionVersion
+        """
+        _, _, name, _, _, from_version, _, to_version = kids
+
+        self.val = qlast.DropExtensionPackageMigration(
+            name=name.val,
+            from_version=from_version.val,
+            to_version=to_version.val,
+        )
+
+
 #
 # EXTENSIONS
 #
@@ -983,6 +1053,10 @@ class ExtensionStmt(Nonterm):
     def reduce_CreateExtensionStmt(self, *kids):
         pass
 
+    @parsing.inline(0)
+    def reduce_AlterExtensionStmt(self, *kids):
+        pass
+
     @parsing.inline(0)
     def reduce_DropExtensionStmt(self, *kids):
         pass
@@ -1011,6 +1085,29 @@ def reduce_CreateExtensionStmt(self, *kids):
             commands=kids[4].val,
         )
 
+#
+# ALTER EXTENSION
+#
+
+
+commands_block(
+    'AlterExtension',
+    SetFieldStmt,
+)
+
+
+class AlterExtensionStmt(Nonterm):
+
+    def reduce_AlterExtensionStmt(self, *kids):
+        r"""%reduce ALTER EXTENSION ShortNodeName
+                    TO ExtensionVersion
+        """
+        _, _, name, _, ver = kids
+        self.val = qlast.AlterExtension(
+            name=name.val,
+            to_version=ver.val,
+        )
+
 
 #
 # DROP EXTENSION

edb/edgeql/qltypes.py

@@ -259,6 +259,7 @@ class SchemaObjectClass(s_enum.StrEnum):
     DATABASE = 'DATABASE'
     EXTENSION = 'EXTENSION'
     EXTENSION_PACKAGE = 'EXTENSION PACKAGE'
+    EXTENSION_PACKAGE_MIGRATION = 'EXTENSION PACKAGE MIGRATION'
     FUTURE = 'FUTURE'
     FUNCTION = 'FUNCTION'
     GLOBAL = 'GLOBAL'

edb/lib/ext/ai.edgeql

@@ -19,7 +19,7 @@
 
 CREATE EXTENSION PACKAGE ai VERSION '1.0' {
     set ext_module := "ext::ai";
-    set dependencies := ["pgvector==0.7"];
+    set dependencies := ["pgvector>=0.7"];
 
     create module ext::ai;
 

edb/lib/ext/auth.edgeql

@@ -19,7 +19,7 @@
 
 CREATE EXTENSION PACKAGE auth VERSION '1.0' {
     set ext_module := "ext::auth";
-    set dependencies := ["pgcrypto==1.3"];
+    set dependencies := ["pgcrypto>=1.3"];
 
     create module ext::auth;
 

edb/lib/sys.edgeql

@@ -57,13 +57,36 @@ CREATE TYPE sys::ExtensionPackage EXTENDING
         sys::SystemObject,
         schema::AnnotationSubject {
     CREATE REQUIRED PROPERTY script -> str;
-    CREATE REQUIRED PROPERTY version -> tuple<
-                                            major: std::int64,
-                                            minor: std::int64,
-                                            stage: sys::VersionStage,
-                                            stage_no: std::int64,
-                                            local: array<std::str>,
-                                        >;
+    CREATE REQUIRED PROPERTY version ->
+        tuple<
+             major: std::int64,
+             minor: std::int64,
+             stage: sys::VersionStage,
+             stage_no: std::int64,
+             local: array<std::str>,
+         >;
+};
+
+CREATE TYPE sys::ExtensionPackageMigration EXTENDING
+        sys::SystemObject,
+        schema::AnnotationSubject {
+    CREATE REQUIRED PROPERTY script -> str;
+    CREATE REQUIRED PROPERTY from_version ->
+        tuple<
+             major: std::int64,
+             minor: std::int64,
+             stage: sys::VersionStage,
+             stage_no: std::int64,
+             local: array<std::str>,
+         >;
+    CREATE REQUIRED PROPERTY to_version ->
+        tuple<
+             major: std::int64,
+             minor: std::int64,
+             stage: sys::VersionStage,
+             stage_no: std::int64,
+             local: array<std::str>,
+         >;
 };