com.amazonaws.util.json.Jackson Java Examples

private void precomputeAdditionalMetadata(Set<String> columnsToIgnore, Map<String, ValueSet> predicates, List<AttributeValue> accumulator,
        IncrementingValueNameProducer valueNameProducer, SchemaBuilder partitionsSchemaBuilder, DDBRecordMetadata recordMetadata)
{
    // precompute non-key filter
    String filterExpression = DDBPredicateUtils.generateFilterExpression(columnsToIgnore, predicates, accumulator, valueNameProducer, recordMetadata);
    if (filterExpression != null) {
        partitionsSchemaBuilder.addMetadata(NON_KEY_FILTER_METADATA, filterExpression);
    }

    if (!accumulator.isEmpty()) {
        // add in mappings for aliased columns and value placeholders
        Map<String, String> aliasedColumns = new HashMap<>();
        for (String column : predicates.keySet()) {
            aliasedColumns.put(DDBPredicateUtils.aliasColumn(column), column);
        }
        Map<String, AttributeValue> expressionValueMapping = new HashMap<>();
        // IncrementingValueNameProducer is repeatable for simplicity
        IncrementingValueNameProducer valueNameProducer2 = new IncrementingValueNameProducer();
        for (AttributeValue value : accumulator) {
            expressionValueMapping.put(valueNameProducer2.getNext(), value);
        }
        partitionsSchemaBuilder.addMetadata(EXPRESSION_NAMES_METADATA, Jackson.toJsonString(aliasedColumns));
        partitionsSchemaBuilder.addMetadata(EXPRESSION_VALUES_METADATA, Jackson.toJsonString(expressionValueMapping));
    }
}
 

@Test
public void doGetTableLayoutScan()
        throws Exception
{
    Map<String, ValueSet> constraintsMap = new HashMap<>();
    constraintsMap.put("col_3",
            EquatableValueSet.newBuilder(allocator, new ArrowType.Bool(), true, true)
                    .add(true).build());

    GetTableLayoutRequest req = new GetTableLayoutRequest(TEST_IDENTITY,
            TEST_QUERY_ID,
            TEST_CATALOG_NAME,
            new TableName(TEST_CATALOG_NAME, TEST_TABLE),
            new Constraints(constraintsMap),
            SchemaBuilder.newBuilder().build(),
            Collections.EMPTY_SET);

    GetTableLayoutResponse res = handler.doGetTableLayout(allocator, req);

    logger.info("doGetTableLayout schema - {}", res.getPartitions().getSchema());
    logger.info("doGetTableLayout partitions - {}", res.getPartitions());

    assertThat(res.getPartitions().getSchema().getCustomMetadata().get(PARTITION_TYPE_METADATA), equalTo(SCAN_PARTITION_TYPE));
    // no hash key constraints, so look for segment count column
    assertThat(res.getPartitions().getSchema().findField(SEGMENT_COUNT_METADATA) != null, is(true));
    assertThat(res.getPartitions().getRowCount(), equalTo(1));

    assertThat(res.getPartitions().getSchema().getCustomMetadata().get(NON_KEY_FILTER_METADATA), equalTo("(#col_3 = :v0 OR attribute_not_exists(#col_3) OR #col_3 = :v1)"));

    ImmutableMap<String, String> expressionNames = ImmutableMap.of("#col_3", "col_3");
    assertThat(res.getPartitions().getSchema().getCustomMetadata().get(EXPRESSION_NAMES_METADATA), equalTo(Jackson.toJsonString(expressionNames)));

    ImmutableMap<String, AttributeValue> expressionValues = ImmutableMap.of(":v0", ItemUtils.toAttributeValue(true), ":v1", ItemUtils.toAttributeValue(null));
    assertThat(res.getPartitions().getSchema().getCustomMetadata().get(EXPRESSION_VALUES_METADATA), equalTo(Jackson.toJsonString(expressionValues)));
}
 

public TripEvent(String payload)  {
  super(payload);

  JsonNode json = Jackson.fromJsonString(this.payload, JsonNode.class);
  this.tripId = json.get(TRIP_ID).asLong();
  this.timestamp = new DateTime(json.get(DROPOFF_DATETIME).asText()).getMillis();
}
 

public static TripEvent fromStringShiftOrigin(String payload, Duration timeDelta) {
  ObjectNode json = (ObjectNode) Jackson.fromJsonString(payload, JsonNode.class);

  DateTime pickupTime = new DateTime(json.get(PICKUP_DATETIME).asText());
  DateTime dropoffTime = new DateTime(json.get(DROPOFF_DATETIME).asText());

  json.put(PICKUP_DATETIME, pickupTime.plus(timeDelta).toString());
  json.put(DROPOFF_DATETIME, dropoffTime.plus(timeDelta).toString());

  return new TripEvent(json.toString());
}
 

