Revisit Insert methods

- introduce separate ifNotExists() property - fix bug when return type is entity and conditional insert succeeds - allow Optional return types
apache · Jun 21, 2019 · 2838543 · 2838543
1 parent 478535d
commit 2838543
Show file tree

Hide file tree

Showing 8 changed files with 89 additions and 24 deletions.
diff --git a/integration-tests/src/test/java/com/datastax/oss/driver/mapper/InsertEntityIT.java b/integration-tests/src/test/java/com/datastax/oss/driver/mapper/InsertEntityIT.java
@@ -24,11 +24,13 @@
 import com.datastax.oss.driver.api.testinfra.ccm.CcmRule;
 import com.datastax.oss.driver.api.testinfra.session.SessionRule;
 import com.datastax.oss.driver.categories.ParallelizableTests;
+import com.datastax.oss.driver.mapper.model.inventory.Dimensions;
 import com.datastax.oss.driver.mapper.model.inventory.InventoryFixtures;
 import com.datastax.oss.driver.mapper.model.inventory.InventoryMapper;
 import com.datastax.oss.driver.mapper.model.inventory.InventoryMapperBuilder;
 import com.datastax.oss.driver.mapper.model.inventory.Product;
 import com.datastax.oss.driver.mapper.model.inventory.ProductDao;
+import java.util.Optional;
 import org.junit.Before;
 import org.junit.BeforeClass;
 import org.junit.ClassRule;
@@ -109,4 +111,24 @@ public void should_insert_entity_with_custom_clause() {
     // Then
     assertThat(writeTime).isEqualTo(timestamp);
   }
+
+  @Test
+  public void should_insert_entity_if_not_exists() {
+    // Given
+    Product product = InventoryFixtures.FLAMETHROWER.entity;
+
+    // When
+    Optional<Product> maybeExisting = productDao.saveIfNotExists(product);
+
+    // Then
+    assertThat(maybeExisting).isEmpty();
+
+    // When
+    Product otherProduct =
+        new Product(product.getId(), "Other description", new Dimensions(1, 1, 1));
+    maybeExisting = productDao.saveIfNotExists(otherProduct);
+
+    // Then
+    assertThat(maybeExisting).contains(product);
+  }
 }
