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

Example 1 with DefaultMapping

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

the class SchemaSettingsTest method should_use_default_writetime_var_name.

@Test
void should_use_default_writetime_var_name() {
    Config config = TestConfigUtils.createTestConfig("dsbulk.schema", "keyspace", "ks", "table", "t1", "mapping", "\" *=*, f1 = writetime(*) \"");
    SchemaSettings settings = new SchemaSettings(config, MAP_AND_WRITE);
    settings.init(session, codecFactory, false, true);
    RecordMapper mapper = settings.createRecordMapper(session, recordMetadata, false);
    DefaultMapping mapping = (DefaultMapping) getInternalState(mapper, "mapping");
    assertThat(mapping).isNotNull();
    @SuppressWarnings("unchecked") Set<CQLWord> writeTimeVariables = (Set<CQLWord>) getInternalState(mapping, "writeTimeVariables");
    assertThat(writeTimeVariables).containsOnly(CQLWord.fromInternal("writetime(*)"));
}

Example 2 with DefaultMapping

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

the class SchemaSettingsTest method should_detect_writetime_var_in_query.

@Test
void should_detect_writetime_var_in_query() {
    ColumnDefinitions definitions = mockColumnDefinitions(mockColumnDefinition("c1", DataTypes.TEXT), mockColumnDefinition("c2", DataTypes.TEXT), mockColumnDefinition("c3", DataTypes.TEXT));
    when(ps.getVariableDefinitions()).thenReturn(definitions);
    when(table.getColumn(CqlIdentifier.fromInternal("c2"))).thenReturn(Optional.of(col2));
    Config config = TestConfigUtils.createTestConfig("dsbulk.schema", "query", "\"INSERT INTO ks.t1 (c1,c2) VALUES (:c1, :c2) USING TIMESTAMP :c3\"", "mapping", "\" f1 = c1 , f2 = c2 , f3 = c3 \" ");
    SchemaSettings settings = new SchemaSettings(config, MAP_AND_WRITE);
    settings.init(session, codecFactory, false, true);
    RecordMapper mapper = settings.createRecordMapper(session, recordMetadata, false);
    DefaultMapping mapping = (DefaultMapping) getInternalState(mapper, "mapping");
    assertThat(mapping).isNotNull();
    @SuppressWarnings("unchecked") Set<CQLWord> writeTimeVariables = (Set<CQLWord>) getInternalState(mapping, "writeTimeVariables");
    assertThat(writeTimeVariables).containsOnly(CQLWord.fromInternal(C3.asInternal()));
}

Example 3 with DefaultMapping

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

the class SchemaSettingsTest method should_detect_positional_writetime_var_in_query.

@Test
void should_detect_positional_writetime_var_in_query() {
    ColumnDefinitions definitions = mockColumnDefinitions(mockColumnDefinition("c1", DataTypes.TEXT), mockColumnDefinition("c2", DataTypes.TEXT), mockColumnDefinition("\"This is a quoted \\\" variable name\"", DataTypes.TEXT));
    when(ps.getVariableDefinitions()).thenReturn(definitions);
    when(table.getColumn(CqlIdentifier.fromInternal("c2"))).thenReturn(Optional.of(col2));
    Config config = TestConfigUtils.createTestConfig("dsbulk.schema", "query", "\"INSERT INTO ks.t1 (c1,c2) VALUES (?, ?) USING TTL 123 AND tImEsTaMp ?\"");
    SchemaSettings settings = new SchemaSettings(config, MAP_AND_WRITE);
    settings.init(session, codecFactory, false, true);
    RecordMapper mapper = settings.createRecordMapper(session, recordMetadata, false);
    DefaultMapping mapping = (DefaultMapping) getInternalState(mapper, "mapping");
    assertThat(mapping).isNotNull();
    @SuppressWarnings("unchecked") Set<CQLWord> writeTimeVariables = (Set<CQLWord>) getInternalState(mapping, "writeTimeVariables");
    assertThat(writeTimeVariables).containsOnly(INTERNAL_TIMESTAMP_VARNAME);
}

Example 4 with DefaultMapping

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

the class SchemaSettingsTest method should_detect_quoted_writetime_var_in_query.

@Test
void should_detect_quoted_writetime_var_in_query() {
    ColumnDefinitions definitions = mockColumnDefinitions(mockColumnDefinition("c1", DataTypes.TEXT), mockColumnDefinition("c2", DataTypes.TEXT), mockColumnDefinition("\"This is a quoted \\\" variable name\"", DataTypes.TEXT));
    when(ps.getVariableDefinitions()).thenReturn(definitions);
    when(table.getColumn(CqlIdentifier.fromInternal("c2"))).thenReturn(Optional.of(col2));
    Config config = TestConfigUtils.createTestConfig("dsbulk.schema", "query", "\"INSERT INTO ks.t1 (c1,c2) VALUES (:c1, :c2) USING TTL 123 AND tImEsTaMp     :\\\"This is a quoted \\\"\\\" variable name\\\"\"");
    SchemaSettings settings = new SchemaSettings(config, MAP_AND_WRITE);
    settings.init(session, codecFactory, false, true);
    RecordMapper mapper = settings.createRecordMapper(session, recordMetadata, false);
    DefaultMapping mapping = (DefaultMapping) getInternalState(mapper, "mapping");
    assertThat(mapping).isNotNull();
    @SuppressWarnings("unchecked") Set<CQLWord> writeTimeVariables = (Set<CQLWord>) getInternalState(mapping, "writeTimeVariables");
    assertThat(writeTimeVariables).containsOnly(CQLWord.fromInternal("This is a quoted \" variable name"));
}

Example 5 with DefaultMapping

use of com.datastax.oss.dsbulk.mapping.DefaultMapping 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()));
}