Change how we establish connection + Rails 6 support

.rubocop.yml

@@ -12,6 +12,12 @@ Style/StringLiteralsInInterpolation:
 Layout/LineLength:
   Max: 180
 
+Metrics/MethodLength:
+  Enabled: false
+
+Metrics/AbcSize:
+  Enabled: false
+
 Metrics/BlockLength:
   Enabled: false
 

lib/planetscale_rails/tasks/psdb.rake

@@ -7,25 +7,42 @@ require "colorize"
 
 databases = ActiveRecord::Tasks::DatabaseTasks.setup_initial_database_yaml
 
-def shared_deps(name = nil)
-  return [:environment, :check_ci, "setup_pscale:#{name}".to_sym] if name
-  return %i[environment check_ci] if name.nil?
-end
-
 def puts_deploy_request_instructions
   ps_config = YAML.load_file(".pscale.yml")
   database = ps_config["database"]
   branch = ps_config["branch"]
   org = ps_config["org"]
 
-  puts "Create a deploy request for '#{branch}' by running:\n"
+  puts "Create a deploy request for '#{branch.colorize(:blue)}' by running:\n"
   puts "     pscale deploy-request create #{database} #{branch} --org #{org}\n\n"
 end
 
 def kill_pscale_process
   Process.kill("TERM", ENV["PSCALE_PID"].to_i) if ENV["PSCALE_PID"]
 end
 
+def delete_password
+  password_id = ENV["PSCALE_PASSWORD_ID"]
+  return unless password_id
+
+  ps_config = YAML.load_file(".pscale.yml")
+  database = ps_config["database"]
+  branch = ps_config["branch"]
+
+  command = "pscale password delete #{database} #{branch} #{password_id} #{ENV["SERVICE_TOKEN_CONFIG"]} --force"
+  output = `#{command}`
+
+  return if $CHILD_STATUS.success?
+
+  puts "Failed to cleanup credentials used for PlanetScale connection. Password ID: #{password_id}".colorize(:red)
+  puts "Command: #{command}"
+  puts "Output: #{output}"
+end
+
+def db_branch_colorized(database, branch)
+  "#{database.colorize(:blue)}/#{branch.colorize(:blue)}"
+end
+
 namespace :psdb do
   task check_ci: :environment do
     use_service_token = ENV["PSCALE_SERVICE_TOKEN"] && ENV["PSCALE_SERVICE_TOKEN_ID"]
@@ -38,165 +55,130 @@ namespace :psdb do
     end
   end
 
-  desc "Setup a proxy to connect to PlanetScale"
-  task "setup_pscale" => shared_deps do
-    ENV["ENABLE_PSDB"] = "1"
-    ENV["DISABLE_PS_GEM"] = "1"
+  def create_connection_string
     ps_config = YAML.load_file(".pscale.yml")
     database = ps_config["database"]
     branch = ps_config["branch"]
 
     raise "You must have `pscale` installed on your computer".colorize(:red) unless command?("pscale")
     if branch.blank? || database.blank?
-      raise "Your branch is not properly setup, please switch to a branch by using the CLI.".colorize(:red)
+      raise "Could not determine which PlanetScale branch to use from .pscale.yml. Please switch to a branch by using: `pscale switch database-name branch-name`".colorize(:red)
     end
 
-    r, = PTY.open
-
-    puts "Connecting to PlanetScale..."
-
-    # Spawns the process in the background
-    pid = Process.spawn("pscale connect #{database} #{branch} --port 3305 #{ENV["SERVICE_TOKEN_CONFIG"]}", out: r)
-    ENV["PSCALE_PID"] = pid.to_s
-
-    out = ""
-    start_time = Time.current
-    time_elapsed = Time.current - start_time
-    sleep(1)
-    while out.blank? && time_elapsed < 10.seconds
-      PTY.check(pid, true)
-      out = r.gets
-      time_elapsed = Time.current - start_time
+    short_hash = SecureRandom.hex(2)[0, 4]
+    password_name = "planetscale-rails-#{short_hash}"
+    command = "pscale password create #{database} #{branch} #{password_name} -f json #{ENV["SERVICE_TOKEN_CONFIG"]}"
+
+    output = `#{command}`
+
+    if $CHILD_STATUS.success?
+      response = JSON.parse(output)
+      puts "Successfully created credentials for PlanetScale #{db_branch_colorized(database, branch)}"
+      host = response["access_host_url"]
+      username = response["username"]
+      password = response["plain_text"]
+      ENV["PSCALE_PASSWORD_ID"] = response["id"]
+
+      "mysql2://#{username}:#{password}@#{host}:3306/#{database}"
+    else
+      puts "Failed to create credentials for PlanetScale #{db_branch_colorized(database, branch)}"
+      puts "Command: #{command}"
+      puts "Output: #{output}"
+      puts "Please verify that you have the correct permissions to create a password for this branch."
+      exit 1
     end