public static final String proxyIntegration(int statusCode, String[] headers, String body) {
	ObjectMapper mapper = Jackson.getObjectMapper();
	ObjectNode json = mapper.createObjectNode();
	
    json.put("statusCode", statusCode);
    
    ArrayNode array = mapper.createArrayNode();
    Arrays.stream(headers).forEach(header -> array.add(header));
    json.put("headers", array);

    json.put("body", body);

    return json.toString();
}
 

private AmazonWebServiceRequest buildReadRequest(Split split, String tableName, Schema schema)
{
    validateExpectedMetadata(split.getProperties());
    // prepare filters
    String rangeKeyFilter = split.getProperty(RANGE_KEY_FILTER_METADATA);
    String nonKeyFilter = split.getProperty(NON_KEY_FILTER_METADATA);
    Map<String, String> expressionAttributeNames = new HashMap<>();
    Map<String, AttributeValue> expressionAttributeValues = new HashMap<>();
    if (rangeKeyFilter != null || nonKeyFilter != null) {
        try {
            expressionAttributeNames.putAll(Jackson.getObjectMapper().readValue(split.getProperty(EXPRESSION_NAMES_METADATA), STRING_MAP_TYPE_REFERENCE));
            expressionAttributeValues.putAll(Jackson.getObjectMapper().readValue(split.getProperty(EXPRESSION_VALUES_METADATA), ATTRIBUTE_VALUE_MAP_TYPE_REFERENCE));
        }
        catch (IOException e) {
            throw new RuntimeException(e);
        }
    }

    // Only read columns that are needed in the query
    String projectionExpression = schema.getFields()
            .stream()
            .map(field -> {
                String aliasedName = DDBPredicateUtils.aliasColumn(field.getName());
                expressionAttributeNames.put(aliasedName, field.getName());
                return aliasedName;
            })
            .collect(Collectors.joining(","));

    boolean isQuery = split.getProperty(SEGMENT_ID_PROPERTY) == null;

    if (isQuery) {
        // prepare key condition expression
        String indexName = split.getProperty(INDEX_METADATA);
        String hashKeyName = split.getProperty(HASH_KEY_NAME_METADATA);
        String hashKeyAlias = DDBPredicateUtils.aliasColumn(hashKeyName);
        String keyConditionExpression = hashKeyAlias + " = " + HASH_KEY_VALUE_ALIAS;
        if (rangeKeyFilter != null) {
            keyConditionExpression += " AND " + rangeKeyFilter;
        }
        expressionAttributeNames.put(hashKeyAlias, hashKeyName);
        expressionAttributeValues.put(HASH_KEY_VALUE_ALIAS, Jackson.fromJsonString(split.getProperty(hashKeyName), AttributeValue.class));

        return new QueryRequest()
                .withTableName(tableName)
                .withIndexName(indexName)
                .withKeyConditionExpression(keyConditionExpression)
                .withFilterExpression(nonKeyFilter)
                .withExpressionAttributeNames(expressionAttributeNames)
                .withExpressionAttributeValues(expressionAttributeValues)
                .withProjectionExpression(projectionExpression);
    }
    else {
        int segmentId = Integer.parseInt(split.getProperty(SEGMENT_ID_PROPERTY));
        int segmentCount = Integer.parseInt(split.getProperty(SEGMENT_COUNT_METADATA));

        return new ScanRequest()
                .withTableName(tableName)
                .withSegment(segmentId)
                .withTotalSegments(segmentCount)
                .withFilterExpression(nonKeyFilter)
                .withExpressionAttributeNames(expressionAttributeNames.isEmpty() ? null : expressionAttributeNames)
                .withExpressionAttributeValues(expressionAttributeValues.isEmpty() ? null : expressionAttributeValues)
                .withProjectionExpression(projectionExpression);
    }
}
 

