Enforce static schemas at preprocessor level

Author: kx-chen

astra/src/main/java/com/slack/astra/bulkIngestApi/BulkIngestApi.java

@@ -6,6 +6,7 @@
 import com.linecorp.armeria.common.HttpStatus;
 import com.linecorp.armeria.server.annotation.Post;
 import com.slack.astra.bulkIngestApi.opensearch.BulkApiRequestParser;
+import com.slack.astra.proto.config.AstraConfigs;
 import com.slack.astra.proto.schema.Schema;
 import com.slack.service.murron.trace.Trace;
 import io.micrometer.core.instrument.Counter;
@@ -25,8 +26,13 @@
  */
 public class BulkIngestApi {
   private static final Logger LOG = LoggerFactory.getLogger(BulkIngestApi.class);
+
+  public static final String SCHEMA_ENFORCEMENT_FLAG = "astra.schema.enforcement.enabled";
+
   private final BulkIngestKafkaProducer bulkIngestKafkaProducer;
   private final DatasetRateLimitingService datasetRateLimitingService;
+  private final Schema.IngestSchema schema;
+  private final AstraConfigs.SchemaMode schemaMode;
   private final MeterRegistry meterRegistry;
   private final Counter incomingByteTotal;
   private final Counter incomingDocsTotal;
@@ -36,7 +42,6 @@ public class BulkIngestApi {
   private final String BULK_INGEST_ERROR = "astra_preprocessor_error";
   private final String BULK_INGEST_TIMER = "astra_preprocessor_bulk_ingest";
   private final int rateLimitExceededErrorCode;
-  private final Schema.IngestSchema schema;
 
   private final Counter bulkIngestErrorCounter;
 
@@ -45,10 +50,13 @@ public BulkIngestApi(
       DatasetRateLimitingService datasetRateLimitingService,
       MeterRegistry meterRegistry,
       int rateLimitExceededErrorCode,
-      Schema.IngestSchema schema) {
+      Schema.IngestSchema schema,
+      AstraConfigs.SchemaMode schemaMode) {
 
     this.bulkIngestKafkaProducer = bulkIngestKafkaProducer;
     this.datasetRateLimitingService = datasetRateLimitingService;
+    this.schema = schema;
+    this.schemaMode = schemaMode;
     this.meterRegistry = meterRegistry;
     this.incomingByteTotal = meterRegistry.counter(BULK_INGEST_INCOMING_BYTE_TOTAL);
     this.incomingDocsTotal = meterRegistry.counter(BULK_INGEST_INCOMING_BYTE_DOCS);
@@ -58,7 +66,6 @@ public BulkIngestApi(
     } else {
       this.rateLimitExceededErrorCode = rateLimitExceededErrorCode;
     }
-    this.schema = schema;
     this.bulkIngestErrorCounter = meterRegistry.counter(BULK_INGEST_ERROR);
   }
 
@@ -75,7 +82,11 @@ public HttpResponse addDocument(String bulkRequest) {
       incomingByteTotal.increment(bulkRequestBytes.length);
       Map<String, List<Trace.Span>> docs = Map.of();
       try {
-        docs = BulkApiRequestParser.parseRequest(bulkRequestBytes, schema);
+        if (Boolean.getBoolean(SCHEMA_ENFORCEMENT_FLAG)) {
+          docs = BulkApiRequestParser.parseRequest(bulkRequestBytes, schema, schemaMode);
+        } else {
+          docs = BulkApiRequestParser.parseRequest(bulkRequestBytes, schema);
+        }
       } catch (Exception e) {
         LOG.error("Request failed ", e);
         bulkIngestErrorCounter.increment();

astra/src/main/java/com/slack/astra/bulkIngestApi/opensearch/BulkApiRequestParser.java

@@ -2,8 +2,10 @@
 
 import com.google.common.annotations.VisibleForTesting;
 import com.google.protobuf.ByteString;
+import com.slack.astra.bulkIngestApi.BulkIngestApi;
 import com.slack.astra.logstore.LogMessage;
 import com.slack.astra.logstore.schema.ReservedFields;
+import com.slack.astra.proto.config.AstraConfigs;
 import com.slack.astra.proto.schema.Schema;
 import com.slack.astra.writer.SpanFormatter;
 import com.slack.service.murron.trace.Trace;
@@ -37,7 +39,13 @@ public class BulkApiRequestParser {
 
   public static Map<String, List<Trace.Span>> parseRequest(
       byte[] postBody, Schema.IngestSchema schema) throws IOException {
-    return convertIndexRequestToTraceFormat(parseBulkRequest(postBody), schema);
+    return parseRequest(postBody, schema, AstraConfigs.SchemaMode.SCHEMA_MODE_DYNAMIC);
+  }
+
+  public static Map<String, List<Trace.Span>> parseRequest(
+      byte[] postBody, Schema.IngestSchema schema, AstraConfigs.SchemaMode schemaMode)
+      throws IOException {
+    return convertIndexRequestToTraceFormat(parseBulkRequest(postBody), schema, schemaMode);
   }
 
   /**
@@ -64,6 +72,14 @@ public static long getTimestampFromIngestDocument(Map<String, Object> sourceAndM
   @VisibleForTesting
   public static Trace.Span fromIngestDocument(
       IngestDocument ingestDocument, Schema.IngestSchema schema) {
+    return fromIngestDocument(ingestDocument, schema, AstraConfigs.SchemaMode.SCHEMA_MODE_DYNAMIC);
+  }
+
+  @VisibleForTesting
+  public static Trace.Span fromIngestDocument(
+      IngestDocument ingestDocument,
+      Schema.IngestSchema schema,
+      AstraConfigs.SchemaMode schemaMode) {
 
     Map<String, Object> sourceAndMetadata = ingestDocument.getSourceAndMetadata();
 
@@ -140,8 +156,21 @@ public static Trace.Span fromIngestDocument(
     sourceAndMetadata.remove(IngestDocument.Metadata.ID.getFieldName());
     sourceAndMetadata.remove(IngestDocument.Metadata.INDEX.getFieldName());
 
+    boolean dropUnknown =
+        Boolean.getBoolean(BulkIngestApi.SCHEMA_ENFORCEMENT_FLAG)
+            && schemaMode == AstraConfigs.SchemaMode.SCHEMA_MODE_DROP_UNKNOWN;
+
     boolean tagsContainServiceName = false;
     for (Map.Entry<String, Object> kv : sourceAndMetadata.entrySet()) {
+      if (dropUnknown) {
+        if (!schema.containsFields(kv.getKey())) {
+          continue;
+        }
+        Schema.SchemaField fieldDef = schema.getFieldsMap().get(kv.getKey());
+        if (!SpanFormatter.isTypeCompatible(kv.getValue(), fieldDef.getType())) {
+          continue;
+        }
+      }
       if (!tagsContainServiceName && kv.getKey().equals(SERVICE_NAME_KEY)) {
         tagsContainServiceName = true;
       }
@@ -151,7 +180,7 @@ public static Trace.Span fromIngestDocument(
         spanBuilder.addAllTags(tags);
       }
     }
-    if (!tagsContainServiceName) {
+    if (!tagsContainServiceName && (!dropUnknown || schema.containsFields(SERVICE_NAME_KEY))) {
       spanBuilder.addTags(
           Trace.KeyValue.newBuilder()
               .setKey(SERVICE_NAME_KEY)
@@ -165,6 +194,14 @@ public static Trace.Span fromIngestDocument(
 
   protected static Map<String, List<Trace.Span>> convertIndexRequestToTraceFormat(
       List<IndexRequest> indexRequests, Schema.IngestSchema schema) {
+    return convertIndexRequestToTraceFormat(
+        indexRequests, schema, AstraConfigs.SchemaMode.SCHEMA_MODE_DYNAMIC);
+  }
+
+  protected static Map<String, List<Trace.Span>> convertIndexRequestToTraceFormat(
+      List<IndexRequest> indexRequests,
+      Schema.IngestSchema schema,
+      AstraConfigs.SchemaMode schemaMode) {
     // key - index. value - list of docs to be indexed
     Map<String, List<Trace.Span>> indexDocs = new HashMap<>();
 
@@ -175,7 +212,7 @@ protected static Map<String, List<Trace.Span>> convertIndexRequestToTraceFormat(
       }
       IngestDocument ingestDocument = convertRequestToDocument(indexRequest);
       List<Trace.Span> docs = indexDocs.computeIfAbsent(index, key -> new ArrayList<>());
-      docs.add(BulkApiRequestParser.fromIngestDocument(ingestDocument, schema));
+      docs.add(BulkApiRequestParser.fromIngestDocument(ingestDocument, schema, schemaMode));
     }
     return indexDocs;
   }

astra/src/main/java/com/slack/astra/server/Astra.java

@@ -601,7 +601,8 @@ private static Set<Service> getServices(
               datasetRateLimitingService,
               meterRegistry,
               preprocessorConfig.getRateLimitExceededErrorCode(),
-              schema);
+              schema,
+              preprocessorConfig.getSchemaMode());
       armeriaServiceBuilder.withAnnotatedService(openSearchBulkApiService);
       services.add(armeriaServiceBuilder.build());
     }

astra/src/main/java/com/slack/astra/writer/SpanFormatter.java

@@ -39,6 +39,41 @@ public static Timestamp parseDate(String dateStr, Schema.SchemaFieldType type) {
         .build();
   }
 
+  public static boolean isTypeCompatible(Object value, Schema.SchemaFieldType type) {
+    try {
+      switch (type) {
+        case KEYWORD, TEXT, IP, BINARY, BOOLEAN -> {
+          return true;
+        }
+        case DATE -> {
+          Instant.parse(value.toString());
+          return true;
+        }
+        case DOUBLE -> {
+          Double.parseDouble(value.toString());
+          return true;
+        }
+        case FLOAT, HALF_FLOAT -> {
+          Float.parseFloat(value.toString());
+          return true;
+        }
+        case INTEGER, SHORT, BYTE -> {
+          Integer.parseInt(value.toString());
+          return true;
+        }
+        case LONG, SCALED_LONG -> {
+          Long.parseLong(value.toString());
+          return true;
+        }
+        default -> {
+          return true;
+        }
+      }
+    } catch (Exception e) {
+      return false;
+    }
+  }
+
   public static Trace.KeyValue makeTraceKV(String key, Object value, Schema.SchemaFieldType type) {
     Trace.KeyValue.Builder tagBuilder = Trace.KeyValue.newBuilder();
     tagBuilder.setKey(key);

astra/src/main/proto/astra_configs.proto

@@ -14,6 +14,15 @@ enum NodeRole {
   PREPROCESSOR = 5;
 };
 
+// Schema enforcement mode for the preprocessor.
+// Controls how fields not defined in the schema are handled during ingestion.
+enum SchemaMode {
+  // Default: infer types for unknown fields using defaults/heuristics
+  SCHEMA_MODE_DYNAMIC = 0;
+  // Drop unknown fields silently - only index fields defined in schema
+  SCHEMA_MODE_DROP_UNKNOWN = 1;
+};
+
 // Astra is a single binary consisting of multiple components.
 // AstraConfig is the uber config object for all of Astra.
 // This config object controls the role a node plays and it's config.
@@ -328,4 +337,9 @@ message PreprocessorConfig {
   // refresh of the dataset rate limits, whether or not a ZK event
   // has triggered it
   int32 dataset_rate_limit_period_secs = 14;
+
+  // Schema enforcement mode - controls how unknown fields are handled
+  // SCHEMA_MODE_DYNAMIC (default): infer types for unknown fields
+  // SCHEMA_MODE_DROP_UNKNOWN: drop fields not defined in schema
+  SchemaMode schema_mode = 15;
 }