-
-    raise "Timed out waiting for PlanetScale connection to be established".colorize(:red) if time_elapsed > 10.seconds
-  ensure
-    r&.close
   end
 
-  namespace :setup_pscale do
-    ActiveRecord::Tasks::DatabaseTasks.for_each(databases) do |name|
-      desc "Setup a proxy to connect to #{name} in PlanetScale"
-      task name => shared_deps do
-        ENV["ENABLE_PSDB"] = "1"
-        ENV["DISABLE_PS_GEM"] = "1"
-        ps_config = YAML.load_file(".pscale.yml")
-        database = ps_config["database"]
-        branch = ps_config["branch"]
-
-        raise "You must have `pscale` installed on your computer" unless command?("pscale")
-        if branch.blank? || database.blank?
-          raise "Your branch is not properly setup, please switch to a branch by using the CLI."
-        end
-
-        config = Rails.configuration.database_configuration[Rails.env][name]
-
-        raise "Database #{name} is not configured for the current environment" unless config
-
-        r, = PTY.open
-
-        puts "Connecting to PlanetScale..."
-
-        # Spawns the process in the background
-        pid = Process.spawn("pscale connect #{database} #{branch} --port 3305 #{ENV["SERVICE_TOKEN_CONFIG"]}", out: r)
-        ENV["PSCALE_PID"] = pid.to_s
-
-        out = ""
-        start_time = Time.current
-        time_elapsed = Time.current - start_time
-        sleep(1)
-        while out.blank? && time_elapsed < 10.seconds
-          PTY.check(pid, true)
-          out = r.gets
-          time_elapsed = Time.current - start_time
-        end
-
-        raise "Timed out waiting for PlanetScale connection to be established" if time_elapsed > 10.seconds
-
-        # Comment out for now, this messes up when running migrations.
-        # Kernel.system("bundle exec rails db:environment:set RAILS_ENV=development")
-      ensure
-        r&.close
-      end
-    end
+  desc "Create credentials for PlanetScale and sets them to ENV['PSCALE_DATABASE_URL']"
+  task "create_creds" => %i[environment check_ci] do
+    ENV["PSCALE_DATABASE_URL"] = create_connection_string
   end
 
-  desc "Migrate the database for current environment"
-  task migrate: %i[environment check_ci setup_pscale] do
+  desc "Connects to the current PlanetScale branch and runs rails db:migrate"
+  task migrate: %i[environment check_ci create_creds] do
     db_configs = ActiveRecord::Base.configurations.configs_for(env_name: ActiveRecord::Tasks::DatabaseTasks.env)
 
     unless db_configs.size == 1
       raise "Found multiple database configurations, please specify which database you want to migrate using `psdb:migrate:<database_name>`".colorize(:red)
     end
 
     puts "Running migrations..."
-    Kernel.system("bundle exec rails db:migrate")
+
+    Kernel.system("DATABASE_URL=#{ENV["PSCALE_DATABASE_URL"]} bundle exec rails db:migrate")
+
     puts "Finished running migrations\n".colorize(:green)
     puts_deploy_request_instructions
   ensure
-    kill_pscale_process
+    delete_password
   end
 
   namespace :migrate do
     ActiveRecord::Tasks::DatabaseTasks.for_each(databases) do |name|
-      desc "Migrate #{name} database for current environment"
-      task name => shared_deps(name) do
+      desc "Connects to the current PlanetScale branch and runs rails db:migrate:#{name}"
+      task name => %i[environment check_ci create_creds] do
         puts "Running migrations..."
