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

The following examples show how to use com.streamsets.pipeline.api.Record#Header . You can vote up the ones you like or vote down the ones you don't like, and go to the original project or source file by following the links above each example. You may check out the related API usage on the sidebar.
Example 1
Source File: SqsConsumerWorkerCallable.java    From datacollector with Apache License 2.0 7 votes vote down vote up
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;
  }
}
 
Example 2
Source File: RecordCloner.java    From datacollector with Apache License 2.0 6 votes vote down vote up
/**
 * 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);
  }
}
 
Example 3
Source File: TestRestServicePushSource.java    From datacollector with Apache License 2.0 6 votes vote down vote up
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"));
}
 
Example 4
Source File: TestJdbcMetadata.java    From datacollector with Apache License 2.0 6 votes vote down vote up
@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;
}
 
Example 5
Source File: HttpProcessorIT.java    From datacollector with Apache License 2.0 6 votes vote down vote up
@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();
  }
}
 
Example 6
Source File: TestDataLakeGeneratorManager.java    From datacollector with Apache License 2.0 6 votes vote down vote up
@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));
}
 
Example 7
Source File: PubSubTarget.java    From datacollector with Apache License 2.0 6 votes vote down vote up
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));
}
 
Example 8
Source File: HttpProcessor.java    From datacollector with Apache License 2.0 6 votes vote down vote up
/**
 * 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;
}
 
Example 9
Source File: TestJdbcMetadata.java    From datacollector with Apache License 2.0 6 votes vote down vote up
@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;
}
 
Example 10
Source File: HeaderImpl.java    From datacollector with Apache License 2.0 5 votes vote down vote up
public void copyErrorFrom(Record record) {
  Record.Header header = record.getHeader();
  setError(
    header.getErrorStage(),
    header.getErrorStageLabel(),
    header.getErrorCode(),
    header.getErrorMessage(),
    header.getErrorTimestamp(),
    header.getErrorStackTrace()
  );
}
 
Example 11
Source File: ScriptingProcessorTestUtil.java    From datacollector with Apache License 2.0 5 votes vote down vote up
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"));
}
 
Example 12
Source File: TestRestServicePushSource.java    From datacollector with Apache License 2.0 5 votes vote down vote up
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());
}
 
Example 13
Source File: HttpProcessorIT.java    From datacollector with Apache License 2.0 5 votes vote down vote up
@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();
  }
}
 
Example 14
Source File: JdbcUtil.java    From datacollector with Apache License 2.0 5 votes vote down vote up
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));
}
 
Example 15
Source File: TestTimeNewEL.java    From datacollector with Apache License 2.0 5 votes vote down vote up
@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));
}
 
Example 16
Source File: TestTimeNewEL.java    From datacollector with Apache License 2.0 5 votes vote down vote up
@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));
}
 
Example 17
Source File: WorkbookParser.java    From datacollector with Apache License 2.0 4 votes vote down vote up
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, ", "));
  }
}
 
Example 18
Source File: RabbitSource.java    From datacollector with Apache License 2.0 4 votes vote down vote up
private void setHeaderIfNotNull(Record.Header header, String key, Object val) {
  if (val != null) {
    header.setAttribute(key, convToString(val));
  }
}
 
Example 19
Source File: CompressionDataParser.java    From datacollector with Apache License 2.0 4 votes vote down vote up
@Override
public void wrapRecordHeaders(Record.Header header, String offset) {
  //NO OP
}
 
Example 20
Source File: TestRecordImpl.java    From datacollector with Apache License 2.0 4 votes vote down vote up
@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();
}