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 index 44cb2e7255..6fdb9d93bf 100644 --- 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 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 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 batch) { builder.separator().defaultVariable(defaultGetter); } builder.leave(); + if ((identityIdGenerator != null || sequenceIdGenerator != null) && + sqlClient.getDialect().isInsertedIdReturningRequired()) { + builder.sql(" returning ") + .sql( + batch.shape().getType().getIdProp() + .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 index 362101dbf0..a5ce65d44e 100644 --- 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 index 5c6c5c859f..2e8e059e9b 100644 --- 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 index 4797c5c7c1..a45fdba7f1 100644 --- 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 index 9b6bbc947d..3f5f1cba7b 100644 --- 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 -) \ No newline at end of file +); + +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); \ No newline at end of file