activerecord-cockroachdb-adapter 7.0.2 → 7.1.0

data/lib/active_record/connection_adapters/cockroachdb/database_statements.rb

@@ -2,21 +2,6 @@ module ActiveRecord
   module ConnectionAdapters
     module CockroachDB
       module DatabaseStatements
-        # Since CockroachDB will run all transactions with serializable isolation,
-        # READ UNCOMMITTED, READ COMMITTED, and REPEATABLE READ are all aliases
-        # for SERIALIZABLE. This lets the adapter support all isolation levels,
-        # but READ UNCOMMITTED has been removed from this list because the
-        # ActiveRecord transaction isolation test fails for READ UNCOMMITTED.
-        # See https://www.cockroachlabs.com/docs/v19.2/transactions.html#isolation-levels
-        def transaction_isolation_levels
-          {
-            read_committed:   "READ COMMITTED",
-            repeatable_read:  "REPEATABLE READ",
-            serializable:     "SERIALIZABLE",
-            read_uncommitted: "READ UNCOMMITTED"
-          }
-        end
-
         # Overridden to avoid using transactions for schema creation.
         def insert_fixtures_set(fixture_set, tables_to_delete = [])
           fixture_inserts = build_fixture_statements(fixture_set)
@@ -29,74 +14,6 @@ module ActiveRecord
             end
           end
         end
