Java Code Examples for com.streamsets.pipeline.api.Record#Header

private void setSqsAttributesOnRecord(Message message, Record record, String queueUrl, String queueNamePrefix) {
  final Record.Header header = record.getHeader();

  switch (sqsAttributesOption) {
    case ALL:
      header.setAttribute(SQS_QUEUE_URL_ATTRIBUTE, queueUrl);
      Optional.of(message.getMessageAttributes()).ifPresent(attrs -> {
        attrs.forEach((name, val) -> {
          final String stringValue = val.getStringValue();
          if (stringValue != null) {
            header.setAttribute(SQS_MESSAGE_ATTRIBUTE_PREFIX + name, stringValue);
          }
        });
      });
      final String body = message.getBody();
      if (body != null) {
        header.setAttribute(SQS_MESSAGE_BODY_ATTRIBUTE, body);
      }
      final String bodyMd5 = message.getMD5OfBody();
      if (bodyMd5 != null) {
        header.setAttribute(SQS_MESSAGE_BODY_MD5_ATTRIBUTE, bodyMd5);
      }
      final String attrsMd5 = message.getMD5OfMessageAttributes();
      if (attrsMd5 != null) {
        header.setAttribute(SQS_MESSAGE_ATTRIBUTE_MD5_ATTRIBUTE, attrsMd5);
      }
      // fall through
    case BASIC:
      header.setAttribute(SQS_MESSAGE_ID_ATTRIBUTE, message.getMessageId());
      header.setAttribute(SQS_QUEUE_NAME_PREFIX_ATTRIBUTE, queueNamePrefix);
      header.setAttribute(SQS_REGION_ATTRIBUTE, awsRegionLabel);
      break;
    case NONE:
      // empty block
      break;
  }
}
 

/**
 * Kryo loads the RecordImpl in Spark's classloader. So this one clones it to this stage's classloader.
 *
 * @param record Record to be cloned
 * @param context The context of the {@linkplain Processor} to use to clone the record
 * @return Cloned record
 */
@SuppressWarnings("unchecked")
public static Record clone(Object record, Processor.Context context) {
  Record newRecord = context.createRecord("dummyId");
  try {
    Object origHeaders = record.getClass().getMethod("getHeader").invoke(record);
    Map<String, Object> headers =
        (Map<String, Object>) origHeaders.getClass().getMethod("getAllAttributes").invoke(origHeaders);
    Record.Header newHeaders = newRecord.getHeader();
    newHeaders.getClass().getMethod("overrideUserAndSystemAttributes", Map.class).invoke(newHeaders, headers);
    newRecord.set(RecordCloner.cloneField(record.getClass().getMethod("get").invoke(record)));
    return newRecord;
  } catch(Exception ex) {
    throw new RuntimeException(ex);
  }
}
 

private void testEmptyPayloadRequest(String method, String httpServerUrl, List<Record> requestRecords) {
  Response response = ClientBuilder.newClient()
      .property(ClientProperties.SUPPRESS_HTTP_COMPLIANCE_VALIDATION, true)
      .property(HttpUrlConnectorProvider.SET_METHOD_WORKAROUND, true)
      .target(httpServerUrl)
      .request()
      .header(Constants.X_SDC_APPLICATION_ID_HEADER, "id")
      .method(method);

  Assert.assertEquals(HttpURLConnection.HTTP_OK, response.getStatus());
  String responseBody = response.readEntity(String.class);

  Assert.assertEquals(1, requestRecords.size());
  Record.Header emptyPayloadRecordHeader = requestRecords.get(0).getHeader();
  Assert.assertEquals(
      "true",
      emptyPayloadRecordHeader.getAttribute(RestServiceReceiver.EMPTY_PAYLOAD_RECORD_HEADER_ATTR_NAME)
  );
  Assert.assertEquals(method, emptyPayloadRecordHeader.getAttribute(RestServiceReceiver.METHOD_HEADER));


  // check custom HTTP Response header
  Assert.assertNotNull(response.getHeaders().getFirst("test"));
  Assert.assertEquals("value", response.getHeaders().getFirst("test"));
}
 

@NotNull
private Record makeListRecord(Map<String, Pair<Field.Type, Object>> fieldMap) {
  Record record = RecordCreator.create();
  Record.Header header = record.getHeader();
  ArrayList<Field> fields = new ArrayList<>();
  for (Map.Entry<String, Pair<Field.Type, Object>> entry : fieldMap.entrySet()) {
    String fieldName = entry.getKey();
    Field.Type fieldType = entry.getValue().getLeft();
    Field field = Field.create(fieldType, entry.getValue().getRight());
    if (fieldType == Field.Type.DECIMAL) {
      field.setAttribute(HeaderAttributeConstants.ATTR_SCALE, SCALE);
      field.setAttribute(HeaderAttributeConstants.ATTR_PRECISION, PRECISION);
    }
    fields.add(field);
  }
  record.set(Field.create(fields));
  header.setAttribute("table", tableName);
  return record;
}
 

@Test
public void testHttpHead() throws Exception {
  HttpProcessorConfig conf = new HttpProcessorConfig();
  conf.httpMethod = HttpMethod.HEAD;
  conf.dataFormat = DataFormat.TEXT;
  conf.resourceUrl = getBaseUri() + "test/head";
  conf.headerOutputLocation = HeaderOutputLocation.HEADER;

  List<Record> records = createRecords("test/head");
  ProcessorRunner runner = createProcessorRunner(conf);
  try {
    StageRunner.Output output = runner.runProcess(records);
    getOutputField(output);
    Record.Header header = getOutputRecord(output).getHeader();
    assertEquals("StreamSets", header.getAttribute("x-test-header"));
    assertEquals("[a, b]", header.getAttribute("x-list-header"));
  } finally {
    runner.runDestroy();
  }
}
 

@Test
public void testShouldRollWithRollHeader() throws Exception {
  final String rollHeaderName = "roll";
  final boolean rollIfHeader = true;
  final String dirPath = "";
  Record record = RecordCreator.create();
  Record.Header header = record.getHeader();
  header.setAttribute(rollHeaderName, rollHeaderName);

  DataLakeGeneratorManager dataLakeGeneratorManager = new DataLakeGeneratorManagerTestBuilder()
      .rollHeaderName(rollHeaderName)
      .rollIfHeader(rollIfHeader)
      .build();

  Assert.assertTrue(dataLakeGeneratorManager.shouldRoll(record, dirPath));
}
 

private void publish(Record record) throws StageException {
  ByteArrayOutputStream os = new ByteArrayOutputStream();
  try (DataGenerator generator = generatorFactory.getGenerator(os)) {
    generator.write(record);
  } catch (IOException | DataGeneratorException e) {
    errorRecordHandler.onError(new OnRecordErrorException(record, Errors.PUBSUB_06, e.toString(), e));
    return;
  }

  ByteString data = ByteString.copyFrom(os.toByteArray());

  Map<String, String> attributes = new HashMap<>();
  Record.Header header = record.getHeader();
  header.getAttributeNames().forEach(k -> attributes.put(k, header.getAttribute(k)));

  PubsubMessage message = PubsubMessage.newBuilder().setData(data).putAllAttributes(attributes).build();

  ApiFuture<String> messageIdFuture = publisher.publish(message);
  pendingMessages.add(new PendingMessage(record, messageIdFuture));
}
 

/**
 * Populates HTTP response headers to the configured location
 *
 * @param record current record to populate
 * @param response HTTP response
 * @throws StageException when writing headers to a field path that already exists
 */
private Field createResponseHeaders(Record record, Response response) throws StageException {
  if (conf.headerOutputLocation == HeaderOutputLocation.NONE) {
    return null;
  }

  Record.Header header = record.getHeader();
  header.setAttribute(REQUEST_STATUS_CONFIG_NAME,String.format("%d",response.getStatus()));

  if (conf.headerOutputLocation == HeaderOutputLocation.FIELD) {
    return createResponseHeaderField(record, response);
  } else if (conf.headerOutputLocation == HeaderOutputLocation.HEADER) {
    createResponseHeaderToRecordHeader(response, header);
    return null;
  }
  return null;
}
 

