Search in sources :

Example 1 with INTERNAL

use of com.datastax.oss.dsbulk.mapping.CQLRenderMode.INTERNAL in project dsbulk by datastax.

the class SchemaSettings method inferBatchInsertQuery.

private String inferBatchInsertQuery(ImmutableMultimap<MappingField, CQLFragment> fieldsToVariables) {
    List<CQLWord> pks = primaryKeyColumns();
    Set<CQLFragment> allSpecificVariables = new LinkedHashSet<>();
    Map<CQLWord, WriteTimeAndTTL> specificWriteTimesAndTTLs = new LinkedHashMap<>();
    boolean hasGlobalWritetime = false;
    boolean hasGlobalTTL = false;
    for (CQLFragment variable : fieldsToVariables.values()) {
        if (variable instanceof FunctionCall) {
            FunctionCall functionCall = (FunctionCall) variable;
            if (functionCall.getFunctionName().equals(WRITETIME)) {
                for (CQLFragment arg : functionCall.getArgs()) {
                    if (arg.equals(STAR)) {
                        if (preserveTimestamp) {
                            throw new IllegalStateException("Invalid mapping: writetime(*) is not allowed when schema.preserveTimestamp is true.");
                        }
                        hasGlobalWritetime = true;
                    } else {
                        CQLWord col = (CQLWord) arg;
                        if (pks.contains(col)) {
                            throw new IllegalStateException("Invalid mapping: writetime() function arg must be either '*' or a non-primary key column name.");
                        }
                        if (fieldsToVariables.containsValue(col)) {
                            allSpecificVariables.add(col);
                            allSpecificVariables.add(functionCall);
                            specificWriteTimesAndTTLs.compute(col, (k, v) -> {
                                if (v == null) {
                                    v = new WriteTimeAndTTL();
                                    MappingField colField = fieldsToVariables.inverse().get(col).iterator().next();
                                    v.value = colField instanceof CQLFragment ? (CQLFragment) colField : col;
                                }
                                MappingField writetimeField = fieldsToVariables.inverse().get(functionCall).iterator().next();
                                v.writetime = writetimeField instanceof CQLLiteral ? (CQLLiteral) writetimeField : CQLWord.fromInternal(functionCall.render(INTERNAL));
                                return v;
                            });
                        } else {
                            throw new IllegalStateException(String.format("Invalid mapping: target column %s must be present if %s is also present.", col.render(VARIABLE), functionCall.render(INTERNAL)));
                        }
                    }
                }
            } else if (functionCall.getFunctionName().equals(TTL)) {
                for (CQLFragment arg : functionCall.getArgs()) {
                    if (arg.equals(STAR)) {
                        if (preserveTtl) {
                            throw new IllegalStateException("Invalid mapping: ttl(*) is not allowed when schema.preserveTtl is true.");
                        }
                        hasGlobalTTL = true;
                    } else {
                        CQLWord col = (CQLWord) arg;
                        if (pks.contains(col)) {
                            throw new IllegalStateException("Invalid mapping: ttl() function arg must be either '*' or a non-primary key column name.");
                        }
                        if (fieldsToVariables.containsValue(col)) {
                            allSpecificVariables.add(col);
                            allSpecificVariables.add(functionCall);
                            specificWriteTimesAndTTLs.compute((CQLWord) arg, (k, v) -> {
                                if (v == null) {
                                    v = new WriteTimeAndTTL();
                                    MappingField colField = fieldsToVariables.inverse().get(col).iterator().next();
                                    v.value = colField instanceof CQLFragment ? (CQLFragment) colField : col;
                                }
                                MappingField ttlField = fieldsToVariables.inverse().get(functionCall).iterator().next();
                                v.ttl = ttlField instanceof CQLLiteral ? (CQLLiteral) ttlField : CQLWord.fromInternal(functionCall.render(INTERNAL));
                                return v;
                            });
                        } else {
                            throw new IllegalStateException(String.format("Invalid mapping: target column %s must be present if %s is also present.", col.render(VARIABLE), functionCall.render(INTERNAL)));
                        }
                    }
                }
            }
        }
    }
    ImmutableMultimap.Builder<MappingField, CQLFragment> defaultFieldsToVariablesBuilder = ImmutableMultimap.builder();
    for (Entry<MappingField, CQLFragment> entry : fieldsToVariables.entries()) {
        CQLFragment variable = entry.getValue();
        if (!allSpecificVariables.contains(variable)) {
            defaultFieldsToVariablesBuilder.put(entry);
        }
    }
    ImmutableMultimap<MappingField, CQLFragment> defaultFieldsToVariables = defaultFieldsToVariablesBuilder.build();
    boolean hasRegularColumnsWithoutSpecificWritetimeAndTTL = defaultFieldsToVariables.values().stream().filter(CQLWord.class::isInstance).map(CQLWord.class::cast).anyMatch(variable -> !pks.contains(variable));
    if (!hasRegularColumnsWithoutSpecificWritetimeAndTTL) {
        if (hasGlobalWritetime) {
            throw new IllegalStateException("Invalid mapping: writetime(*) function has no target column.");
        }
        if (hasGlobalTTL) {
            throw new IllegalStateException("Invalid mapping: ttl(*) function has no target column.");
        }
    }
    StringBuilder sb = new StringBuilder();
    if (!hasRegularColumnsWithoutSpecificWritetimeAndTTL && specificWriteTimesAndTTLs.size() == 1) {
        // edge case: there is only one regular column in the table,
        // and it has specific writetime or tll: no need for a BATCH as there is only one child
        // statement.
        Entry<CQLWord, WriteTimeAndTTL> entry = specificWriteTimesAndTTLs.entrySet().iterator().next();
        appendBatchChildQuery(sb, entry.getKey(), entry.getValue().value, entry.getValue().writetime, entry.getValue().ttl, pks);
    } else {
        sb.append("BEGIN UNLOGGED BATCH ");
        // generate a first INSERT INTO child query similar to the ones generated for simple INSERTs.
        if (hasRegularColumnsWithoutSpecificWritetimeAndTTL) {
            sb.append(inferInsertQuery(defaultFieldsToVariables)).append("; ");
        }
        // generate a specific INSERT INTO query for that variable only + its TTL and/or writetime.
        for (Entry<CQLWord, WriteTimeAndTTL> entry : specificWriteTimesAndTTLs.entrySet()) {
            appendBatchChildQuery(sb, entry.getKey(), entry.getValue().value, entry.getValue().writetime, entry.getValue().ttl, pks);
            sb.append("; ");
        }
        sb.append("APPLY BATCH");
    }
    return sb.toString();
}
Also used : LinkedHashSet(java.util.LinkedHashSet) CQLLiteral(com.datastax.oss.dsbulk.mapping.CQLLiteral) TypedCQLLiteral(com.datastax.oss.dsbulk.mapping.TypedCQLLiteral) MAPPED_OR_INDEXED(com.datastax.oss.dsbulk.mapping.MappingPreference.MAPPED_OR_INDEXED) CqlIdentifier(com.datastax.oss.driver.api.core.CqlIdentifier) MAPPED_ONLY(com.datastax.oss.dsbulk.mapping.MappingPreference.MAPPED_ONLY) DefaultMapping(com.datastax.oss.dsbulk.mapping.DefaultMapping) GenericType(com.datastax.oss.driver.api.core.type.reflect.GenericType) TokenRangeReadStatementGenerator(com.datastax.oss.dsbulk.partitioner.TokenRangeReadStatementGenerator) ConfigUtils(com.datastax.oss.dsbulk.config.ConfigUtils) ALIASED_SELECTOR(com.datastax.oss.dsbulk.mapping.CQLRenderMode.ALIASED_SELECTOR) BatchType(com.datastax.oss.driver.api.core.cql.BatchType) Map(java.util.Map) VisibleForTesting(com.datastax.oss.driver.shaded.guava.common.annotations.VisibleForTesting) ProtocolVersion(com.datastax.oss.driver.api.core.ProtocolVersion) EnumSet(java.util.EnumSet) CQLRenderMode(com.datastax.oss.dsbulk.mapping.CQLRenderMode) VARIABLE(com.datastax.oss.dsbulk.mapping.CQLRenderMode.VARIABLE) MapType(com.datastax.oss.driver.api.core.type.MapType) GraphUtils(com.datastax.oss.dsbulk.workflow.commons.utils.GraphUtils) Set(java.util.Set) ImmutableList(com.datastax.oss.driver.shaded.guava.common.collect.ImmutableList) Field(com.datastax.oss.dsbulk.connectors.api.Field) Stream(java.util.stream.Stream) ConfigException(com.typesafe.config.ConfigException) CQLWord(com.datastax.oss.dsbulk.mapping.CQLWord) INDEXED_ONLY(com.datastax.oss.dsbulk.mapping.MappingPreference.INDEXED_ONLY) MICROSECONDS(java.util.concurrent.TimeUnit.MICROSECONDS) CQLLiteral(com.datastax.oss.dsbulk.mapping.CQLLiteral) ImmutableSet(com.datastax.oss.driver.shaded.guava.common.collect.ImmutableSet) Preconditions(com.datastax.oss.driver.shaded.guava.common.base.Preconditions) STAR(com.datastax.oss.dsbulk.mapping.MappingInspector.STAR) ViewMetadata(com.datastax.oss.driver.api.core.metadata.schema.ViewMetadata) Supplier(java.util.function.Supplier) ArrayList(java.util.ArrayList) LinkedHashMap(java.util.LinkedHashMap) CqlSession(com.datastax.oss.driver.api.core.CqlSession) DefaultReadResultMapper(com.datastax.oss.dsbulk.workflow.commons.schema.DefaultReadResultMapper) FunctionCall(com.datastax.oss.dsbulk.mapping.FunctionCall) StreamSupport(java.util.stream.StreamSupport) Metadata(com.datastax.oss.driver.api.core.metadata.Metadata) LinkedHashSet(java.util.LinkedHashSet) DseEdgeMetadata(com.datastax.dse.driver.api.core.metadata.schema.DseEdgeMetadata) TableMetadata(com.datastax.oss.driver.api.core.metadata.schema.TableMetadata) RelationMetadata(com.datastax.oss.driver.api.core.metadata.schema.RelationMetadata) MappingPreference(com.datastax.oss.dsbulk.mapping.MappingPreference) Config(com.typesafe.config.Config) NestedBatchException(com.datastax.oss.dsbulk.workflow.commons.schema.NestedBatchException) DataType(com.datastax.oss.driver.api.core.type.DataType) ConvertingCodecFactory(com.datastax.oss.dsbulk.codecs.api.ConvertingCodecFactory) DefaultRecordMapper(com.datastax.oss.dsbulk.workflow.commons.schema.DefaultRecordMapper) RecordMapper(com.datastax.oss.dsbulk.workflow.commons.schema.RecordMapper) DseTableMetadata(com.datastax.dse.driver.api.core.metadata.schema.DseTableMetadata) CQLFragment(com.datastax.oss.dsbulk.mapping.CQLFragment) WRITETIME(com.datastax.oss.dsbulk.mapping.MappingInspector.WRITETIME) ReadResultMapper(com.datastax.oss.dsbulk.workflow.commons.schema.ReadResultMapper) DseVertexMetadata(com.datastax.dse.driver.api.core.metadata.schema.DseVertexMetadata) Nullable(edu.umd.cs.findbugs.annotations.Nullable) StatisticsMode(com.datastax.oss.dsbulk.workflow.commons.settings.StatsSettings.StatisticsMode) DefaultReadResultCounter(com.datastax.oss.dsbulk.workflow.commons.schema.DefaultReadResultCounter) CodecUtils.instantToNumber(com.datastax.oss.dsbulk.codecs.api.util.CodecUtils.instantToNumber) IndexedMappingField(com.datastax.oss.dsbulk.mapping.IndexedMappingField) LoggerFactory(org.slf4j.LoggerFactory) MappingField(com.datastax.oss.dsbulk.mapping.MappingField) QueryInspector(com.datastax.oss.dsbulk.workflow.commons.schema.QueryInspector) Mapping(com.datastax.oss.dsbulk.mapping.Mapping) ReadResultCounter(com.datastax.oss.dsbulk.workflow.commons.schema.ReadResultCounter) Lists(com.datastax.oss.driver.shaded.guava.common.collect.Lists) ConvertingCodec(com.datastax.oss.dsbulk.codecs.api.ConvertingCodec) NonNull(edu.umd.cs.findbugs.annotations.NonNull) Predicates(com.datastax.oss.driver.shaded.guava.common.base.Predicates) URI(java.net.URI) IndexMetadata(com.datastax.oss.driver.api.core.metadata.schema.IndexMetadata) Record(com.datastax.oss.dsbulk.connectors.api.Record) ImmutableMultimap(com.datastax.oss.driver.shaded.guava.common.collect.ImmutableMultimap) Predicate(java.util.function.Predicate) Collection(java.util.Collection) Instant(java.time.Instant) Collectors(java.util.stream.Collectors) NAMED_ASSIGNMENT(com.datastax.oss.dsbulk.mapping.CQLRenderMode.NAMED_ASSIGNMENT) Objects(java.util.Objects) KeyspaceMetadata(com.datastax.oss.driver.api.core.metadata.schema.KeyspaceMetadata) DataTypes(com.datastax.oss.driver.api.core.type.DataTypes) List(java.util.List) Entry(java.util.Map.Entry) Optional(java.util.Optional) SetType(com.datastax.oss.driver.api.core.type.SetType) TIMESTAMP_PATTERN(com.datastax.oss.dsbulk.codecs.api.CommonConversionContext.TIMESTAMP_PATTERN) MappingInspector(com.datastax.oss.dsbulk.mapping.MappingInspector) DseGraphKeyspaceMetadata(com.datastax.dse.driver.api.core.metadata.schema.DseGraphKeyspaceMetadata) ListType(com.datastax.oss.driver.api.core.type.ListType) HashSet(java.util.HashSet) RecordMetadata(com.datastax.oss.dsbulk.connectors.api.RecordMetadata) ColumnDefinitions(com.datastax.oss.driver.api.core.cql.ColumnDefinitions) MappedMappingField(com.datastax.oss.dsbulk.mapping.MappedMappingField) ImmutableSetMultimap(com.datastax.oss.driver.shaded.guava.common.collect.ImmutableSetMultimap) INTERNAL(com.datastax.oss.dsbulk.mapping.CQLRenderMode.INTERNAL) Logger(org.slf4j.Logger) Iterator(java.util.Iterator) TypedCQLLiteral(com.datastax.oss.dsbulk.mapping.TypedCQLLiteral) PreparedStatement(com.datastax.oss.driver.api.core.cql.PreparedStatement) DseGraphTableMetadata(com.datastax.dse.driver.api.core.metadata.schema.DseGraphTableMetadata) ColumnMetadata(com.datastax.oss.driver.api.core.metadata.schema.ColumnMetadata) DefaultProtocolVersion(com.datastax.oss.driver.api.core.DefaultProtocolVersion) UserDefinedType(com.datastax.oss.driver.api.core.type.UserDefinedType) TTL(com.datastax.oss.dsbulk.mapping.MappingInspector.TTL) Multimap(com.datastax.oss.driver.shaded.guava.common.collect.Multimap) EPOCH(java.time.Instant.EPOCH) Collections(java.util.Collections) Statement(com.datastax.oss.driver.api.core.cql.Statement) CQLFragment(com.datastax.oss.dsbulk.mapping.CQLFragment) IndexedMappingField(com.datastax.oss.dsbulk.mapping.IndexedMappingField) MappingField(com.datastax.oss.dsbulk.mapping.MappingField) MappedMappingField(com.datastax.oss.dsbulk.mapping.MappedMappingField) LinkedHashMap(java.util.LinkedHashMap) CQLWord(com.datastax.oss.dsbulk.mapping.CQLWord) FunctionCall(com.datastax.oss.dsbulk.mapping.FunctionCall) ImmutableMultimap(com.datastax.oss.driver.shaded.guava.common.collect.ImmutableMultimap)