@Test
public void doGetTableLayoutQueryIndex()
        throws Exception
{
    Map<String, ValueSet> constraintsMap = new HashMap<>();
    SortedRangeSet.Builder dateValueSet = SortedRangeSet.newBuilder(Types.MinorType.DATEDAY.getType(), false);
    SortedRangeSet.Builder timeValueSet = SortedRangeSet.newBuilder(Types.MinorType.DATEMILLI.getType(), false);
    LocalDateTime dateTime = new LocalDateTime().withYear(2019).withMonthOfYear(9).withDayOfMonth(23).withHourOfDay(11).withMinuteOfHour(18).withSecondOfMinute(37);
    MutableDateTime epoch = new MutableDateTime();
    epoch.setDate(0); //Set to Epoch time
    dateValueSet.add(Range.equal(allocator, Types.MinorType.DATEDAY.getType(), Days.daysBetween(epoch, dateTime.toDateTime()).getDays()));
    LocalDateTime dateTime2 = dateTime.plusHours(26);
    dateValueSet.add(Range.equal(allocator, Types.MinorType.DATEDAY.getType(), Days.daysBetween(epoch, dateTime2.toDateTime()).getDays()));
    long startTime = dateTime.toDateTime().getMillis();
    long endTime = dateTime2.toDateTime().getMillis();
    timeValueSet.add(Range.range(allocator, Types.MinorType.DATEMILLI.getType(), startTime, true,
            endTime, true));
    constraintsMap.put("col_4", dateValueSet.build());
    constraintsMap.put("col_5", timeValueSet.build());

    GetTableLayoutResponse res = handler.doGetTableLayout(allocator, new GetTableLayoutRequest(TEST_IDENTITY,
            TEST_QUERY_ID,
            TEST_CATALOG_NAME,
            TEST_TABLE_NAME,
            new Constraints(constraintsMap),
            SchemaBuilder.newBuilder().build(),
            Collections.EMPTY_SET));

    logger.info("doGetTableLayout schema - {}", res.getPartitions().getSchema());
    logger.info("doGetTableLayout partitions - {}", res.getPartitions());

    assertThat(res.getPartitions().getSchema().getCustomMetadata().get(PARTITION_TYPE_METADATA), equalTo(QUERY_PARTITION_TYPE));
    assertThat(res.getPartitions().getSchema().getCustomMetadata().containsKey(INDEX_METADATA), is(true));
    assertThat(res.getPartitions().getSchema().getCustomMetadata().get(INDEX_METADATA), equalTo("test_index"));
    assertThat(res.getPartitions().getSchema().getCustomMetadata().get(HASH_KEY_NAME_METADATA), equalTo("col_4"));
    assertThat(res.getPartitions().getRowCount(), equalTo(2));
    assertThat(res.getPartitions().getSchema().getCustomMetadata().get(RANGE_KEY_NAME_METADATA), equalTo("col_5"));
    assertThat(res.getPartitions().getSchema().getCustomMetadata().get(RANGE_KEY_FILTER_METADATA), equalTo("(#col_5 >= :v0 AND #col_5 <= :v1)"));

    ImmutableMap<String, String> expressionNames = ImmutableMap.of("#col_4", "col_4", "#col_5", "col_5");
    assertThat(res.getPartitions().getSchema().getCustomMetadata().get(EXPRESSION_NAMES_METADATA), equalTo(Jackson.toJsonString(expressionNames)));

    ImmutableMap<String, AttributeValue> expressionValues = ImmutableMap.of(":v0", ItemUtils.toAttributeValue(startTime), ":v1", ItemUtils.toAttributeValue(endTime));
    assertThat(res.getPartitions().getSchema().getCustomMetadata().get(EXPRESSION_VALUES_METADATA), equalTo(Jackson.toJsonString(expressionValues)));
}
 

