Last active
July 18, 2019 21:58
-
-
Save olim7t/296121eb398fb151ac1b70fff204dc3f to your computer and use it in GitHub Desktop.
Example of batch with `@QueryProvider` in driver 4 mapper (https://docs.datastax.com/en/developer/java-driver/4.1/manual/mapper/daos/queryprovider/)
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
package com.datastax.oss.driver.mapper; | |
import static com.datastax.oss.driver.api.querybuilder.QueryBuilder.bindMarker; | |
import com.datastax.oss.driver.api.core.CqlSession; | |
import com.datastax.oss.driver.api.core.cql.BatchStatement; | |
import com.datastax.oss.driver.api.core.cql.BatchStatementBuilder; | |
import com.datastax.oss.driver.api.core.cql.BoundStatementBuilder; | |
import com.datastax.oss.driver.api.core.cql.DefaultBatchType; | |
import com.datastax.oss.driver.api.core.cql.PreparedStatement; | |
import com.datastax.oss.driver.api.mapper.MapperBuilder; | |
import com.datastax.oss.driver.api.mapper.MapperContext; | |
import com.datastax.oss.driver.api.mapper.annotations.Dao; | |
import com.datastax.oss.driver.api.mapper.annotations.DaoFactory; | |
import com.datastax.oss.driver.api.mapper.annotations.DaoKeyspace; | |
import com.datastax.oss.driver.api.mapper.annotations.Entity; | |
import com.datastax.oss.driver.api.mapper.annotations.Mapper; | |
import com.datastax.oss.driver.api.mapper.annotations.PartitionKey; | |
import com.datastax.oss.driver.api.mapper.annotations.QueryProvider; | |
import com.datastax.oss.driver.api.mapper.entity.EntityHelper; | |
import com.datastax.oss.driver.api.mapper.entity.saving.NullSavingStrategy; | |
import com.datastax.oss.driver.api.querybuilder.insert.Insert; | |
import com.datastax.oss.driver.shaded.guava.common.collect.ImmutableList; | |
import java.util.List; | |
public class BatchExample { | |
public static void main(String[] args) { | |
try (CqlSession session = CqlSession.builder().withLocalDatacenter("datacenter1").build()) { | |
BatchExample.TestMapper mapper = TestMapper.builder(session).build(); | |
BatchExample.TestDao dao = mapper.dao("test"); | |
dao.save( | |
ImmutableList.of(new EntityA(1), new EntityA(2)), | |
ImmutableList.of(new EntityB(1), new EntityB(2)), | |
3600000); | |
} | |
} | |
@Dao | |
public interface TestDao { | |
@QueryProvider( | |
// Reference our custom provider class: | |
providerClass = BatchInsertProvider.class, | |
// Ask the mapper to inject "entity helpers" that will build/bind our queries: | |
entityHelpers = {EntityA.class, EntityB.class}) | |
void save(List<EntityA> as, List<EntityB> bs, int ttl); | |
} | |
public static class BatchInsertProvider { | |
private final CqlSession session; | |
private final EntityHelper<EntityA> aHelper; | |
private final EntityHelper<EntityB> bHelper; | |
private final PreparedStatement preparedInsertA; | |
private final PreparedStatement preparedInsertB; | |
// The mapper will call this once for every new DAO instance | |
public BatchInsertProvider( | |
MapperContext context, | |
// one for each class we requested in the annotation: | |
EntityHelper<EntityA> aHelper, EntityHelper<EntityB> bHelper) { | |
this.session = context.getSession(); | |
this.aHelper = aHelper; | |
this.bHelper = bHelper; | |
// Since statements don't change across invocations, we can generate and prepare them now. | |
// EntityHelper knows how to generate an INSERT query with all of the entity's fields. Since | |
// you had a TTL in your example I'm also adding one here. The helper returns a Query Builder | |
// object, so we just chain a TTL clause with a named bind marker: | |
Insert insertAQuery = aHelper.insert().usingTtl(bindMarker("ttl")); | |
Insert insertBQuery = bHelper.insert().usingTtl(bindMarker("ttl")); | |
// Now prepare and cache the result: | |
preparedInsertA = session.prepare(insertAQuery.asCql()); | |
preparedInsertB = session.prepare(insertBQuery.asCql()); | |
} | |
// The DAO's save() method will call this every time. | |
// By default, the mapper expects the same name, you can also use a different one if the | |
// annotation specifies `providerMethod`. | |
// The parameters list must be the same. | |
public void save(List<EntityA> as, List<EntityB> bs, int ttl) { | |
BatchStatementBuilder batchBuilder = BatchStatement.builder(DefaultBatchType.LOGGED); | |
for (EntityA a : as) { | |
BoundStatementBuilder boundStatementBuilder = preparedInsertA.boundStatementBuilder(); | |
// EntityHelper also knows how to bind all the fields of an entity to a statement: | |
boundStatementBuilder = | |
aHelper.set(a, boundStatementBuilder, NullSavingStrategy.DO_NOT_SET); | |
// We also bind our additional TTL placeholder: | |
boundStatementBuilder = boundStatementBuilder.setInt("ttl", ttl); | |
batchBuilder.addStatement(boundStatementBuilder.build()); | |
} | |
for (EntityB b : bs) { | |
BoundStatementBuilder boundStatementBuilder = preparedInsertB.boundStatementBuilder(); | |
boundStatementBuilder = | |
bHelper.set(b, boundStatementBuilder, NullSavingStrategy.DO_NOT_SET); | |
boundStatementBuilder = boundStatementBuilder.setInt("ttl", ttl); | |
batchBuilder.addStatement(boundStatementBuilder.build()); | |
} | |
session.execute(batchBuilder.build()); | |
} | |
} | |
// The rest is boilerplate (mapper + entity definitions) | |
@Mapper | |
public interface TestMapper { | |
@DaoFactory | |
TestDao dao(@DaoKeyspace String keyspace); | |
static MapperBuilder<TestMapper> builder(CqlSession session) { | |
return new BatchExample_TestMapperBuilder(session); | |
} | |
} | |
@Entity | |
public static class EntityA { | |
@PartitionKey private int id; | |
public EntityA() {} | |
public EntityA(int id) { | |
this.id = id; | |
} | |
public int getId() { | |
return id; | |
} | |
public void setId(int id) { | |
this.id = id; | |
} | |
} | |
@Entity | |
public static class EntityB { | |
@PartitionKey private int id; | |
public EntityB() {} | |
public EntityB(int id) { | |
this.id = id; | |
} | |
public int getId() { | |
return id; | |
} | |
public void setId(int id) { | |
this.id = id; | |
} | |
} | |
} |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment