1 files changed, 37 insertions, 0 deletions
diff --git a/lib/tasks/gitlab/db.rake b/lib/tasks/gitlab/db.rake
index 4917d496d07..61318570fd5 100644
--- a/lib/tasks/gitlab/db.rake
+++ b/lib/tasks/gitlab/db.rake
@@ -39,6 +39,11 @@ namespace :gitlab do
       # PG: http://www.postgresql.org/docs/current/static/ddl-depend.html
       # Add `IF EXISTS` because cascade could have already deleted a table.
       tables.each { |t| connection.execute("DROP TABLE IF EXISTS #{connection.quote_table_name(t)} CASCADE") }
+
+      # Drop all extra schema objects GitLab owns
+      Gitlab::Database::EXTRA_SCHEMAS.each do |schema|
+        connection.execute("DROP SCHEMA IF EXISTS #{connection.quote_table_name(schema)}")
+      end
     end
 
     desc 'GitLab | DB | Configures the database by running migrate, or by loading the schema and seeding if needed'
@@ -129,5 +134,37 @@ namespace :gitlab do
     Rake::Task['db:structure:load'].enhance do
       Rake::Task['gitlab:db:load_custom_structure'].invoke
     end
+
+    desc 'Create missing dynamic database partitions'
+    task :create_dynamic_partitions do
+      Gitlab::Database::Partitioning::PartitionCreator.new.create_partitions
+    end
+
+    # This is targeted towards deploys and upgrades of GitLab.
+    # Since we're running migrations already at this time,
+    # we also check and create partitions as needed here.
+    Rake::Task['db:migrate'].enhance do
+      Rake::Task['gitlab:db:create_dynamic_partitions'].invoke
+    end
+
+    # When we load the database schema from db/structure.sql
+    # we don't have any dynamic partitions created. We don't really need to
+    # because application initializers/sidekiq take care of that, too.
+    # However, the presence of partitions for a table has influence on their
+    # position in db/structure.sql (which is topologically sorted).
+    #
+    # Other than that it's helpful to create partitions early when bootstrapping
+    # a new installation.
+    Rake::Task['db:structure:load'].enhance do
+      Rake::Task['gitlab:db:create_dynamic_partitions'].invoke
+    end
+
+    # During testing, db:test:load restores the database schema from scratch
+    # which does not include dynamic partitions. We cannot rely on application
+    # initializers here as the application can continue to run while
+    # a rake task reloads the database schema.
+    Rake::Task['db:test:load'].enhance do
+      Rake::Task['gitlab:db:create_dynamic_partitions'].invoke
+    end
   end
 end