-
-        private
-          def execute_batch(statements, name = nil)
-            statements.each do |statement|
-              execute(statement, name)
-            end
-          end
-
-          DEFAULT_INSERT_VALUE = Arel.sql("DEFAULT").freeze
-          private_constant :DEFAULT_INSERT_VALUE
-
-          def default_insert_value(column)
-            DEFAULT_INSERT_VALUE
-          end
-
-          def build_fixture_sql(fixtures, table_name)
-            columns = schema_cache.columns_hash(table_name)
-
-            values_list = fixtures.map do |fixture|
-              fixture = fixture.stringify_keys
-
-              unknown_columns = fixture.keys - columns.keys
-              if unknown_columns.any?
-                raise Fixture::FixtureError, %(table "#{table_name}" has no columns named #{unknown_columns.map(&:inspect).join(', ')}.)
-              end
-
-              columns.map do |name, column|
-                if fixture.key?(name)
-                  type = lookup_cast_type_from_column(column)
-                  with_yaml_fallback(type.serialize(fixture[name]))
-                else
-                  default_insert_value(column)
-                end
-              end
-            end
-
-            table = Arel::Table.new(table_name)
-            manager = Arel::InsertManager.new
-            manager.into(table)
-
-            if values_list.size == 1
-              values = values_list.shift
-              new_values = []
-              columns.each_key.with_index { |column, i|
-                unless values[i].equal?(DEFAULT_INSERT_VALUE)
-                  new_values << values[i]
-                  manager.columns << table[column]
-                end
-              }
-              values_list << new_values
-            else
-              columns.each_key { |column| manager.columns << table[column] }
-            end
-
-            manager.values = manager.create_values_list(values_list)
-            manager.to_sql
-          end
-
-          def build_fixture_statements(fixture_set)
-            fixture_set.map do |table_name, fixtures|
-              next if fixtures.empty?
-              build_fixture_sql(fixtures, table_name)
-            end.compact
-          end
-
-          def with_multi_statements
-            yield
-          end
       end
     end
   end

data/lib/active_record/connection_adapters/cockroachdb/database_tasks.rb

@@ -5,12 +5,96 @@ module ActiveRecord
     module CockroachDB
       class DatabaseTasks < ActiveRecord::Tasks::PostgreSQLDatabaseTasks
         def structure_dump(filename, extra_flags=nil)
-          raise "db:structure:dump 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
+
+          # "See https://github.com/cockroachdb/cockroach/issues/26443."
+          search_path =
+            case ActiveRecord.dump_schemas
+            when :schema_search_path
+              configuration_hash[:schema_search_path]
+            when :all
+              nil
+            when String
+              ActiveRecord.dump_schemas
+            end
+
+          conn = ActiveRecord::Base.connection
+          begin
+            old_search_path = conn.schema_search_path
+            conn.schema_search_path = search_path
+            File.open(filename, "w") do |file|
+              %w(SCHEMAS TYPES).each do |object_kind|
+                ActiveRecord::Base.connection.execute("SHOW CREATE ALL #{object_kind}").each_row { file.puts _1 }
+              end
+
+              ignore_tables = ActiveRecord::SchemaDumper.ignore_tables.to_set
+
+              conn.execute("SHOW CREATE ALL TABLES").each_row do |(sql)|
+                if sql.start_with?("CREATE")
+                  table_name = sql[/CREATE TABLE (?:.*?\.)?\"?(.*?)[\" ]/, 1]
+                  next if ignore_tables.member?(table_name)
+                elsif sql.start_with?("ALTER")
+                  table_name = sql[/ALTER TABLE (?:.*?\.)?\"?(.*?)[\" ]/, 1]
+                  ref_table_name = sql[/REFERENCES (?:.*?\.)?\"?(.*?)[\" ]/, 1]
+                  next if ignore_tables.member?(table_name) || ignore_tables.member?(ref_table_name)
+                end
+
+                file.puts sql
+              end
+              file.puts "SET seach_path TO #{conn.schema_search_path};\n\n"
+            end
+          ensure
+            conn.schema_search_path = old_search_path
+          end
         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

data/lib/active_record/connection_adapters/cockroachdb/quoting.rb

@@ -19,6 +19,12 @@ module ActiveRecord
         # converting to WKB, so this does it automatically.
         def quote(value)
           if value.is_a?(Numeric)
+            # NOTE: The fact that integers are quoted is important and helps
+            # mitigate a potential vulnerability.
+            #
+            # See
+            # - https://nvd.nist.gov/vuln/detail/CVE-2022-44566
+            # - https://github.com/cockroachdb/activerecord-cockroachdb-adapter/pull/280#discussion_r1288692977
             "'#{quote_string(value.to_s)}'"
           elsif RGeo::Feature::Geometry.check_type(value)
             "'#{RGeo::WKRep::WKBGenerator.new(hex_format: true, type_format: :ewkb, emit_ewkb_srid: true).generate(value)}'"

data/lib/active_record/connection_adapters/cockroachdb/referential_integrity.rb

@@ -11,6 +11,14 @@ module ActiveRecord
   module ConnectionAdapters
     module CockroachDB
       module ReferentialIntegrity
+        # CockroachDB will raise a `PG::ForeignKeyViolation` when re-enabling
+        # referential integrity (e.g: adding a foreign key with invalid data
+        # raises).
+        # So foreign keys should always be valid for that matter.
+        def check_all_foreign_keys_valid!
+          true
+        end
+
         def disable_referential_integrity
           foreign_keys = tables.map { |table| foreign_keys(table) }.flatten
 
@@ -31,16 +39,12 @@ module ActiveRecord
 
           begin
             foreign_keys.each do |foreign_key|
-              begin
-                add_foreign_key(foreign_key.from_table, foreign_key.to_table, **foreign_key.options)
-              rescue ActiveRecord::StatementInvalid => error
-                if error.cause.class == PG::DuplicateObject
-                  # This error is safe to ignore because the yielded caller
-                  # already re-added the foreign key constraint.
-                else
-                  raise error
-                end
-              end
+              # Avoid having PG:DuplicateObject error if a test is ran in transaction.
+              # TODO: verify that there is no cache issue related to running this (e.g: fk
+              #   still in cache but not in db)
+              next if foreign_key_exists?(foreign_key.from_table, name: foreign_key.options[:name])
+
+              add_foreign_key(foreign_key.from_table, foreign_key.to_table, **foreign_key.options)
             end
           ensure
             ActiveRecord::Base.table_name_prefix = old_prefix

data/lib/active_record/connection_adapters/cockroachdb/schema_statements.rb

@@ -36,16 +36,25 @@ module ActiveRecord
         # Modified version of the postgresql foreign_keys method.
         # Replaces t2.oid::regclass::text with t2.relname since this is
         # more efficient in CockroachDB.
+        # Also, CockroachDB does not append the schema name in relname,
+        # so we append it manually.
         def foreign_keys(table_name)
           scope = quoted_scope(table_name)
           fk_info = exec_query(<<~SQL, "SCHEMA")
-            SELECT t2.relname AS to_table, a1.attname AS column, a2.attname AS primary_key, c.conname AS name, c.confupdtype AS on_update, c.confdeltype AS on_delete, c.convalidated AS valid
+            SELECT CASE
+              WHEN n2.nspname = current_schema()
+              THEN ''
+              ELSE n2.nspname || '.'
+            END || t2.relname AS to_table,
+            a1.attname AS column, a2.attname AS primary_key, c.conname AS name, c.confupdtype AS on_update, c.confdeltype AS on_delete, c.convalidated AS valid, c.condeferrable AS deferrable, c.condeferred AS deferred,
+            c.conkey, c.confkey, c.conrelid, c.confrelid
             FROM pg_constraint c
             JOIN pg_class t1 ON c.conrelid = t1.oid
             JOIN pg_class t2 ON c.confrelid = t2.oid
             JOIN pg_attribute a1 ON a1.attnum = c.conkey[1] AND a1.attrelid = t1.oid
             JOIN pg_attribute a2 ON a2.attnum = c.confkey[1] AND a2.attrelid = t2.oid
             JOIN pg_namespace t3 ON c.connamespace = t3.oid
+            JOIN pg_namespace n2 ON t2.relnamespace = n2.oid
             WHERE c.contype = 'f'
               AND t1.relname = #{scope[:name]}
               AND t3.nspname = #{scope[:schema]}
@@ -53,17 +62,31 @@ module ActiveRecord
           SQL
 
           fk_info.map do |row|
+            to_table = PostgreSQL::Utils.unquote_identifier(row["to_table"])
+            conkey = row["conkey"].scan(/\d+/).map(&:to_i)
+            confkey = row["confkey"].scan(/\d+/).map(&:to_i)
+
+            if conkey.size > 1
+              column = column_names_from_column_numbers(row["conrelid"], conkey)
+              primary_key = column_names_from_column_numbers(row["confrelid"], confkey)
+            else
+              column = PostgreSQL::Utils.unquote_identifier(row["column"])
+              primary_key = row["primary_key"]
+            end
+
             options = {
-              column: row["column"],
+              column: column,
               name: row["name"],
-              primary_key: row["primary_key"]
+              primary_key: primary_key
             }
-
             options[:on_delete] = extract_foreign_key_action(row["on_delete"])
             options[:on_update] = extract_foreign_key_action(row["on_update"])
+            options[:deferrable] = extract_constraint_deferrable(row["deferrable"], row["deferred"])
+
             options[:validate] = row["valid"]
+            to_table = PostgreSQL::Utils.unquote_identifier(row["to_table"])
 
-            ForeignKeyDefinition.new(table_name, row["to_table"], options)
+            ForeignKeyDefinition.new(table_name, to_table, options)
           end
         end
 
@@ -76,21 +99,25 @@ module ActiveRecord
 
         # override
         # https://github.com/rails/rails/blob/6-0-stable/activerecord/lib/active_record/connection_adapters/postgresql/schema_statements.rb#L624
-        def new_column_from_field(table_name, field)
-          column_name, type, default, notnull, oid, fmod, collation, comment, generated, hidden = field
+        def new_column_from_field(table_name, field, _definition)
+          column_name, type, default, notnull, oid, fmod, collation, comment, identity, attgenerated, hidden = field
           type_metadata = fetch_type_metadata(column_name, type, oid.to_i, fmod.to_i)
           default_value = extract_value_from_default(default)
-          default_function = extract_default_function(default_value, default)
 
-          serial =
-            if (match = default_function&.match(/\Anextval\('"?(?<sequence_name>.+_(?<suffix>seq\d*))"?'::regclass\)\z/))
-              sequence_name_from_parts(table_name, column_name, match[:suffix]) == match[:sequence_name]
-            end
+          if attgenerated.present?
+            default_function = default
+          else
+            default_function = extract_default_function(default_value, default)
+          end
+
+          if match = default_function&.match(/\Anextval\('"?(?<sequence_name>.+_(?<suffix>seq\d*))"?'::regclass\)\z/)
+            serial = sequence_name_from_parts(table_name, column_name, match[:suffix]) == match[:sequence_name]
+          end
 
           # {:dimension=>2, :has_m=>false, :has_z=>false, :name=>"latlon", :srid=>0, :type=>"GEOMETRY"}
           spatial = spatial_column_info(table_name).get(column_name, type_metadata.sql_type)
 
-          PostgreSQL::Column.new(
+          CockroachDB::Column.new(
             column_name,
             default_value,
             type_metadata,
@@ -99,8 +126,9 @@ module ActiveRecord
             collation: collation,
             comment: comment.presence,
             serial: serial,
+            identity: identity.presence,
             spatial: spatial,
-            generated: generated,
+            generated: attgenerated,
             hidden: hidden
           )
         end
@@ -112,7 +140,7 @@ module ActiveRecord
         # since type alone is not enough to format the column.
         # Ex. type_to_sql(:geography, limit: "Point,4326")
         # => "geography(Point,4326)"
-        # 
+        #
         def type_to_sql(type, limit: nil, precision: nil, scale: nil, array: nil, **) # :nodoc:
           sql = \
             case type.to_s

data/lib/active_record/connection_adapters/cockroachdb/type.rb

@@ -1,12 +1,16 @@
 module ActiveRecord
   module Type
-    class << self
+    module CRDBExt
       # Return :postgresql instead of :cockroachdb for current_adapter_name so
       # we can continue using the ActiveRecord::Types defined in
       # PostgreSQLAdapter.
-      def adapter_name_from(_model)
-        :postgresql
+      def adapter_name_from(model)
+        name = super
+        return :postgresql if name == :cockroachdb
+
+        name
       end
     end
+    singleton_class.prepend CRDBExt
   end
 end