diff --git a/lib/activerecord-multi-tenant/model_extensions.rb b/lib/activerecord-multi-tenant/model_extensions.rb
index f0acaa2..44fce9e 100644
--- a/lib/activerecord-multi-tenant/model_extensions.rb
+++ b/lib/activerecord-multi-tenant/model_extensions.rb
@@ -42,20 +42,15 @@ def partition_key
                                         .try(:instance_variable_get, :@partition_key)
           end
 
-          # Avoid primary_key errors when using composite primary keys (e.g. id, tenant_id)
-          def primary_key
-            if defined?(PRIMARY_KEY_NOT_SET) ? !PRIMARY_KEY_NOT_SET.equal?(@primary_key) : @primary_key
-              return @primary_key
-            end
-
+          def reset_primary_key
             primary_object_keys = Array.wrap(connection.schema_cache.primary_keys(table_name)) - [partition_key]
 
-            @primary_key = if primary_object_keys.size == 1
-                             primary_object_keys.first
-                           elsif table_name &&
-                                 connection.schema_cache.columns_hash(table_name).include?(DEFAULT_ID_FIELD)
-                             DEFAULT_ID_FIELD
-                           end
+            self.primary_key = if primary_object_keys.size == 1
+                                 primary_object_keys.first
+                               elsif table_name &&
+                                     connection.schema_cache.columns_hash(table_name).include?(DEFAULT_ID_FIELD)
+                                 DEFAULT_ID_FIELD
+                               end
           end
 
           def inherited(subclass)