diff --git a/lib/rom/sql/commands/create.rb b/lib/rom/sql/commands/create.rb
index 44627995..e548d233 100644
--- a/lib/rom/sql/commands/create.rb
+++ b/lib/rom/sql/commands/create.rb
@@ -46,7 +46,7 @@ def finalize(tuples, *)
         # @api private
         def insert(tuples)
           pks = tuples.map { |tuple| relation.insert(tuple) }
-          relation.where(relation.primary_key => pks).to_a
+          relation.dataset.where(relation.primary_key => pks).to_a
         end
 
         # Executes multi_insert statement and returns inserted tuples
@@ -54,7 +54,7 @@ def insert(tuples)
         # @api private
         def multi_insert(tuples)
           pks = relation.multi_insert(tuples, return: :primary_key)
-          relation.where(relation.primary_key => pks).to_a
+          relation.dataset.where(relation.primary_key => pks).to_a
         end
 
         # Yields tuples for insertion or return an enumerator
diff --git a/spec/integration/commands/create_spec.rb b/spec/integration/commands/create_spec.rb
index 94da1975..a922d734 100644
--- a/spec/integration/commands/create_spec.rb
+++ b/spec/integration/commands/create_spec.rb
@@ -187,6 +187,22 @@ def self.[](input)
       }.to raise_error(ROM::SQL::NotNullConstraintError)
     end
 
+    context 'with json notes' do
+      include_context 'json_notes'
+
+      before do
+        conf.commands(:json_notes) do
+          define(:create)
+        end
+      end
+
+      it 'writes and reads back custom type' do
+        json_notes = commands[:json_notes]
+
+        expect(json_notes[:create].call(note: 'this is my note')).to eq([{id: 1, note: 'this is my note'}])
+      end
+    end
+
     # Because Oracle doesn't have boolean in SQL
     unless metadata[:oracle]
       context 'with puppies' do
diff --git a/spec/integration/commands/update_spec.rb b/spec/integration/commands/update_spec.rb
index cd6a0553..a03d6263 100644
--- a/spec/integration/commands/update_spec.rb
+++ b/spec/integration/commands/update_spec.rb
@@ -110,6 +110,25 @@ def by_name(name)
           { id: 2, name: 'Josie' }
         ])
       end
+
+      context "with json notes" do
+        include_context "json_notes"
+
+        before do
+          conf.commands(:json_notes) do
+            define(:update)
+          end
+        end
+
+        let(:json_notes) { container.relations[:json_notes] }
+
+        it "writes and reads back custom type" do
+          note_id = conn[:json_notes].insert(note: "note version 1")
+          result = json_notes.by_pk(note_id).command(:update).call(note: "note version 2")
+
+          expect(result).to eq([{id: 1, note: "note version 2"}])
+        end
+      end
     end
   end
 end
diff --git a/spec/shared/json_notes.rb b/spec/shared/json_notes.rb
new file mode 100644
index 00000000..35232038
--- /dev/null
+++ b/spec/shared/json_notes.rb
@@ -0,0 +1,23 @@
+# frozen_string_literal: true
+
+RSpec.shared_context "json_notes" do
+  before do
+    inferrable_relations.concat %i[json_notes]
+  end
+
+  before do |_example|
+    conn.create_table :json_notes do
+      primary_key :id
+      String :note
+    end
+
+    write_type = Dry.Types.Constructor(String) { |value| JSON.dump({content: value}) }
+    read_type = Dry.Types.Constructor(String) { |value| JSON.parse(value)["content"] }
+
+    conf.relation(:json_notes) do
+      schema(infer: true) do
+        attribute :note, write_type, read: read_type
+      end
+    end
+  end
+end