Merge branch 'jp-embedding-helper' into 'master'

Add support for embedding to migration helpers

See merge request https://gitlab.com/gitlab-org/gitlab/-/merge_requests/138143



Merged-by: Dmitry Gruzd <dgruzd@gitlab.com>
Approved-by: Nicolas Dular <ndular@gitlab.com>
Approved-by: Dmitry Gruzd <dgruzd@gitlab.com>
Reviewed-by: Nicolas Dular <ndular@gitlab.com>
Co-authored-by: Jan Provaznik <jprovaznik@gitlab.com>

ee/spec/support/helpers/ee/migrations_helpers.rb

@@ -12,8 +12,8 @@ def reset_column_information(klass)
 
     override :active_record_base
     def active_record_base(...)
-      if geo_migration?
-        ::Geo::TrackingBase
+      if custom_migration?
+        db_base_model
       else
         super
       end
@@ -21,8 +21,8 @@ def active_record_base(...)
 
     override :migrations_paths
     def migrations_paths
-      if geo_migration?
-        geo_db_config.configuration_hash[:migrations_paths]
+      if custom_migration?
+        db_base_model.connection_db_config.configuration_hash[:migrations_paths]
       else
         super
       end
@@ -44,18 +44,18 @@ def migrate!
     end
 
     def with_db_config(&block)
-      if geo_migration?
-        with_added_geo_connection { yield }
+      if custom_migration?
+        with_custom_connection { yield }
       else
         yield
       end
     end
 
-    def with_added_geo_connection
+    def with_custom_connection
       with_reestablished_active_record_base(reconnect: true) do
         reconfigure_db_connection(
-          name: :geo,
-          config_model: Geo::TrackingBase,
+          name: geo_migration? ? :geo : :embedding,
+          config_model: db_base_model,
           model: ActiveRecord::Base
         )
 
@@ -63,12 +63,26 @@ def with_added_geo_connection
       end
     end
 
+    def custom_migration?
+      geo_migration? || embedding_migration?
+    end
+
     def geo_migration?
       self.class.metadata[:geo]
     end
 
-    def geo_db_config
-      Geo::TrackingBase.connection_db_config
+    def embedding_migration?
+      self.class.metadata[:embedding]
+    end
+
+    def db_base_model
+      if geo_migration?
+        Geo::TrackingBase
+      elsif embedding_migration?
+        ::Embedding::ApplicationRecord
+      else
+        raise "unknown database migration"
+      end
     end
   end
 end