,

babyfish-ct · Oct 30, 2024 · 044789e · 044789e
1 parent 3ca6ba8
commit 044789e
Show file tree

Hide file tree

Showing 5 changed files with 65 additions and 4 deletions.
diff --git a/project/jimmer-sql/src/main/java/org/babyfish/jimmer/sql/ast/impl/mutation/Operator.java b/project/jimmer-sql/src/main/java/org/babyfish/jimmer/sql/ast/impl/mutation/Operator.java
@@ -62,6 +62,7 @@ public void insert(Batch<DraftSpi> batch) {
                 defaultGetters.add(getter);
             }
         }
+        IdentityIdGenerator identityIdGenerator = null;
         SequenceIdGenerator sequenceIdGenerator = null;
         UserIdGenerator<?> userIdGenerator = null;
         if (batch.shape().getIdGetters().isEmpty()) {
@@ -70,7 +71,9 @@ public void insert(Batch<DraftSpi> batch) {
                 sequenceIdGenerator = (SequenceIdGenerator) idGenerator;
             } else if (idGenerator instanceof UserIdGenerator<?>) {
                 userIdGenerator = (UserIdGenerator<?>) idGenerator;
-            } else if (!(idGenerator instanceof IdentityIdGenerator)) {
+            } else if (idGenerator instanceof IdentityIdGenerator) {
+                identityIdGenerator = (IdentityIdGenerator) idGenerator;
+            } else {
                 ctx.throwIllegalIdGenerator(
                         "In order to insert object without id, the id generator must be identity or sequence"
                 );
@@ -125,6 +128,15 @@ public void insert(Batch<DraftSpi> batch) {
             builder.separator().defaultVariable(defaultGetter);
         }
         builder.leave();
+        if ((identityIdGenerator != null || sequenceIdGenerator != null) &&
+        sqlClient.getDialect().isInsertedIdReturningRequired()) {
+            builder.sql(" returning ")
+                    .sql(
+                            batch.shape().getType().getIdProp()
+                                    .<SingleColumn>getStorage(sqlClient.getMetadataStrategy())
+                                    .getName()
+                    );
+        }
 
         MutationTrigger trigger = ctx.trigger;
         if (trigger != null) {

diff --git a/project/jimmer-sql/src/main/java/org/babyfish/jimmer/sql/dialect/Dialect.java b/project/jimmer-sql/src/main/java/org/babyfish/jimmer/sql/dialect/Dialect.java
@@ -115,6 +115,10 @@ default boolean isIdFetchableByKeyUpdate() {
         return false;
     }
 
+    default boolean isInsertedIdReturningRequired() {
+        return false;
+    }
+
     default boolean isUpsertSupported() {
         return false;
     }

diff --git a/project/jimmer-sql/src/main/java/org/babyfish/jimmer/sql/dialect/PostgresDialect.java b/project/jimmer-sql/src/main/java/org/babyfish/jimmer/sql/dialect/PostgresDialect.java
@@ -153,6 +153,11 @@ public boolean isIdFetchableByKeyUpdate() {
         return true;
     }
 
+    @Override
+    public boolean isInsertedIdReturningRequired() {
+        return true;
+    }
+
     @Override
     public boolean isUpsertSupported() {
         return true;

diff --git a/project/jimmer-sql/src/test/java/org/babyfish/jimmer/sql/mutation/IdentityTest.java b/project/jimmer-sql/src/test/java/org/babyfish/jimmer/sql/mutation/IdentityTest.java
@@ -441,12 +441,18 @@ public void testInsertPostgres() {
                         .setAssociatedModeAll(AssociatedSaveMode.APPEND),
                 ctx -> {
                     ctx.statement(it -> {
-                        it.sql("insert into DEPARTMENT(NAME, DELETED_MILLIS) values(?, ?)");
+                        it.sql(
+                                "insert into DEPARTMENT(NAME, DELETED_MILLIS) " +
+                                        "values(?, ?) returning ID"
+                        );
                         it.batchVariables(0, "Develop", 0L);
                         it.batchVariables(1, "Sales", 0L);
                     });
                     ctx.statement(it -> {
-                        it.sql("insert into EMPLOYEE(NAME, DELETED_MILLIS, DEPARTMENT_ID) values(?, ?, ?)");
+                        it.sql(
+                                "insert into EMPLOYEE(NAME, DELETED_MILLIS, DEPARTMENT_ID) " +
+                                        "values(?, ?, ?) returning ID"
+                        );
                         it.batchVariables(0, "Jacob", 0L, 100L);
                         it.batchVariables(1, "Tania", 0L, 100L);
                         it.batchVariables(2, "Oakes", 0L, 101L);

diff --git a/project/jimmer-sql/src/test/resources/database-postgres.sql b/project/jimmer-sql/src/test/resources/database-postgres.sql
@@ -181,4 +181,38 @@ create table pg_json_wrapper(
     json_3 jsonb,
     json_4 jsonb,
     json_5 jsonb
-)
+);
+
+create table department(
+    name text not null,
+    deleted_millis bigint not null,
+    id bigint not null generated by default as identity(start with 1 increment by 1)
+);
+alter table department
+    add constraint pk_department
+        primary key(id);
+alter table department
+    add constraint uq_department
+        unique(name, deleted_millis);
+
+create table employee(
+    name text not null,
+    deleted_millis bigint not null,
+    department_id bigint,
+    id bigint not null generated by default as identity(start with 1 increment by 1)
+);
+alter table employee
+    add constraint pk_employee
+        primary key(id);
+alter table employee
+    add constraint uq_employee
+        unique(name, deleted_millis);
+alter table employee
+    add constraint fk_employee__department
+        foreign key(department_id)
+            references department(id);
+
+insert into department(id, name, deleted_millis)
+    values(1, 'Market', 0);
+insert into employee(id, name, department_id, deleted_millis)
+    values(1, 'Sam', 1, 0), (2, 'Jessica', 1, 0);