Fix deteministic fetch when table has a composite primary key (#936)

Even though Rails does not support tables with primary keys, it still allows
those tables to be defined without any problems.

The SQL Server adapter will always try to retrieve records in a deterministic
way but it currently fails if the table has a composite primary key, which
shouldn't happen.

This PR fixes how the adapter determines which column (generally the PK) is
used as the default order when multiple primary keys are detected.

It first checks if one of the primary keys is also an identity column. If
there's an identity column, it will choose it as the "main primary key",
otherwise it will just fallback to the first of the primary keys.

Author: wpolicarpo

CHANGELOG.md

@@ -5,6 +5,7 @@
 - [#933](https://github.com/rails-sqlserver/activerecord-sqlserver-adapter/pull/933) Conditionally apply SQL Server monkey patches to ActiveRecord so that it is safe to use this gem alongside other database adapters (e.g. PostgreSQL) in a multi-database Rails app
 - [#935](https://github.com/rails-sqlserver/activerecord-sqlserver-adapter/pull/935) Fix schema cache generation
   (**breaking change**)
+- [#936](https://github.com/rails-sqlserver/activerecord-sqlserver-adapter/pull/936) Fix deteministic fetch when table has a composite primary key
 
 #### Changed
 

lib/arel/visitors/sqlserver.rb

@@ -296,8 +296,21 @@ def table_From_Statement(o)
       def primary_Key_From_Table(t)
         return unless t
 
-        column_name = @connection.schema_cache.primary_keys(t.name) ||
-          @connection.schema_cache.columns_hash(t.name).first.try(:second).try(:name)
+        primary_keys = @connection.schema_cache.primary_keys(t.name)
+        column_name = nil
+
+        case primary_keys
+        when NilClass
+          column_name = @connection.schema_cache.columns_hash(t.name).first.try(:second).try(:name)
+        when String
+          column_name = primary_keys
+        when Array
+          candidate_columns = @connection.schema_cache.columns_hash(t.name).slice(*primary_keys).values
+          candidate_column = candidate_columns.find(&:is_identity?)
+          candidate_column ||= candidate_columns.first
+          column_name = candidate_column.try(:name)
+        end
+
         column_name ? t[column_name] : nil
       end
 

test/cases/fetch_test_sqlserver.rb

@@ -49,3 +49,21 @@ def create_10_books
     @books = (1..10).map { |i| Book.create! name: "Name-#{i}" }
   end
 end
+
+class DeterministicFetchWithCompositePkTestSQLServer < ActiveRecord::TestCase
+  it "orders by the identity column if table has one" do
+    SSCompositePkWithIdentity.delete_all
+    SSCompositePkWithIdentity.create(pk_col_two: 2)
+    SSCompositePkWithIdentity.create(pk_col_two: 1)
+
+    _(SSCompositePkWithIdentity.take(1).map(&:pk_col_two)).must_equal [2]
+  end
+
+  it "orders by the first column if table has no identity column" do
+    SSCompositePkWithoutIdentity.delete_all
+    SSCompositePkWithoutIdentity.create(pk_col_one: 2, pk_col_two: 2)
+    SSCompositePkWithoutIdentity.create(pk_col_one: 1, pk_col_two: 1)
+
+    _(SSCompositePkWithoutIdentity.take(1).map(&:pk_col_two)).must_equal [1]
+  end
+end

test/models/sqlserver/composite_pk.rb

@@ -0,0 +1,9 @@
+# frozen_string_literal: true
+
+class SSCompositePkWithoutIdentity < ActiveRecord::Base
+  self.table_name = :sst_composite_without_identity
+end
+
+class SSCompositePkWithIdentity < ActiveRecord::Base
+  self.table_name = :sst_composite_with_identity
+end

test/schema/sqlserver_specific_schema.rb

@@ -294,4 +294,22 @@
       CONSTRAINT PK_UNIQUE_KEY PRIMARY KEY (id)
     );
   SQLSERVERUNIQUEKEYS
+
+  execute "IF EXISTS(SELECT * FROM INFORMATION_SCHEMA.TABLES WHERE TABLE_NAME = 'sst_composite_without_identity') DROP TABLE sst_composite_without_identity"
+  execute <<-COMPOSITE_WITHOUT_IDENTITY
+    CREATE TABLE sst_composite_without_identity (
+      pk_col_one int NOT NULL,
+      pk_col_two int NOT NULL,
+      CONSTRAINT PK_sst_composite_without_identity PRIMARY KEY (pk_col_one, pk_col_two)
+    );
+  COMPOSITE_WITHOUT_IDENTITY
+
+  execute "IF EXISTS(SELECT * FROM INFORMATION_SCHEMA.TABLES WHERE TABLE_NAME = 'sst_composite_with_identity') DROP TABLE sst_composite_with_identity"
+  execute <<-COMPOSITE_WITH_IDENTITY
+    CREATE TABLE sst_composite_with_identity (
+      pk_col_one int IDENTITY NOT NULL,
+      pk_col_two int NOT NULL,
+      CONSTRAINT PK_sst_composite_with_identity PRIMARY KEY (pk_col_one, pk_col_two)
+    );
+  COMPOSITE_WITH_IDENTITY
 end