diff --git a/...ration-tests/src/test/java/com/datastax/oss/driver/mapper/model/inventory/ProductDao.java b/...ration-tests/src/test/java/com/datastax/oss/driver/mapper/model/inventory/ProductDao.java
@@ -63,8 +63,11 @@ public interface ProductDao {
   @Insert
   void save(Product product);
 
-  @Insert(customClause = "USING TIMESTAMP :timestamp")
-  Product saveWithBoundTimestamp(Product product, long timestamp);
+  @Insert(customUsingClause = "USING TIMESTAMP :timestamp")
+  void saveWithBoundTimestamp(Product product, long timestamp);
+
+  @Insert(ifNotExists = true)
+  Optional<Product> saveIfNotExists(Product product);
 
   @Select
   Product findById(UUID productId);

diff --git a/.../java/com/datastax/oss/driver/internal/mapper/processor/dao/DaoInsertMethodGenerator.java b/.../java/com/datastax/oss/driver/internal/mapper/processor/dao/DaoInsertMethodGenerator.java
@@ -17,7 +17,9 @@
 
 import static com.datastax.oss.driver.internal.mapper.processor.dao.ReturnTypeKind.ENTITY;
 import static com.datastax.oss.driver.internal.mapper.processor.dao.ReturnTypeKind.FUTURE_OF_ENTITY;
+import static com.datastax.oss.driver.internal.mapper.processor.dao.ReturnTypeKind.FUTURE_OF_OPTIONAL_ENTITY;
 import static com.datastax.oss.driver.internal.mapper.processor.dao.ReturnTypeKind.FUTURE_OF_VOID;
+import static com.datastax.oss.driver.internal.mapper.processor.dao.ReturnTypeKind.OPTIONAL_ENTITY;
 import static com.datastax.oss.driver.internal.mapper.processor.dao.ReturnTypeKind.VOID;
 
 import com.datastax.oss.driver.api.core.cql.BoundStatement;
@@ -39,7 +41,13 @@
 public class DaoInsertMethodGenerator extends DaoMethodGenerator {
 
   private static final EnumSet<ReturnTypeKind> SUPPORTED_RETURN_TYPES =
-      EnumSet.of(VOID, FUTURE_OF_VOID, ENTITY, FUTURE_OF_ENTITY);
+      EnumSet.of(
+          VOID,
+          FUTURE_OF_VOID,
+          ENTITY,
+          FUTURE_OF_ENTITY,
+          OPTIONAL_ENTITY,
+          FUTURE_OF_OPTIONAL_ENTITY);
 
   public DaoInsertMethodGenerator(
       ExecutableElement methodElement,
@@ -145,13 +153,19 @@ public Optional<MethodSpec> generate() {
   private void generatePrepareRequest(
       MethodSpec.Builder methodBuilder, String requestName, String helperFieldName) {
     methodBuilder.addCode(
-        "$[$1T $2L = $1T.newInstance($3L.insert().asCql()",
+        "$[$1T $2L = $1T.newInstance($3L.insert()",
         SimpleStatement.class,
         requestName,
         helperFieldName);
-    String customClause = methodElement.getAnnotation(Insert.class).customClause();
-    if (!customClause.isEmpty()) {
-      methodBuilder.addCode(" + $S", " " + customClause);
+    Insert annotation = methodElement.getAnnotation(Insert.class);
+    if (annotation.ifNotExists()) {
+      methodBuilder.addCode(".ifNotExists()");
+    }
+    methodBuilder.addCode(".asCql()");
+
+    String customUsingClause = annotation.customUsingClause();
+    if (!customUsingClause.isEmpty()) {
+      methodBuilder.addCode(" + $S", " " + customUsingClause);
     }
     methodBuilder.addCode(")$];\n");
   }

diff --git a/...tastax/oss/driver/internal/mapper/processor/entity/EntityHelperInsertMethodGenerator.java b/...tastax/oss/driver/internal/mapper/processor/entity/EntityHelperInsertMethodGenerator.java
@@ -16,9 +16,9 @@
 package com.datastax.oss.driver.internal.mapper.processor.entity;
 
 import com.datastax.oss.driver.api.core.CqlIdentifier;
-import com.datastax.oss.driver.api.querybuilder.BuildableQuery;
 import com.datastax.oss.driver.api.querybuilder.QueryBuilder;
 import com.datastax.oss.driver.api.querybuilder.insert.InsertInto;
+import com.datastax.oss.driver.api.querybuilder.insert.RegularInsert;
 import com.datastax.oss.driver.internal.mapper.processor.MethodGenerator;
 import com.datastax.oss.driver.internal.mapper.processor.ProcessorContext;
 import com.squareup.javapoet.MethodSpec;
@@ -42,7 +42,7 @@ public Optional<MethodSpec> generate() {
         MethodSpec.methodBuilder("insert")
             .addAnnotation(Override.class)
             .addModifiers(Modifier.PUBLIC)
-            .returns(BuildableQuery.class)
+            .returns(RegularInsert.class)
             .addStatement("$T keyspaceId = context.getKeyspaceId()", CqlIdentifier.class)
             .addStatement("$T tableId = context.getTableId()", CqlIdentifier.class)
             .beginControlFlow("if (tableId == null)")

diff --git a/...a/com/datastax/oss/driver/internal/mapper/processor/dao/DaoInsertMethodGeneratorTest.java b/...a/com/datastax/oss/driver/internal/mapper/processor/dao/DaoInsertMethodGeneratorTest.java
@@ -24,6 +24,7 @@
 import com.tngtech.java.junit.dataprovider.DataProvider;
 import com.tngtech.java.junit.dataprovider.DataProviderRunner;
 import com.tngtech.java.junit.dataprovider.UseDataProvider;
+import java.util.Optional;
 import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.CompletionStage;
 import javax.lang.model.element.Modifier;
@@ -133,6 +134,26 @@ public static Object[][] validSignatures() {
                     ClassName.get(CompletableFuture.class), ENTITY_CLASS_NAME))
             .build()
       },
+      // Returns an optional of the entity class, or a future thereof:
+      {
+        MethodSpec.methodBuilder("insert")
+            .addAnnotation(Insert.class)
+            .addModifiers(Modifier.PUBLIC, Modifier.ABSTRACT)
+            .addParameter(ParameterSpec.builder(ENTITY_CLASS_NAME, "entity").build())
+            .returns(ParameterizedTypeName.get(ClassName.get(Optional.class), ENTITY_CLASS_NAME))
+            .build()
+      },
+      {
+        MethodSpec.methodBuilder("insert")
+            .addAnnotation(Insert.class)
+            .addModifiers(Modifier.PUBLIC, Modifier.ABSTRACT)
+            .addParameter(ParameterSpec.builder(ENTITY_CLASS_NAME, "entity").build())
+            .returns(
+                ParameterizedTypeName.get(
+                    ClassName.get(CompletionStage.class),
+                    ParameterizedTypeName.get(ClassName.get(Optional.class), ENTITY_CLASS_NAME)))
+            .build()
+      },
       // Extra parameters in addition to the entity (to bind into the request):
       {
         MethodSpec.methodBuilder("insert")

diff --git a/mapper-runtime/src/main/java/com/datastax/oss/driver/api/mapper/annotations/Insert.java b/mapper-runtime/src/main/java/com/datastax/oss/driver/api/mapper/annotations/Insert.java
@@ -23,5 +23,7 @@
 @Target(ElementType.METHOD)
 @Retention(RetentionPolicy.CLASS)
 public @interface Insert {
-  String customClause() default "";
+  boolean ifNotExists() default false;
+
+  String customUsingClause() default "";
 }
diff --git a/mapper-runtime/src/main/java/com/datastax/oss/driver/api/mapper/entity/EntityHelper.java b/mapper-runtime/src/main/java/com/datastax/oss/driver/api/mapper/entity/EntityHelper.java
@@ -26,6 +26,7 @@
 import com.datastax.oss.driver.api.mapper.annotations.PartitionKey;
 import com.datastax.oss.driver.api.querybuilder.BuildableQuery;
 import com.datastax.oss.driver.api.querybuilder.delete.Delete;
+import com.datastax.oss.driver.api.querybuilder.insert.RegularInsert;
 import com.datastax.oss.driver.api.querybuilder.select.Select;
 
 /**
@@ -104,7 +105,7 @@ public interface EntityHelper<EntityT> {
    * if the DAO was built without a specific keyspace and table, the query doesn't specify a
    * keyspace, and the table name is inferred from the naming strategy.
    */
-  BuildableQuery insert();
+  RegularInsert insert();
 
   /**
    * Builds a select query to fetch an instance of the entity by primary key (partition key +

diff --git a/mapper-runtime/src/main/java/com/datastax/oss/driver/internal/mapper/DaoBase.java b/mapper-runtime/src/main/java/com/datastax/oss/driver/internal/mapper/DaoBase.java
@@ -46,6 +46,8 @@ public class DaoBase {
   /** The qualified table id placeholder in {@link Query#value()}. */
   public static final String QUALIFIED_TABLE_ID_PLACEHOLDER = "${qualifiedTableId}";
 
+  private static final CqlIdentifier APPLIED = CqlIdentifier.fromInternal("[applied]");
+
   protected static CompletionStage<PreparedStatement> prepare(
       SimpleStatement statement, MapperContext context) {
     if (statement == null) {
@@ -153,8 +155,17 @@ protected Row executeAndExtractFirstRow(Statement<?> statement) {
   protected <EntityT> EntityT executeAndMapToSingleEntity(
       Statement<?> statement, EntityHelper<EntityT> entityHelper) {
     ResultSet rs = execute(statement);
-    Row row = rs.one();
-    return (row == null) ? null : entityHelper.get(row);
+    return asEntity(rs.one(), entityHelper);
+  }
+
+  private <EntityT> EntityT asEntity(Row row, EntityHelper<EntityT> entityHelper) {
+    return (row == null
+            // Special case for INSERT IF NOT EXISTS. If the row did not exists, the query returns
+            // only [applied], we want to return null to indicate there was no previous entity
+            || (row.getColumnDefinitions().size() == 1
+                && row.getColumnDefinitions().get(0).getName().equals(APPLIED)))
+        ? null
+        : entityHelper.get(row);
   }
 
   protected <EntityT> Optional<EntityT> executeAndMapToOptionalEntity(
@@ -199,22 +210,13 @@ protected CompletableFuture<Row> executeAsyncAndExtractFirstRow(Statement<?> sta
 
   protected <EntityT> CompletableFuture<EntityT> executeAsyncAndMapToSingleEntity(
       Statement<?> statement, EntityHelper<EntityT> entityHelper) {
-    return executeAsync(statement)
-        .thenApply(
-            rs -> {
-              Row row = rs.one();
-              return (row == null) ? null : entityHelper.get(row);
-            });
+    return executeAsync(statement).thenApply(rs -> asEntity(rs.one(), entityHelper));
   }
 
   protected <EntityT> CompletableFuture<Optional<EntityT>> executeAsyncAndMapToOptionalEntity(
       Statement<?> statement, EntityHelper<EntityT> entityHelper) {
     return executeAsync(statement)
-        .thenApply(
-            rs -> {
-              Row row = rs.one();
-              return (row == null) ? Optional.empty() : Optional.of(entityHelper.get(row));
-            });
+        .thenApply(rs -> Optional.ofNullable(asEntity(rs.one(), entityHelper)));
   }
 
   protected <EntityT>