Add support for STRING and BYTES to query step (#106)

datastax · Jul 15, 2023 · 364f215 · 364f215
1 parent edb2daa
commit 364f215
Show file tree

Hide file tree

Showing 2 changed files with 44 additions and 0 deletions.
diff --git a/streaming-ai/src/main/java/com/datastax/oss/streaming/ai/QueryStep.java b/streaming-ai/src/main/java/com/datastax/oss/streaming/ai/QueryStep.java
@@ -90,6 +90,9 @@ public void process(TransformContext transformContext) {
   private Object getField(
       String key, String field, TransformSchemaType keySchemaType, Object keyObject) {
     String fieldName = field.substring((key.length() + 1));
+    if (keyObject instanceof Map) {
+      return ((Map<String, Object>) keyObject).get(fieldName);
+    }
     switch (keySchemaType) {
       case AVRO:
         GenericRecord avroRecord = (GenericRecord) keyObject;

diff --git a/streaming-ai/src/test/java/com/datastax/oss/streaming/ai/QueryStepTest.java b/streaming-ai/src/test/java/com/datastax/oss/streaming/ai/QueryStepTest.java
@@ -33,6 +33,8 @@
 import org.apache.pulsar.client.api.schema.RecordSchemaBuilder;
 import org.apache.pulsar.client.api.schema.SchemaBuilder;
 import org.apache.pulsar.client.impl.schema.AutoConsumeSchema;
+import org.apache.pulsar.common.schema.KeyValue;
+import org.apache.pulsar.common.schema.KeyValueEncodingType;
 import org.apache.pulsar.common.schema.SchemaInfo;
 import org.apache.pulsar.common.schema.SchemaType;
 import org.apache.pulsar.functions.api.Record;
@@ -134,6 +136,45 @@ public List<Map<String, String>> fetchData(String query, List<Object> params) {
     Utils.process(record, queryStep);
   }
 
+  @Test
+  void testKVStringJson() throws Exception {
+    Schema<KeyValue<String, String>> keyValueSchema =
+        Schema.KeyValue(Schema.STRING, Schema.STRING, KeyValueEncodingType.SEPARATED);
+
+    String key = "{\"keyField1\": \"key1\", \"keyField2\": \"key2\", \"keyField3\": \"key3\"}";
+    String value =
+        "{\"valueField1\": \"value1\", \"valueField2\": \"value2\", \"valueField3\": \"value3\"}";
+
+    KeyValue<String, String> keyValue = new KeyValue<>(key, value);
+
+    Record<GenericObject> record =
+        new Utils.TestRecord<>(
+            keyValueSchema,
+            AutoConsumeSchema.wrapPrimitiveObject(keyValue, SchemaType.KEY_VALUE, new byte[] {}),
+            null);
+
+    QueryStepDataSource dataSource =
+        new QueryStepDataSource() {
+          @Override
+          public List<Map<String, String>> fetchData(String query, List<Object> params) {
+            assertEquals(query, "select 1");
+            List<Object> expectedParams = List.of("value1", "key2");
+            assertEquals(params, expectedParams);
+            return List.of(Map.of());
+          }
+        };
+    List<String> fields = Arrays.asList("value.valueField1", "key.keyField2");
+    QueryStep queryStep =
+        QueryStep.builder()
+            .dataSource(dataSource)
+            .outputFieldName("value.result")
+            .query("select 1")
+            .fields(fields)
+            .build();
+
+    Utils.process(record, queryStep);
+  }
+
   @Test
   void testKVAvro() throws Exception {
     Record<GenericObject> record = Utils.createTestAvroKeyValueRecord();