-        # We run it using the Kernel.system here because this properly handles
-        # when exceptions occur whereas Rake::Task.invoke does not.
-        Kernel.system("bundle exec rake db:migrate:#{name}")
+
+        name_env_key = "#{name.upcase}_DATABASE_URL"
+        Kernel.system("#{name_env_key}=#{ENV["PSCALE_DATABASE_URL"]} bundle exec rails db:migrate:#{name}")
 
         puts "Finished running migrations\n".colorize(:green)
         puts_deploy_request_instructions
       ensure
-        kill_pscale_process
-      end
-    end
-  end
-
-  namespace :truncate_all do
-    ActiveRecord::Tasks::DatabaseTasks.for_each(databases) do |name|
-      desc "Truncate all tables in #{name} database for current environment"
-      task name => shared_deps(name) do
-        puts "Truncating database..."
-        # We run it using the Kernel.system here because this properly handles
-        # when exceptions occur whereas Rake::Task.invoke does not.
-        Kernel.system("bundle exec rake db:truncate_all")
-        puts "Finished truncating database."
-      ensure
-        kill_pscale_process
+        delete_password
       end
     end
   end
 
   namespace :schema do
     namespace :load do
       ActiveRecord::Tasks::DatabaseTasks.for_each(databases) do |name|
-        desc "Load the current schema into the #{name} database"
-        task name => shared_deps(name) do
+        desc "Connects to the current PlanetScale branch and runs rails db:schema:load:#{name}"
+        task name => %i[environment check_ci create_creds] do
           puts "Loading schema..."
-          # We run it using the Kernel.system here because this properly handles
-          # when exceptions occur whereas Rake::Task.invoke does not.
-          Kernel.system("bundle exec rake db:schema:load:#{name}")
+
+          name_env_key = "#{name.upcase}_DATABASE_URL"
+          Kernel.system("#{name_env_key}=#{ENV["PSCALE_DATABASE_URL"]} bundle exec rake db:schema:load:#{name}")
+
           puts "Finished loading schema."
         ensure
-          kill_pscale_process
+          delete_password
         end
       end
     end
   end
 
+  desc "Connects to the current PlanetScale branch and runs rails db:rollback"
+  task rollback: %i[environment check_ci create_creds] do
+    db_configs = ActiveRecord::Base.configurations.configs_for(env_name: ActiveRecord::Tasks::DatabaseTasks.env)
+
+    unless db_configs.size == 1
+      raise "Found multiple database configurations, please specify which database you want to rollback using `psdb:rollback:<database_name>`".colorize(:red)
+    end
+
+    Kernel.system("DATABASE_URL=#{ENV["PSCALE_DATABASE_URL"]} bundle exec rails db:rollback")
+  ensure
+    delete_password
+  end
+
   namespace :rollback do
     ActiveRecord::Tasks::DatabaseTasks.for_each(databases) do |name|
-      desc "Rollback #{name} database for current environment"
-      task name => shared_deps(name) do
+      desc "Connects to the current PlanetScale branch and runs rails db:rollback:#{name}"
+      task name => %i[environment check_ci create_creds] do
+        required_version = Gem::Version.new("6.1.0.0")
+        rails_version = Gem::Version.new(Rails.version)
+
+        if rails_version < required_version
+          rails "This version of Rails does not support rollback commands for multi-database Rails apps. Please upgrade to at least Rails 6.1"
+        end
+
         puts "Rolling back migrations..."
-        # We run it using the Kernel.system here because this properly handles
-        # when exceptions occur whereas Rake::Task.invoke does not.
-        Kernel.system("bundle exec rake db:rollback:#{name}")
+
+        name_env_key = "#{name.upcase}_DATABASE_URL"
+        Kernel.system("#{name_env_key}=#{ENV["PSCALE_DATABASE_URL"]} bundle exec rake db:rollback:#{name}")
+
         puts "Finished rolling back migrations.".colorize(:green)
       ensure
-        kill_pscale_process
+        delete_password
       end
     end
   end

planetscale_rails.gemspec

@@ -29,5 +29,5 @@ Gem::Specification.new do |spec|
   spec.require_paths = ["lib"]
 
   spec.add_dependency "colorize", "~> 0.8.1"
-  spec.add_dependency "rails", "~> 7.0"
+  spec.add_dependency "rails", "~> 6.0"
 end