Examples with FunctionCall com.datastax.oss.dsbulk.mapping.FunctionCall used on opensource projects

Example 6 with FunctionCall

use of com.datastax.oss.dsbulk.mapping.FunctionCall in project dsbulk by datastax.

the class SchemaSettings method inferInsertQuery.

private String inferInsertQuery(ImmutableMultimap<MappingField, CQLFragment> fieldsToVariables) {
    ImmutableMultimap.Builder<MappingField, CQLFragment> regularFieldsToVariablesBuilder = ImmutableMultimap.builder();
    CQLFragment writetime = null;
    CQLFragment ttl = null;
    for (Entry<MappingField, CQLFragment> entry : fieldsToVariables.entries()) {
        if (entry.getValue() instanceof FunctionCall) {
            FunctionCall functionCall = (FunctionCall) entry.getValue();
            if (functionCall.getFunctionName().equals(WRITETIME)) {
                assert writetime == null;
                if (entry.getKey() instanceof CQLLiteral) {
                    writetime = (CQLLiteral) entry.getKey();
                } else {
                    writetime = CQLWord.fromInternal(functionCall.render(INTERNAL));
                }
            } else if (functionCall.getFunctionName().equals(TTL)) {
                assert ttl == null;
                if (entry.getKey() instanceof CQLLiteral) {
                    ttl = (CQLLiteral) entry.getKey();
                } else {
                    ttl = CQLWord.fromInternal(functionCall.render(INTERNAL));
                }
            }
        } else {
            regularFieldsToVariablesBuilder.put(entry);
        }
    }
    ImmutableMultimap<MappingField, CQLFragment> regularFieldsToVariables = regularFieldsToVariablesBuilder.build();
    StringBuilder sb = new StringBuilder("INSERT INTO ");
    sb.append(keyspaceName.render(VARIABLE)).append('.').append(tableName.render(VARIABLE)).append(" (");
    appendColumnNames(regularFieldsToVariables, sb, VARIABLE);
    sb.append(") VALUES (");
    Set<CQLFragment> cols = maybeSortCols(regularFieldsToVariables);
    Iterator<CQLFragment> it = cols.iterator();
    while (it.hasNext()) {
        CQLFragment col = it.next();
        // for insert queries there can be only one field mapped to a given column
        MappingField field = fieldsToVariables.inverse().get(col).iterator().next();
        if (field instanceof CQLFragment) {
            sb.append(((CQLFragment) field).render(NAMED_ASSIGNMENT));
        } else {
            sb.append(col.render(NAMED_ASSIGNMENT));
        }
        if (it.hasNext()) {
            sb.append(", ");
        }
    }
    sb.append(')');
    appendWriteTimeAndTTL(sb, writetime, ttl);
    return sb.toString();
}

Example 7 with FunctionCall

use of com.datastax.oss.dsbulk.mapping.FunctionCall in project dsbulk by datastax.

the class SchemaSettings method prepareStatementAndCreateMapping.

