Add JSONB store support to :has_one assocations

.gitignore

@@ -2,3 +2,6 @@ Gemfile.lock
 .bundle/
 log/*.log
 pkg/
+.DS_Store
+
+spec/config/database.yml

Gemfile

@@ -2,7 +2,9 @@ source 'https://rubygems.org'
 
 gemspec
 
-group :test do
-  gem 'pry'
-  gem 'rubocop'
-end
+gem 'pg'
+
+gem 'database_cleaner'
+gem 'factory_bot', '~> 4.8.0'
+gem 'pry'
+gem 'rubocop'

activerecord-jsonb-associations.gemspec

@@ -23,5 +23,6 @@ Gem::Specification.new do |spec|
   spec.required_ruby_version = '~> 2.0'
 
   spec.add_dependency 'activerecord', '~> 5.1.0'
+
   spec.add_development_dependency 'rspec', '~> 3.7.0'
 end

bin/setup

@@ -1,6 +1,31 @@
-#!/usr/bin/env bash
-set -euo pipefail
-IFS=$'\n\t'
-set -vx
+#!/usr/bin/env ruby
 
-bundle install
+require 'pathname'
+require 'fileutils'
+
+def system!(*args)
+  system(*args) || abort("\n== Command #{args} failed ==")
+end
+
+FileUtils.chdir(Pathname.new(File.expand_path('../../', __FILE__))) do
+  puts '== Installing dependencies =='
+  system! 'gem install bundler --conservative'
+  system('bundle check') || system!('bundle install')
+
+  puts "\n== Creating db config file =="
+  unless File.exist? 'spec/config/database.yml'
+    db_config = File.read('spec/config/database.yml.sample')
+
+    print 'Enter your PostgreSQL username: '
+    db_username = gets.chomp
+
+    File.write(
+      'spec/config/database.yml',
+      db_config.gsub('%USERNAME%', db_username)
+    )
+  end
+
+  puts "\n== Creating test database =="
+  system! 'dropdb activerecord_jsonb_associations_test'
+  system! 'createdb activerecord_jsonb_associations_test'
+end

lib/activerecord/jsonb/associations.rb

@@ -1,6 +1,11 @@
 require 'active_record'
+require 'pry'
 
 require 'activerecord/jsonb/associations/builder/belongs_to'
+require 'activerecord/jsonb/associations/builder/has_one'
+require 'activerecord/jsonb/associations/belongs_to_association'
+require 'activerecord/jsonb/associations/has_one_association'
+require 'activerecord/jsonb/associations/association_scope'
 
 module ActiveRecord #:nodoc:
   module JSONB #:nodoc:
@@ -13,4 +18,20 @@ module Associations #:nodoc:
   ::ActiveRecord::Associations::Builder::BelongsTo.extend(
     ActiveRecord::JSONB::Associations::Builder::BelongsTo
   )
+
+  ::ActiveRecord::Associations::Builder::HasOne.extend(
+    ActiveRecord::JSONB::Associations::Builder::HasOne
+  )
+
+  ::ActiveRecord::Associations::BelongsToAssociation.prepend(
+    ActiveRecord::JSONB::Associations::BelongsToAssociation
+  )
+
+  ::ActiveRecord::Associations::HasOneAssociation.prepend(
+    ActiveRecord::JSONB::Associations::HasOneAssociation
+  )
+
+  ::ActiveRecord::Associations::AssociationScope.prepend(
+    ActiveRecord::JSONB::Associations::AssociationScope
+  )
 end

lib/activerecord/jsonb/associations/association_scope.rb

@@ -0,0 +1,34 @@
+module ActiveRecord
+  module JSONB
+    module Associations
+      module AssociationScope #:nodoc:
+        def last_chain_scope(scope, table, reflection, owner)
+          reflection = reflection.instance_variable_get(:@reflection)
+
+          if reflection.options.key?(:foreign_store)
+            join_keys = reflection.join_keys
+            value = transform_value(owner[join_keys.foreign_key])
+
+            if value.is_a?(Integer)
+              scope = apply_jsonb_scope(
+                scope, table, reflection.options[:foreign_store],
+                join_keys.key, value
+              )
+            end
+
+            scope
+          else
+            super
+          end
+        end
+
+        def apply_jsonb_scope(scope, table, jsonb_column, key, value)
+          scope.where!(
+            "(#{table.name}.#{jsonb_column}->>'#{key}')::int = :id",
+            id: value
+          )
+        end
+      end
+    end
+  end
+end

lib/activerecord/jsonb/associations/belongs_to_association.rb

@@ -0,0 +1,18 @@
+module ActiveRecord
+  module JSONB
+    module Associations
+      module BelongsToAssociation #:nodoc:
+        def replace_keys(record)
+          if reflection.options.key?(:store)
+            owner[reflection.options[:store]][reflection.foreign_key] =
+              record._read_attribute(
+                reflection.association_primary_key(record.class)
+              )
+          else
+            super
+          end
+        end
+      end
+    end
+  end
+end

lib/activerecord/jsonb/associations/builder/belongs_to.rb

@@ -6,6 +6,29 @@ module BelongsTo #:nodoc:
           def valid_options(options)
             super + [:store]
           end
+
+          def define_accessors(mixin, reflection)
+            if reflection.options.key?(:store)
+              mixin.attribute reflection.foreign_key, :integer
+              add_association_accessor_methods(mixin, reflection)
+            end
+
+            super
+          end
+
+          def add_association_accessor_methods(mixin, reflection)
+            foreign_key = reflection.foreign_key.to_s
+
+            mixin.class_eval <<-CODE, __FILE__, __LINE__ + 1
+              def #{foreign_key}=(value)
+                #{reflection.options[:store]}['#{foreign_key}'] = value
+              end
+
+              def #{foreign_key}
+                #{reflection.options[:store]}['#{foreign_key}']
+              end
+            CODE
+          end
         end
       end
     end

lib/activerecord/jsonb/associations/builder/has_one.rb

@@ -0,0 +1,13 @@
+module ActiveRecord
+  module JSONB
+    module Associations
+      module Builder
+        module HasOne #:nodoc:
+          def valid_options(options)
+            super + [:foreign_store]
+          end
+        end
+      end
+    end
+  end
+end

lib/activerecord/jsonb/associations/has_one_association.rb

@@ -0,0 +1,30 @@
+module ActiveRecord
+  module JSONB
+    module Associations
+      module HasOneAssociation #:nodoc:
+        def creation_attributes
+          if reflection.options.key?(:foreign_store)
+            attributes = {}
+            jsonb_store = reflection.options[:foreign_store]
+            attributes[jsonb_store] ||= {}
+            attributes[jsonb_store][reflection.foreign_key] =
+              owner[reflection.active_record_primary_key]
+
+            attributes
+          else
+            super
+          end
+        end
+
+        def create_scope
+          super.tap do |scope|
+            next unless options.key?(:foreign_store)
+            scope[options[:foreign_store].to_s] ||= {}
+            scope[options[:foreign_store].to_s][reflection.foreign_key] =
+              owner[reflection.active_record_primary_key]
+          end
+        end
+      end
+    end
+  end
+end

spec/config/database.yml.sample

@@ -0,0 +1,2 @@
+pg:
+  username: "%USERNAME%"

spec/integration/has_one_integration_spec.rb

@@ -0,0 +1,124 @@
+RSpec.shared_examples ':has_one with JSONB store' do
+  let(:child_name) { child_model.model_name.element }
+
+  describe '#association' do
+    before do
+      child_model.update
+    end
+  end
+
+  describe '#association' do
+    before do
+      child_model.send "#{store}=", foreign_key => parent_model.id
+      child_model.save
+    end
+
+    it 'properly loads association from parent model' do
+      expect(parent_model.reload.send(child_name)).to eq(child_model)
+    end
+  end
+
+  describe '#association=' do
+    before do
+      parent_model.send "#{child_name}=", child_model
+    end
+
+    it 'sets and persists foreign key on child model' do
+      expect(
+        child_model.reload.send(store)
+      ).to eq(foreign_key.to_s => parent_model.id)
+    end
+  end
+
+  describe 'association_id' do
+    before do
+      child_model.send(store)[foreign_key.to_s] = parent_model.id
+    end
+
+    it 'reads foreign id from specified :store column by foreign key' do
+      expect(child_model.send(foreign_key)).to eq parent_model.id
+    end
+  end
+
+  describe '#association_id=' do
+    before do
+      child_model.send "#{foreign_key}=", parent_model.id
+    end
+
+    it 'sets foreign id in specified :store column as hash item' do
+      expect(child_model.send(store)[foreign_key.to_s]).to eq(parent_model.id)
+    end
+  end
+
+  describe '#build_association' do
+    let(:built_association) do
+      parent_model.send "build_#{child_name}"
+    end
+
+    it 'sets foreign key on child model' do
+      expect(
+        built_association.send(store)
+      ).to eq(foreign_key.to_s => parent_model.id)
+    end
+  end
+
+  describe '#create_association' do
+    let(:created_association) do
+      parent_model.send "create_#{child_name}"
+    end
+
+    it 'sets and persists foreign key on child model' do
+      expect(
+        created_association.reload.send(store)
+      ).to eq(foreign_key.to_s => parent_model.id)
+    end
+  end
+
+  describe '#reload_association' do
+    before do
+      parent_model.send "#{child_name}=", child_model
+    end
+
+    it 'reloads the association' do
+      expect(parent_model.send("reload_#{child_name}")).to eq(child_model)
+    end
+  end
+end
+
+RSpec.describe ':has_one' do
+  context 'regular association' do
+    let(:parent_model) { User.create }
+    let(:child_model) { Profile.new }
+
+    describe '#create_association' do
+      let(:created_association) do
+        parent_model.send "create_profile"
+      end
+
+      it 'sets and persists foreign key on child model' do
+        expect(
+          created_association.reload.user_id
+        ).to eq(parent_model.id)
+      end
+    end
+  end
+
+  context 'association with :store option set on child model' do
+    let(:child_model) { Account.new }
+    let(:store) { :extra }
+
+    context 'with default options' do
+      let(:parent_model) { User.create }
+      let(:foreign_key) { :user_id }
+
+      include_examples ':has_one with JSONB store'
+    end
+
+    context 'with non-default :options' do
+      let(:parent_model) { GoodsSupplier.create }
+      let(:foreign_key) { :supplier_id }
+
+      include_examples ':has_one with JSONB store'
+    end
+  end
+end

spec/spec_helper.rb

@@ -1,5 +1,24 @@
 require 'bundler/setup'
+require 'database_cleaner'
+require 'factory_bot'
 
 require 'activerecord/jsonb/associations'
 
-ActiveRecord::Base
+require 'support/schema'
+require 'support/models'
+require 'support/factories'
+
+RSpec.configure do |config|
+  config.include FactoryBot::Syntax::Methods
+
+  config.before(:suite) do
+    DatabaseCleaner.strategy = :transaction
+    DatabaseCleaner.clean_with(:truncation)
+  end
+
+  config.around(:each) do |example|
+    DatabaseCleaner.cleaning do
+      example.run
+    end
+  end
+end