@NotNull
private Record makeRecord(Map<String, Pair<Field.Type, Object>> fieldMap) {
  Record record = RecordCreator.create();
  Record.Header header = record.getHeader();
  LinkedHashMap<String, Field> fields = new LinkedHashMap<>();
  for (Map.Entry<String, Pair<Field.Type, Object>> entry : fieldMap.entrySet()) {
    String fieldName = entry.getKey();
    Field.Type fieldType = entry.getValue().getLeft();
    Field field = Field.create(fieldType, entry.getValue().getRight());
    if (fieldType == Field.Type.DECIMAL) {
      field.setAttribute(HeaderAttributeConstants.ATTR_SCALE, SCALE);
      field.setAttribute(HeaderAttributeConstants.ATTR_PRECISION, PRECISION);
    }
    fields.put(fieldName, field);
  }
  record.set(Field.create(fields));
  header.setAttribute("table", tableName);
  return record;
}
 

public void copyErrorFrom(Record record) {
  Record.Header header = record.getHeader();
  setError(
    header.getErrorStage(),
    header.getErrorStageLabel(),
    header.getErrorCode(),
    header.getErrorMessage(),
    header.getErrorTimestamp(),
    header.getErrorStackTrace()
  );
}
 

public static <C extends Processor> void verifyRecordHeaderAttribute(
    Class<C> clazz,
    Processor processor,
    Record record
) throws StageException {
  ProcessorRunner runner = new ProcessorRunner.Builder(clazz, processor)
      .addOutputLane("lane")
      .build();

  runner.runInit();

  // This header should be always removed
  record.getHeader().setAttribute("remove", "please remove me");

  StageRunner.Output output;
  try {
    List<Record> input = Collections.singletonList(record);
    output = runner.runProcess(input);
  } finally {
    runner.runDestroy();
  }
  List<Record> records = output.getRecords().get("lane");
  assertEquals(1, records.size());

  Record.Header outputHeader = records.get(0).getHeader();
  assertEquals(1, outputHeader.getAttributeNames().size());

  final String key = "key1";
  final String value = "value1";
  assertEquals(value, outputHeader.getAttribute(key));

  assertFalse(outputHeader.getAttributeNames().contains("remove"));
}
 

private void testMultiStatusResponse(
    String httpServerUrl,
    List<Record> requestRecords
) {
  Response response = ClientBuilder.newClient()
      .property(ClientProperties.SUPPRESS_HTTP_COMPLIANCE_VALIDATION, true)
      .property(HttpUrlConnectorProvider.SET_METHOD_WORKAROUND, true)
      .target(httpServerUrl)
      .request()
      .header(Constants.X_SDC_APPLICATION_ID_HEADER, "id")
      .method(
          "POST",
          Entity.json("{\"f1\": \"abc\", \"f2\": \"xyz\"}\n{\"f1\": \"abc\", \"f2\": \"xyz\", \"sendToError\": \"Sample Error message\"}")
      );

  // Test Request Records
  Assert.assertEquals(2, requestRecords.size());
  Record.Header payloadRecord = requestRecords.get(0).getHeader();
  Assert.assertEquals("POST", payloadRecord.getAttribute(RestServiceReceiver.METHOD_HEADER));
  Assert.assertNull(payloadRecord.getAttribute(RestServiceReceiver.EMPTY_PAYLOAD_RECORD_HEADER_ATTR_NAME));
  Assert.assertEquals("abc", requestRecords.get(0).get("/f1").getValue());
  Assert.assertEquals("xyz", requestRecords.get(0).get("/f2").getValue());

  // Test Response from REST Service
  Assert.assertEquals(207, response.getStatus());
  ResponseEnvelope responseBody = response.readEntity(ResponseEnvelope.class);
  Assert.assertNotNull(responseBody);
  Assert.assertEquals(207, responseBody.getHttpStatusCode());
  Assert.assertNotNull(responseBody.getData());
  Assert.assertEquals(1, responseBody.getData().size());
  Assert.assertNotNull(responseBody.getError());
  Assert.assertEquals(1, responseBody.getError().size());
  Assert.assertNotNull(responseBody.getErrorMessage());
}
 