Aggregations

DseEdgeMetadata (com.datastax.dse.driver.api.core.metadata.schema.DseEdgeMetadata)1 DseGraphKeyspaceMetadata (com.datastax.dse.driver.api.core.metadata.schema.DseGraphKeyspaceMetadata)1 DseGraphTableMetadata (com.datastax.dse.driver.api.core.metadata.schema.DseGraphTableMetadata)1 DseTableMetadata (com.datastax.dse.driver.api.core.metadata.schema.DseTableMetadata)1 DseVertexMetadata (com.datastax.dse.driver.api.core.metadata.schema.DseVertexMetadata)1 CqlIdentifier (com.datastax.oss.driver.api.core.CqlIdentifier)1 CqlSession (com.datastax.oss.driver.api.core.CqlSession)1 DefaultProtocolVersion (com.datastax.oss.driver.api.core.DefaultProtocolVersion)1 ProtocolVersion (com.datastax.oss.driver.api.core.ProtocolVersion)1 BatchType (com.datastax.oss.driver.api.core.cql.BatchType)1 ColumnDefinitions (com.datastax.oss.driver.api.core.cql.ColumnDefinitions)1 PreparedStatement (com.datastax.oss.driver.api.core.cql.PreparedStatement)1 Statement (com.datastax.oss.driver.api.core.cql.Statement)1 Metadata (com.datastax.oss.driver.api.core.metadata.Metadata)1 ColumnMetadata (com.datastax.oss.driver.api.core.metadata.schema.ColumnMetadata)1 IndexMetadata (com.datastax.oss.driver.api.core.metadata.schema.IndexMetadata)1 KeyspaceMetadata (com.datastax.oss.driver.api.core.metadata.schema.KeyspaceMetadata)1 RelationMetadata (com.datastax.oss.driver.api.core.metadata.schema.RelationMetadata)1 TableMetadata (com.datastax.oss.driver.api.core.metadata.schema.TableMetadata)1 ViewMetadata (com.datastax.oss.driver.api.core.metadata.schema.ViewMetadata)1