private IndexRequest prepareInsert(
DocTableInfo tableInfo, ShardUpsertRequest request, ShardUpsertRequest.Item item)
throws IOException {
List<GeneratedReferenceInfo> generatedReferencesWithValue = new ArrayList<>();
BytesReference source;
if (request.isRawSourceInsert()) {
assert item.insertValues().length > 0 : "empty insert values array";
source = new BytesArray((BytesRef) item.insertValues()[0]);
} else {
XContentBuilder builder = XContentFactory.jsonBuilder().startObject();
for (int i = 0; i < item.insertValues().length; i++) {
Reference ref = request.insertColumns()[i];
if (ref.info().granularity() == RowGranularity.DOC) {
// don't include values for partitions in the _source
// ideally columns with partition granularity shouldn't be part of the request
builder.field(ref.ident().columnIdent().fqn(), item.insertValues()[i]);
if (ref.info() instanceof GeneratedReferenceInfo) {
generatedReferencesWithValue.add((GeneratedReferenceInfo) ref.info());
}
}
}
source = builder.bytes();
}
int generatedColumnSize = 0;
for (GeneratedReferenceInfo generatedReferenceInfo : tableInfo.generatedColumns()) {
if (!tableInfo.partitionedByColumns().contains(generatedReferenceInfo)) {
generatedColumnSize++;
}
}
int numMissingGeneratedColumns = generatedColumnSize - generatedReferencesWithValue.size();
if (numMissingGeneratedColumns > 0
|| (generatedReferencesWithValue.size() > 0 && request.validateGeneratedColumns())) {
// we need to evaluate some generated column expressions
Map<String, Object> sourceMap =
processGeneratedColumnsOnInsert(
tableInfo,
request.insertColumns(),
item.insertValues(),
request.isRawSourceInsert(),
request.validateGeneratedColumns());
source = XContentFactory.jsonBuilder().map(sourceMap).bytes();
}
IndexRequest indexRequest =
Requests.indexRequest(request.index())
.type(request.type())
.id(item.id())
.routing(request.routing())
.source(source)
.create(!request.overwriteDuplicates())
.operationThreaded(false);
if (logger.isTraceEnabled()) {
logger.trace(
"Inserting document with id {}, source: {}", item.id(), indexRequest.source().toUtf8());
}
return indexRequest;
}
@Test
public void testUpdateAssignmentConvertableType() throws Exception {
UpdateAnalyzedStatement.NestedAnalyzedStatement statement =
analyze("update users set other_id=9.9").nestedStatements().get(0);
Reference ref = statement.assignments().keySet().iterator().next();
assertThat(ref, not(instanceOf(DynamicReference.class)));
assertEquals(DataTypes.LONG, ref.info().type());
Symbol value = statement.assignments().entrySet().iterator().next().getValue();
assertThat(value, isLiteral(9L));
}
@Test
public void testUpdateAssignmentNestedDynamicColumn() throws Exception {
UpdateAnalyzedStatement.NestedAnalyzedStatement statement =
analyze("update users set details['arms']=3").nestedStatements().get(0);
assertThat(statement.assignments().size(), is(1));
Reference ref = statement.assignments().keySet().iterator().next();
assertThat(ref, instanceOf(DynamicReference.class));
Assert.assertEquals(DataTypes.LONG, ref.info().type());
assertThat(ref.info().ident().columnIdent().isColumn(), is(false));
assertThat(ref.info().ident().columnIdent().fqn(), is("details.arms"));
}
@Nullable
@Override
public Object referenceValue(Reference reference) {
if (updatedColumnValues == null) {
return super.referenceValue(reference);
}
Object value = updatedColumnValues.get(reference.ident().columnIdent().fqn());
if (value == null && !reference.ident().isColumn()) {
value =
XContentMapValues.extractValue(
reference.ident().columnIdent().fqn(), updatedColumnValues);
}
return reference.valueType().value(value);
}
@Test
public void testUpdateAssignments() throws Exception {
UpdateAnalyzedStatement statement = analyze("update users set name='Trillian'");
UpdateAnalyzedStatement.NestedAnalyzedStatement statement1 =
statement.nestedStatements().get(0);
assertThat(statement1.assignments().size(), is(1));
assertThat(
((DocTableRelation) statement.sourceRelation()).tableInfo().ident(),
is(new TableIdent(Schemas.DEFAULT_SCHEMA_NAME, "users")));
Reference ref = statement1.assignments().keySet().iterator().next();
assertThat(ref.info().ident().tableIdent().name(), is("users"));
assertThat(ref.info().ident().columnIdent().name(), is("name"));
assertTrue(statement1.assignments().containsKey(ref));
Symbol value = statement1.assignments().entrySet().iterator().next().getValue();
assertThat(value, isLiteral("Trillian"));
}