diff --git a/lib/rom/sql/commands/create.rb b/lib/rom/sql/commands/create.rb
index e20ffa6a..dddf8172 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 b59e89dd..e36469aa 100644
--- a/spec/integration/commands/create_spec.rb
+++ b/spec/integration/commands/create_spec.rb
@@ -178,6 +178,22 @@ def self.[](input)
       ])
     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
+
     it "re-raises not-null constraint violation error" do
       expect {
         create_user.call(name: nil)
diff --git a/spec/integration/commands/update_spec.rb b/spec/integration/commands/update_spec.rb
index 24838021..7bc23269 100644
--- a/spec/integration/commands/update_spec.rb
+++ b/spec/integration/commands/update_spec.rb
@@ -109,6 +109,25 @@ def by_name(name)
 
         expect(result).to eq([{id: 1, name: "Josie"}, {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