db.rake 2.4 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475
  1. # frozen_string_literal: true
  2. require_relative '../mastodon/snowflake'
  3. def each_schema_load_environment
  4. # If we're in development, also run this for the test environment.
  5. # This is a somewhat hacky way to do this, so here's why:
  6. # 1. We have to define this before we load the schema, or we won't
  7. # have a timestamp_id function when we get to it in the schema.
  8. # 2. db:setup calls db:schema:load_if_ruby, which calls
  9. # db:schema:load, which we define above as having a prerequisite
  10. # of this task.
  11. # 3. db:schema:load ends up running
  12. # ActiveRecord::Tasks::DatabaseTasks.load_schema_current, which
  13. # calls a private method `each_current_configuration`, which
  14. # explicitly also does the loading for the `test` environment
  15. # if the current environment is `development`, so we end up
  16. # needing to do the same, and we can't even use the same method
  17. # to do it.
  18. if Rails.env.development?
  19. test_conf = ActiveRecord::Base.configurations['test']
  20. if test_conf['database']&.present?
  21. ActiveRecord::Base.establish_connection(:test)
  22. yield
  23. ActiveRecord::Base.establish_connection(Rails.env.to_sym)
  24. end
  25. end
  26. yield
  27. end
  28. namespace :db do
  29. namespace :migrate do
  30. desc 'Setup the db or migrate depending on state of db'
  31. task setup: :environment do
  32. begin
  33. if ActiveRecord::Migrator.current_version.zero?
  34. Rake::Task['db:migrate'].invoke
  35. Rake::Task['db:seed'].invoke
  36. end
  37. rescue ActiveRecord::NoDatabaseError
  38. Rake::Task['db:setup'].invoke
  39. else
  40. Rake::Task['db:migrate'].invoke
  41. end
  42. end
  43. end
  44. # Before we load the schema, define the timestamp_id function.
  45. # Idiomatically, we might do this in a migration, but then it
  46. # wouldn't end up in schema.rb, so we'd need to figure out a way to
  47. # get it in before doing db:setup as well. This is simpler, and
  48. # ensures it's always in place.
  49. Rake::Task['db:schema:load'].enhance ['db:define_timestamp_id']
  50. # After we load the schema, make sure we have sequences for each
  51. # table using timestamp IDs.
  52. Rake::Task['db:schema:load'].enhance do
  53. Rake::Task['db:ensure_id_sequences_exist'].invoke
  54. end
  55. task :define_timestamp_id do
  56. each_schema_load_environment do
  57. Mastodon::Snowflake.define_timestamp_id
  58. end
  59. end
  60. task :ensure_id_sequences_exist do
  61. each_schema_load_environment do
  62. Mastodon::Snowflake.ensure_id_sequences_exist
  63. end
  64. end
  65. end