feat: support sql load

CHANGELOG.md

@@ -2,6 +2,7 @@
 
 ## Ongoing
 
+- Add support for sql load in rake tasks (#275).
 - Add support for sql dump in rake tasks (#273).
 - Add support for table optimize hints (#266).
 

lib/active_record/connection_adapters/cockroachdb/database_tasks.rb

@@ -45,8 +45,49 @@ def structure_dump(filename, extra_flags=nil)
         end
 
         def structure_load(filename, extra_flags=nil)
-          raise "db:structure:load is unimplemented. See https://github.com/cockroachdb/activerecord-cockroachdb-adapter/issues/2"
+          if extra_flags
+            raise "No flag supported yet, please raise an issue if needed. " \
+              "https://github.com/cockroachdb/activerecord-cockroachdb-adapter/issues/new"
+          end
+
+          run_cmd("cockroach", ["sql", "--set", "errexit=false", "--file", filename], "loading")
+        end
+
+        private
+
+        # Adapted from https://github.com/rails/rails/blob/a5fc471b3/activerecord/lib/active_record/tasks/postgresql_database_tasks.rb#L106.
+        # Using https://www.cockroachlabs.com/docs/stable/connection-parameters.html#additional-connection-parameters.
+        def cockroach_env
+          usr_pwd = ""
+          if configuration_hash[:username]
+            usr_pwd += configuration_hash[:username].to_s
+            if configuration_hash[:password]
+              usr_pwd += ":"
+              usr_pwd += configuration_hash[:password].to_s
+            end
+            usr_pwd += "@"
+          end
+
+          port = ""
+          port = ":#{configuration_hash[:port]}" if configuration_hash[:port]
+
+          params = %i(sslmode sslrootcert sslcert sslkey).filter_map do |key|
+            "#{key}=#{configuration_hash[key]}" if configuration_hash[key]
+          end.join("&")
+          params = "?#{params}" unless params.empty?
+
+          url = "postgres://#{usr_pwd}#{db_config.host}#{port}/#{db_config.database}#{params}"
+
+          {
+            # NOTE: sslmode in the url will take precedence over this setting, hence
+            #   we don't need to conditionally set it.
+            "COCKROACH_INSECURE" => "true",
+            "COCKROACH_URL" => url
+          }
         end
+        # The `#run_cmd` method use `psql_env` to set environments variables.
+        # We override it with cockroach env variables.
+        alias_method :psql_env, :cockroach_env
       end
     end
   end

test/cases/tasks/cockroachdb_rake_test.rb

@@ -54,4 +54,62 @@ def test_structure_dump
       refute read.include?("CREATE TABLE public.accounts ("), "\"accounts\" table should be ignored"
     end
   end
+
+  class CockroachDBStructureLoadTest < ActiveRecord::TestCase
+    def setup
+      @configuration = {
+        adapter: "cockroachdb",
+        database: "my-app-db",
+        host: "localhost"
+      }
+    end
+
+    def test_structure_load
+      filename = "awesome-file.sql"
+      assert_called_with(
+        Kernel,
+        :system,
+        [
+          {"COCKROACH_INSECURE"=>"true", "COCKROACH_URL"=>"postgres://localhost/my-app-db"},
+          "cockroach", "sql", "--set", "errexit=false", "--file", filename
+        ],
+        returns: true
+      ) do
+        ActiveRecord::Tasks::DatabaseTasks.structure_load(@configuration, filename)
+      end
+    end
+
+    def test_url_generation
+      assert_correct_url @configuration.merge(
+        %i(sslmode sslrootcert sslcert sslkey).to_h { [_1, "v#{_1}"] }
+      ), "postgres://localhost/my-app-db?sslmode=vsslmode&sslrootcert=vsslrootcert&sslcert=vsslcert&sslkey=vsslkey"
+      assert_correct_url @configuration.merge({
+        username: "root",
+        port: 1234
+      }), "postgres://root@localhost:1234/my-app-db"
+      assert_correct_url @configuration.merge({
+        username: "root",
+        password: "secret"
+      }), "postgres://root:secret@localhost/my-app-db"
+    end
+
+    private
+
+    # Verify that given a config we generate the expected connection URL,
+    # and that if we parse it again, we get the same config. Except the
+    # `adapter` key, that'll changed to postgresql as the url given to
+    # `cockroach sql` must start with the `postrges://` scheme.
+    def assert_correct_url(config, expected_url)
+      db_config = ActiveRecord::DatabaseConfigurations::HashConfig.new("default_env", "primary", config)
+      task_chief = ActiveRecord::ConnectionAdapters::CockroachDB::DatabaseTasks.new(db_config)
+      generated_url = task_chief.send(:cockroach_env)["COCKROACH_URL"]
+
+      conf_from_generated_url = ActiveRecord::Base.
+        configurations.
+        resolve(generated_url).
+        configuration_hash
+      assert_equal expected_url, generated_url
+      assert_equal config.except(:adapter), conf_from_generated_url.except(:adapter)
+    end
+  end
 end

test/support/rake_helpers.rb

@@ -17,8 +17,7 @@ def test_files
     ar_test_files = ar_test_files.
       split(',').
       map { |file| File.join ARTest::CockroachDB.root_activerecord, file.strip }.
-      then { _1.prepend(COCKROACHDB_TEST_HELPER) unless _1.empty? }.
-      prepend(COCKROACHDB_TEST_HELPER)
+      tap { _1.prepend(COCKROACHDB_TEST_HELPER) unless _1.empty? }
 
     cr_test_files = cr_test_files.split(',').map(&:strip)
 
@@ -28,7 +27,7 @@ def test_files
   def all_test_files
     activerecord_test_files = Dir.
       glob("#{ARTest::CockroachDB.root_activerecord}/test/cases/**/*_test.rb").
-      reject { _1.match?(%r(/adapters/(?:mysql2|sqlite3)/) }.
+      reject { _1.match?(%r(/adapters/(?:mysql2|sqlite3)/)) }.
       prepend(COCKROACHDB_TEST_HELPER)
 
     cockroachdb_test_files = Dir.glob('test/cases/**/*_test.rb')