@Test
public void testHttpGetPrefixedHeaderOutput() throws Exception {
  HttpProcessorConfig conf = new HttpProcessorConfig();
  conf.httpMethod = HttpMethod.GET;
  conf.dataFormat = DataFormat.TEXT;
  conf.resourceUrl = getBaseUri() + "test/get";
  conf.headerOutputLocation = HeaderOutputLocation.HEADER;
  conf.headerAttributePrefix = "test-prefix-";

  List<Record> records = createRecords("test/get");
  ProcessorRunner runner = createProcessorRunner(conf);
  try {
    StageRunner.Output output = runner.runProcess(records);
    assertEquals("{\"hello\":\"world!\"}", getOutputField(output).getValueAsString());
    Record.Header header = getOutputRecord(output).getHeader();
    assertEquals(
        "StreamSets",
        header.getAttribute(conf.headerAttributePrefix + "x-test-header")
    );
    assertEquals(
        "[a, b]",
        header.getAttribute(conf.headerAttributePrefix + "x-list-header")
    );
  } finally {
    runner.runDestroy();
  }
}
 

public void setColumnSpecificHeaders(
    Record record,
    Set<String> knownTableNames,
    ResultSetMetaData metaData,
    String jdbcNameSpacePrefix
) throws SQLException {
  Record.Header header = record.getHeader();
  Set<String> tableNames = new HashSet<>();

  for (int i=1; i<=metaData.getColumnCount(); i++) {
    header.setAttribute(jdbcNameSpacePrefix + metaData.getColumnLabel(i) + ".jdbcType", String.valueOf(metaData.getColumnType(i)));

    // Additional headers per various types
    switch(metaData.getColumnType(i)) {
      case Types.DECIMAL:
      case Types.NUMERIC:
        header.setAttribute(jdbcNameSpacePrefix + metaData.getColumnLabel(i) + ".scale", String.valueOf(metaData.getScale(i)));
        header.setAttribute(jdbcNameSpacePrefix + metaData.getColumnLabel(i) + ".precision", String.valueOf(metaData.getPrecision(i)));
        break;
    }

    String tableName = metaData.getTableName(i);

    // Store the column's table name (if not empty)
    if (StringUtils.isNotEmpty(tableName)) {
      tableNames.add(tableName);
    }
  }

  if (tableNames.isEmpty()) {
    tableNames.addAll(knownTableNames);
  }

  header.setAttribute(jdbcNameSpacePrefix + "tables", Joiner.on(",").join(tableNames));
}
 

@Test
public void testMillisecondsToDateTime() throws Exception {

  Record.Header header = Mockito.mock(Record.Header.class);
  Record record = Mockito.mock(Record.class);
  Mockito.when(record.getHeader()).thenReturn(header);
  long val = System.currentTimeMillis();
  Mockito.when(record.get(Mockito.eq("/epochMS"))).thenReturn(Field.create(val));

  RecordEL.setRecordInContext(variables, record);

  Assert.assertEquals(new Date(val), eval.eval(variables, "${time:millisecondsToDateTime(record:value('/epochMS'))}", Date.class));
}
 

@Test
public void testDateTimeToMilliSeconds() throws Exception {
  Record.Header header = Mockito.mock(Record.Header.class);
  Record record = Mockito.mock(Record.class);
  Mockito.when(record.getHeader()).thenReturn(header);
  Long val = System.currentTimeMillis();
  Mockito.when(record.get(Mockito.eq("/date"))).thenReturn(Field.create(Field.Type.DATETIME, new Date(val)));
  RecordEL.setRecordInContext(variables, record);
  Assert.assertEquals(val, eval.eval(variables, "${time:dateTimeToMilliseconds(record:value('/date'))}", Long.class));
}
 