@Test
public void doGetTableLayoutScanWithTypeOverride()
        throws Exception
{
    List<Column> columns = new ArrayList<>();
    columns.add(new Column().withName("col1").withType("int"));
    columns.add(new Column().withName("col2").withType("timestamptz"));
    columns.add(new Column().withName("col3").withType("string"));

    Map<String, String> param = ImmutableMap.of(
            SOURCE_TABLE_PROPERTY, TEST_TABLE,
            COLUMN_NAME_MAPPING_PROPERTY, "col1=Col1",
            DATETIME_FORMAT_MAPPING_PROPERTY, "col1=datetime1,col3=datetime3 ");
    Table table = new Table()
            .withParameters(param)
            .withPartitionKeys()
            .withStorageDescriptor(new StorageDescriptor().withColumns(columns));
    GetTableResult mockResult = new GetTableResult().withTable(table);
    when(glueClient.getTable(any())).thenReturn(mockResult);

    TableName tableName = new TableName(DEFAULT_SCHEMA, "glueTableForTestTable");
    GetTableRequest getTableRequest = new GetTableRequest(TEST_IDENTITY, TEST_QUERY_ID, TEST_CATALOG_NAME, tableName);
    GetTableResponse getTableResponse = handler.doGetTable(allocator, getTableRequest);
    logger.info("validateSourceTableNamePropagation: GetTableResponse[{}]", getTableResponse);
    Map<String, String> customMetadata = getTableResponse.getSchema().getCustomMetadata();
    assertThat(customMetadata.get(SOURCE_TABLE_PROPERTY), equalTo(TEST_TABLE));
    assertThat(customMetadata.get(DATETIME_FORMAT_MAPPING_PROPERTY_NORMALIZED), equalTo("Col1=datetime1,col3=datetime3"));

    Map<String, ValueSet> constraintsMap = new HashMap<>();
    constraintsMap.put("col3",
            EquatableValueSet.newBuilder(allocator, new ArrowType.Bool(), true, true)
                    .add(true).build());
    constraintsMap.put("col2",
            EquatableValueSet.newBuilder(allocator, new ArrowType.Bool(), true, true)
                    .add(true).build());

    GetTableLayoutRequest getTableLayoutRequest = new GetTableLayoutRequest(TEST_IDENTITY,
            TEST_QUERY_ID,
            TEST_CATALOG_NAME,
            tableName,
            new Constraints(constraintsMap),
            getTableResponse.getSchema(),
            Collections.EMPTY_SET);


    GetTableLayoutResponse res = handler.doGetTableLayout(allocator, getTableLayoutRequest);

    logger.info("doGetTableLayoutScanWithTypeOverride schema - {}", res.getPartitions().getSchema());
    logger.info("doGetTableLayoutScanWithTypeOverride partitions - {}", res.getPartitions());

    assertThat(res.getPartitions().getSchema().getCustomMetadata().get(PARTITION_TYPE_METADATA), equalTo(SCAN_PARTITION_TYPE));
    // no hash key constraints, so look for segment count column
    assertThat(res.getPartitions().getSchema().findField(SEGMENT_COUNT_METADATA) != null, is(true));
    assertThat(res.getPartitions().getRowCount(), equalTo(1));

    assertThat(res.getPartitions().getSchema().getCustomMetadata().get(NON_KEY_FILTER_METADATA), equalTo("(#col3 = :v0 OR attribute_not_exists(#col3) OR #col3 = :v1)"));

    ImmutableMap<String, String> expressionNames = ImmutableMap.of("#col3", "col3", "#col2", "col2");
    assertThat(res.getPartitions().getSchema().getCustomMetadata().get(EXPRESSION_NAMES_METADATA), equalTo(Jackson.toJsonString(expressionNames)));

    ImmutableMap<String, AttributeValue> expressionValues = ImmutableMap.of(":v0", ItemUtils.toAttributeValue(true), ":v1", ItemUtils.toAttributeValue(null));
    assertThat(res.getPartitions().getSchema().getCustomMetadata().get(EXPRESSION_VALUES_METADATA), equalTo(Jackson.toJsonString(expressionValues)));
}
 

/**
 * Executes a lambda function and returns the result of the execution.
 */
@Override
public cfData execute( cfSession _session, cfArgStructData argStruct ) throws cfmRunTimeException {

	AmazonKey amazonKey = getAmazonKey( _session, argStruct );

	// Arguments to extract
	String payload = getNamedStringParam( argStruct, "payload", null );
	String functionName = getNamedStringParam( argStruct, "function", null );
	String qualifier = getNamedStringParam( argStruct, "qualifier", null );

	try {

		// Construct the Lambda Client
		InvokeRequest invokeRequest = new InvokeRequest();
		invokeRequest.setInvocationType( InvocationType.RequestResponse );
		invokeRequest.setLogType( LogType.Tail );
		invokeRequest.setFunctionName( functionName );
		invokeRequest.setPayload( payload );
		if ( qualifier != null ) {
			invokeRequest.setQualifier( qualifier );
		}

		// Lambda client must be created with credentials
		BasicAWSCredentials awsCreds = new BasicAWSCredentials( amazonKey.getKey(), amazonKey.getSecret() );
		AWSLambda awsLambda = AWSLambdaClientBuilder.standard()
				.withRegion( amazonKey.getAmazonRegion().toAWSRegion().getName() )
				.withCredentials( new AWSStaticCredentialsProvider( awsCreds ) ).build();

		// Execute and process the results
		InvokeResult result = awsLambda.invoke( invokeRequest );

		// Convert the returned result
		ByteBuffer resultPayload = result.getPayload();
		String resultJson = new String( resultPayload.array(), "UTF-8" );
		Map<String, Object> resultMap = Jackson.fromJsonString( resultJson, Map.class );

		return tagUtils.convertToCfData( resultMap );

	} catch ( Exception e ) {
		throwException( _session, "AmazonLambdaExecute: " + e.getMessage() );
		return cfBooleanData.FALSE;
	}

}
 

public static TripEvent fromStringOverwriteTime(String payload) {
  ObjectNode json = (ObjectNode) Jackson.fromJsonString(payload, JsonNode.class);

  DateTime pickupTime = new DateTime(json.get(PICKUP_DATETIME).asText());
  DateTime dropoffTime = new DateTime(json.get(DROPOFF_DATETIME).asText());

  Duration timeDelta = new Duration(dropoffTime, DateTime.now());

  json.put(PICKUP_DATETIME, pickupTime.plus(timeDelta).toString());
  json.put(DROPOFF_DATETIME, dropoffTime.plus(timeDelta).toString());

  return new TripEvent(json.toString());
}