@NonNull
private Mapping prepareStatementAndCreateMapping(CqlSession session, boolean batchingEnabled, EnumSet<StatisticsMode> modes) {
    ImmutableMultimap<MappingField, CQLFragment> fieldsToVariables = null;
    if (!config.hasPath(QUERY)) {
        // in the absence of user-provided queries, create the mapping *before* query generation and
        // preparation
        List<CQLFragment> columns = table.getColumns().values().stream().filter(col -> !isDSESearchPseudoColumn(col)).flatMap(column -> {
            CQLWord colName = CQLWord.fromCqlIdentifier(column.getName());
            List<CQLFragment> cols = Lists.newArrayList(colName);
            if (schemaGenerationStrategy.isMapping()) {
                if (preserveTimestamp && checkWritetimeTtlSupported(column, WRITETIME)) {
                    cols.add(new FunctionCall(null, WRITETIME, colName));
                }
                if (preserveTtl && checkWritetimeTtlSupported(column, TTL)) {
                    cols.add(new FunctionCall(null, TTL, colName));
                }
            }
            return cols.stream();
        }).collect(Collectors.toList());
        fieldsToVariables = createFieldsToVariablesMap(columns);
        // query generation
        if (schemaGenerationStrategy.isWriting()) {
            if (isCounterTable()) {
                query = inferUpdateCounterQuery(fieldsToVariables);
            } else if (requiresBatchInsertQuery(fieldsToVariables)) {
                query = inferBatchInsertQuery(fieldsToVariables);
            } else {
                query = inferInsertQuery(fieldsToVariables);
            }
        } else if (schemaGenerationStrategy.isReading() && schemaGenerationStrategy.isMapping()) {
            query = inferReadQuery(fieldsToVariables);
        } else if (schemaGenerationStrategy.isReading() && schemaGenerationStrategy.isCounting()) {
            query = inferCountQuery(modes);
        } else {
            throw new IllegalStateException("Unsupported schema generation strategy: " + schemaGenerationStrategy);
        }
        LOGGER.debug("Inferred query: {}", query);
        queryInspector = new QueryInspector(query);
        // validate generated query
        if (schemaGenerationStrategy.isWriting()) {
            validatePrimaryKeyPresent(fieldsToVariables);
        }
    }
    assert query != null;
    assert queryInspector != null;
    if (!queryInspector.getKeyspaceName().isPresent()) {
        session.execute("USE " + keyspaceName);
    }
    // Transform user-provided queries before preparation
    if (config.hasPath(QUERY)) {
        if (schemaGenerationStrategy.isReading() && queryInspector.isParallelizable()) {
            int whereClauseIndex = queryInspector.getFromClauseEndIndex() + 1;
            StringBuilder sb = new StringBuilder(query.substring(0, whereClauseIndex));
            appendTokenRangeRestriction(sb);
            query = sb.append(query.substring(whereClauseIndex)).toString();
        }
        if (schemaGenerationStrategy.isCounting()) {
            if (modes.contains(StatisticsMode.partitions) || modes.contains(StatisticsMode.ranges) || modes.contains(StatisticsMode.hosts)) {
                throw new IllegalArgumentException(String.format("Cannot count with stats.modes = %s when schema.query is provided; " + "only stats.modes = [global] is allowed", modes));
            }
            // reduce row size by only selecting one column
            StringBuilder sb = new StringBuilder("SELECT ");
            sb.append(getGlobalCountSelector());
            query = sb.append(' ').append(query.substring(queryInspector.getFromClauseStartIndex())).toString();
        }
        queryInspector = new QueryInspector(query);
    }
    if (batchingEnabled && queryInspector.isBatch()) {
        preparedStatements = unwrapAndPrepareBatchChildStatements(session);
    } else {
        preparedStatements = Collections.singletonList(session.prepare(query));
    }
    if (config.hasPath(QUERY)) {
        // in the presence of user-provided queries, create the mapping *after* query preparation
        Stream<ColumnDefinitions> variables = getVariables();
        fieldsToVariables = createFieldsToVariablesMap(variables.flatMap(defs -> StreamSupport.stream(defs.spliterator(), false)).map(def -> def.getName().asInternal()).map(CQLWord::fromInternal).collect(Collectors.toList()));
        // validate user-provided query
        if (schemaGenerationStrategy.isWriting()) {
            if (mutatesOnlyStaticColumns()) {
                // DAT-414: mutations that only affect static columns are allowed
                // to skip the clustering columns, only the partition key should be present.
                validatePartitionKeyPresent(fieldsToVariables);
            } else {
                validatePrimaryKeyPresent(fieldsToVariables);
            }
        }
    }
    assert fieldsToVariables != null;
    return new DefaultMapping(transformFieldsToVariables(fieldsToVariables), codecFactory, transformWriteTimeVariables(queryInspector.getWriteTimeVariables()));
}

Example 8 with FunctionCall

use of com.datastax.oss.dsbulk.mapping.FunctionCall in project dsbulk by datastax.

the class SchemaSettings method inferFieldsToVariablesMap.

private ImmutableMultimap<MappingField, CQLFragment> inferFieldsToVariablesMap(Collection<CQLFragment> columns) {
    // use a builder to preserve iteration order
    ImmutableMultimap.Builder<MappingField, CQLFragment> fieldsToVariables = ImmutableMultimap.builder();
    List<CQLFragment> excludedVariables = new ArrayList<>(mapping.getExcludedVariables());
    if (!isCounterTable() && schemaGenerationStrategy.isMapping()) {
        for (CQLWord variable : mapping.getExcludedVariables()) {
            if (preserveTimestamp) {
                excludedVariables.add(new FunctionCall(null, WRITETIME, variable));
            }
            if (preserveTtl) {
                excludedVariables.add(new FunctionCall(null, TTL, variable));
            }
        }
    }
    int i = 0;
    for (CQLFragment colName : columns) {
        if (!excludedVariables.contains(colName)) {
            if (mappingPreference == INDEXED_ONLY) {
                fieldsToVariables.put(new IndexedMappingField(i), colName);
            } else {
                fieldsToVariables.put(new MappedMappingField(colName.render(INTERNAL)), colName);
            }
        }
        i++;
    }
    return fieldsToVariables.build();
}

Example 9 with FunctionCall

use of com.datastax.oss.dsbulk.mapping.FunctionCall in project dsbulk by datastax.

the class QueryInspectorTest method should_detect_writetime_select.

@Test
void should_detect_writetime_select() {
    QueryInspector inspector = new QueryInspector("SELECT WrItEtImE(col1) FROM ks.table1");
    assertThat(inspector.getWriteTimeVariables()).hasSize(1).containsOnly(new FunctionCall(null, WRITETIME, COL_1));
}

Example 10 with FunctionCall

use of com.datastax.oss.dsbulk.mapping.FunctionCall in project dsbulk by datastax.

the class QueryInspectorTest method should_detect_writetime_quoted_select.

@Test
void should_detect_writetime_quoted_select() {
    QueryInspector inspector = new QueryInspector("SELECT WrItEtImE(\"My Col 2\") FROM ks.table1");
    assertThat(inspector.getWriteTimeVariables()).hasSize(1).containsOnly(new FunctionCall(null, WRITETIME, MY_COL_2));
}