private void updateRecordWithCellValues(Row row, Record record) throws DataParserException {
  LinkedHashMap<String, Field> output = new LinkedHashMap<>();
  String sheetName = row.getSheet().getSheetName();
  String columnHeader;
  Set<String> unsupportedCellTypes = new HashSet<>();
  for (int columnNum = row.getFirstCellNum(); columnNum < row.getLastCellNum(); columnNum++) {
    if (headers.isEmpty()) {
      columnHeader = String.valueOf(columnNum);
    } else {
      if (columnNum >= headers.get(sheetName).size() || headers.get(sheetName).get(columnNum) == null) {
        // The current cell doesn't hae any associated header, which we conditionally skip
        if(settings.shouldSkipCellsWithNoHeader()) {
          continue;
        }

        columnHeader = String.valueOf(columnNum);
      } else {
        columnHeader = headers.get(sheetName).get(columnNum).getValueAsString();
      }
    }

    Cell cell = row.getCell(columnNum, Row.MissingCellPolicy.CREATE_NULL_AS_BLANK);
    try {
      output.put(columnHeader, Cells.parseCell(cell, this.evaluator));
    } catch (ExcelUnsupportedCellTypeException e) {
      output.put(columnHeader, Cells.parseCellAsString(cell));
      unsupportedCellTypes.add(e.getCellType().name());
    }
  }


  // Set interesting metadata about the row
  Record.Header hdr = record.getHeader();
  hdr.setAttribute("worksheet", row.getSheet().getSheetName());
  hdr.setAttribute("row",  Integer.toString(row.getRowNum()));
  hdr.setAttribute("firstCol", Integer.toString(row.getFirstCellNum()));
  hdr.setAttribute("lastCol", Integer.toString(row.getLastCellNum()));
  record.set(Field.createListMap(output));
  if (unsupportedCellTypes.size() > 0) {
    throw new RecoverableDataParserException(record, Errors.EXCEL_PARSER_05, StringUtils.join(unsupportedCellTypes, ", "));
  }
}
 

private void setHeaderIfNotNull(Record.Header header, String key, Object val) {
  if (val != null) {
    header.setAttribute(key, convToString(val));
  }
}
 

@Override
public void wrapRecordHeaders(Record.Header header, String offset) {
  //NO OP
}
 

@Test
public void testHeaderMethods() {
  RecordImpl record = new RecordImpl("stage", "source", null, null);
  Record.Header header = record.getHeader();
  Assert.assertNull(header.getRaw());
  Assert.assertNull(header.getRawMimeType());

  record = new RecordImpl("stage", "source", new byte[0], "M");
  header = record.getHeader();
  Assert.assertArrayEquals(new byte[0], header.getRaw());
  Assert.assertEquals("M", header.getRawMimeType());

  Assert.assertEquals("stage", header.getStageCreator());
  Assert.assertEquals("source", header.getSourceId());
  Assert.assertNull(header.getStagesPath());

  record.addStageToStagePath("x");
  Assert.assertEquals("x", header.getStagesPath());

  record.addStageToStagePath("y");
  Assert.assertEquals("x:y", header.getStagesPath());

  Assert.assertTrue(header.getAttributeNames().isEmpty());
  try {
    header.setAttribute("a", null);
    Assert.fail();
  } catch (NullPointerException ex) {
    //expected
  }
  Assert.assertTrue(header.getAttributeNames().isEmpty());
  HeaderImpl headerImpl = ((HeaderImpl)header);
  Assert.assertTrue(headerImpl.getValues().isEmpty());
  header.setAttribute("a", "A");
  Assert.assertEquals(ImmutableSet.of("a"), header.getAttributeNames());
  Assert.assertEquals("A", header.getAttribute("a"));
  Assert.assertEquals(1, headerImpl.getValues().size());
  Assert.assertEquals("A", headerImpl.getValues().get("a"));
  header.deleteAttribute("a");
  Assert.assertNull(header.getAttribute("a"));
  Assert.assertTrue(header.getAttributeNames().isEmpty());
  Assert.assertTrue(headerImpl.getValues().isEmpty());

  record.toString();
}