Field-level Redaction (#1122)

Co-authored-by: Bryan Burkholder <[email protected]>

astra/src/main/java/com/slack/astra/chunk/ReadOnlyChunkImpl.java

@@ -14,6 +14,7 @@
 import com.slack.astra.metadata.cache.CacheSlotMetadata;
 import com.slack.astra.metadata.cache.CacheSlotMetadataStore;
 import com.slack.astra.metadata.core.AstraMetadataStoreChangeListener;
+import com.slack.astra.metadata.fieldredaction.FieldRedactionMetadataStore;
 import com.slack.astra.metadata.replica.ReplicaMetadata;
 import com.slack.astra.metadata.replica.ReplicaMetadataStore;
 import com.slack.astra.metadata.schema.ChunkSchema;
@@ -50,6 +51,7 @@
  * BlobFs.
  */
 public class ReadOnlyChunkImpl<T> implements Chunk<T> {
+
   private static final Logger LOG = LoggerFactory.getLogger(ReadOnlyChunkImpl.class);
   public static final String ASTRA_S3_STREAMING_FLAG = "astra.s3Streaming.enabled";
 
@@ -71,6 +73,7 @@ public class ReadOnlyChunkImpl<T> implements Chunk<T> {
   private final SnapshotMetadataStore snapshotMetadataStore;
   private final SearchMetadataStore searchMetadataStore;
   private CacheNodeAssignmentStore cacheNodeAssignmentStore;
+  private final FieldRedactionMetadataStore fieldRedactionMetadataStore;
   private final MeterRegistry meterRegistry;
   private final BlobStore blobStore;
 
@@ -104,7 +107,8 @@ public ReadOnlyChunkImpl(
       SearchMetadataStore searchMetadataStore,
       CacheNodeAssignmentStore cacheNodeAssignmentStore,
       CacheNodeAssignment assignment,
-      SnapshotMetadata snapshotMetadata)
+      SnapshotMetadata snapshotMetadata,
+      FieldRedactionMetadataStore fieldRedactionMetadataStore)
       throws Exception {
     this(
         curatorFramework,
@@ -117,7 +121,8 @@ public ReadOnlyChunkImpl(
         cacheSlotMetadataStore,
         replicaMetadataStore,
         snapshotMetadataStore,
-        searchMetadataStore);
+        searchMetadataStore,
+        fieldRedactionMetadataStore);
     this.assignment = assignment;
     this.lastKnownAssignmentState = assignment.state;
     this.snapshotMetadata = snapshotMetadata;
@@ -135,7 +140,8 @@ public ReadOnlyChunkImpl(
       CacheSlotMetadataStore cacheSlotMetadataStore,
       ReplicaMetadataStore replicaMetadataStore,
       SnapshotMetadataStore snapshotMetadataStore,
-      SearchMetadataStore searchMetadataStore)
+      SearchMetadataStore searchMetadataStore,
+      FieldRedactionMetadataStore fieldRedactionMetadataStore)
       throws Exception {
     this.meterRegistry = meterRegistry;
     this.blobStore = blobStore;
@@ -147,6 +153,7 @@ public ReadOnlyChunkImpl(
     this.replicaMetadataStore = replicaMetadataStore;
     this.snapshotMetadataStore = snapshotMetadataStore;
     this.searchMetadataStore = searchMetadataStore;
+    this.fieldRedactionMetadataStore = fieldRedactionMetadataStore;
 
     if (!Boolean.getBoolean(ASTRA_NG_DYNAMIC_CHUNK_SIZES_FLAG)) {
       CacheSlotMetadata cacheSlotMetadata =
@@ -230,7 +237,8 @@ public void downloadChunkData() {
         this.logSearcher =
             (LogIndexSearcher<T>)
                 new LogIndexSearcherImpl(
-                    LogIndexSearcherImpl.searcherManagerFromChunkId(chunkInfo.chunkId, blobStore),
+                    LogIndexSearcherImpl.searcherManagerFromChunkId(
+                        chunkInfo.chunkId, blobStore, fieldRedactionMetadataStore),
                     chunkSchema.fieldDefMap);
       } else {
         // get data directory
@@ -264,7 +272,8 @@ public void downloadChunkData() {
         this.logSearcher =
             (LogIndexSearcher<T>)
                 new LogIndexSearcherImpl(
-                    LogIndexSearcherImpl.searcherManagerFromPath(dataDirectory),
+                    LogIndexSearcherImpl.searcherManagerFromPath(
+                        dataDirectory, fieldRedactionMetadataStore),
                     chunkSchema.fieldDefMap);
       }
 
@@ -414,7 +423,8 @@ private void handleChunkAssignment(CacheSlotMetadata cacheSlotMetadata) {
       this.logSearcher =
           (LogIndexSearcher<T>)
               new LogIndexSearcherImpl(
-                  LogIndexSearcherImpl.searcherManagerFromPath(dataDirectory),
+                  LogIndexSearcherImpl.searcherManagerFromPath(
+                      dataDirectory, fieldRedactionMetadataStore),
                   chunkSchema.fieldDefMap);
 
       // we first mark the slot LIVE before registering the search metadata as available

astra/src/main/java/com/slack/astra/chunk/RecoveryChunkFactoryImpl.java

@@ -5,6 +5,7 @@
 
 import com.slack.astra.logstore.LogStore;
 import com.slack.astra.logstore.LuceneIndexStoreImpl;
+import com.slack.astra.metadata.fieldredaction.FieldRedactionMetadataStore;
 import com.slack.astra.metadata.search.SearchMetadataStore;
 import com.slack.astra.metadata.snapshot.SnapshotMetadataStore;
 import com.slack.astra.proto.config.AstraConfigs;
@@ -24,6 +25,7 @@ public class RecoveryChunkFactoryImpl<T> implements ChunkFactory<T> {
   private final MeterRegistry meterRegistry;
   private final SearchMetadataStore searchMetadataStore;
   private final SnapshotMetadataStore snapshotMetadataStore;
+  private final FieldRedactionMetadataStore fieldRedactionMetadataStore;
   private final SearchContext searchContext;
   private final AstraConfigs.IndexerConfig indexerConfig;
   private String kafkaPartitionId = null;
@@ -34,13 +36,15 @@ public RecoveryChunkFactoryImpl(
       MeterRegistry meterRegistry,
       SearchMetadataStore searchMetadataStore,
       SnapshotMetadataStore snapshotMetadataStore,
+      FieldRedactionMetadataStore fieldRedactionMetadataStore,
       SearchContext searchContext) {
     checkNotNull(indexerConfig, "indexerConfig can't be null");
     this.indexerConfig = indexerConfig;
     this.chunkDataPrefix = chunkDataPrefix;
     this.meterRegistry = meterRegistry;
     this.searchMetadataStore = searchMetadataStore;
     this.snapshotMetadataStore = snapshotMetadataStore;
+    this.fieldRedactionMetadataStore = fieldRedactionMetadataStore;
     this.searchContext = searchContext;
   }
 
@@ -49,9 +53,13 @@ public ReadWriteChunk<T> makeChunk() throws IOException {
     ensureNonNullString(kafkaPartitionId, "kafkaPartitionId can't be null and should be set.");
     ensureNonNullString(indexerConfig.getDataDirectory(), "The data directory shouldn't be empty");
     final File dataDirectory = new File(indexerConfig.getDataDirectory());
+
     LogStore logStore =
         LuceneIndexStoreImpl.makeLogStore(
-            dataDirectory, indexerConfig.getLuceneConfig(), meterRegistry);
+            dataDirectory,
+            indexerConfig.getLuceneConfig(),
+            meterRegistry,
+            fieldRedactionMetadataStore);
 
     return new RecoveryChunkImpl<>(
         logStore,

astra/src/main/java/com/slack/astra/chunkManager/CachingChunkManager.java

@@ -14,6 +14,7 @@
 import com.slack.astra.metadata.cache.CacheNodeMetadataStore;
 import com.slack.astra.metadata.cache.CacheSlotMetadataStore;
 import com.slack.astra.metadata.core.AstraMetadataStoreChangeListener;
+import com.slack.astra.metadata.fieldredaction.FieldRedactionMetadataStore;
 import com.slack.astra.metadata.replica.ReplicaMetadataStore;
 import com.slack.astra.metadata.search.SearchMetadataStore;
 import com.slack.astra.metadata.snapshot.SnapshotMetadata;
@@ -56,6 +57,7 @@ public class CachingChunkManager<T> extends ChunkManagerBase<T> {
   private SnapshotMetadataStore snapshotMetadataStore;
   private SearchMetadataStore searchMetadataStore;
   private CacheSlotMetadataStore cacheSlotMetadataStore;
+  private FieldRedactionMetadataStore fieldRedactionMetadataStore;
 
   // for flag "astra.ng.dynamicChunkSizes"
   private final String cacheNodeId;
@@ -101,6 +103,7 @@ protected void startUp() throws Exception {
     cacheNodeAssignmentStore =
         new CacheNodeAssignmentStore(curatorFramework, zkConfig, cacheNodeId);
     cacheNodeMetadataStore = new CacheNodeMetadataStore(curatorFramework, zkConfig);
+    fieldRedactionMetadataStore = new FieldRedactionMetadataStore(curatorFramework, zkConfig, true);
 
     if (Boolean.getBoolean(ASTRA_NG_DYNAMIC_CHUNK_SIZES_FLAG)) {
       cacheNodeAssignmentStore.addListener(cacheNodeAssignmentChangeListener);
@@ -122,7 +125,8 @@ protected void startUp() throws Exception {
                 cacheSlotMetadataStore,
                 replicaMetadataStore,
                 snapshotMetadataStore,
-                searchMetadataStore);
+                searchMetadataStore,
+                fieldRedactionMetadataStore);
 
         chunkMap.put(newChunk.getSlotId(), newChunk);
       }
@@ -155,6 +159,7 @@ protected void shutDown() throws Exception {
     searchMetadataStore.close();
     snapshotMetadataStore.close();
     replicaMetadataStore.close();
+    fieldRedactionMetadataStore.close();
 
     LOG.info("Closed caching chunk manager.");
   }
@@ -235,7 +240,8 @@ private void onAssignmentHandler(CacheNodeAssignment assignment) {
                     searchMetadataStore,
                     cacheNodeAssignmentStore,
                     assignment,
-                    snapshotsBySnapshotId.get(assignment.snapshotId));
+                    snapshotsBySnapshotId.get(assignment.snapshotId),
+                    fieldRedactionMetadataStore);
             executorService.submit(newChunk::downloadChunkData);
             chunkMap.put(assignment.assignmentId, newChunk);
           }

astra/src/main/java/com/slack/astra/chunkManager/IndexingChunkManager.java

@@ -22,6 +22,7 @@
 import com.slack.astra.logstore.LogMessage;
 import com.slack.astra.logstore.LogStore;
 import com.slack.astra.logstore.LuceneIndexStoreImpl;
+import com.slack.astra.metadata.fieldredaction.FieldRedactionMetadataStore;
 import com.slack.astra.metadata.search.SearchMetadataStore;
 import com.slack.astra.metadata.snapshot.SnapshotMetadataStore;
 import com.slack.astra.proto.config.AstraConfigs;
@@ -93,6 +94,7 @@ public class IndexingChunkManager<T> extends ChunkManagerBase<T> {
   private SnapshotMetadataStore snapshotMetadataStore;
 
   private SearchMetadataStore searchMetadataStore;
+  private FieldRedactionMetadataStore fieldRedactionMetadataStore;
 
   /**
    * For capacity planning, we want to control how many roll overs are in progress at the same time.
@@ -261,7 +263,10 @@ private ReadWriteChunk<T> getOrCreateActiveChunk(
       @SuppressWarnings("unchecked")
       LogStore logStore =
           LuceneIndexStoreImpl.makeLogStore(
-              dataDirectory, indexerConfig.getLuceneConfig(), meterRegistry);
+              dataDirectory,
+              indexerConfig.getLuceneConfig(),
+              meterRegistry,
+              fieldRedactionMetadataStore);
 
       chunkRollOverStrategy.setActiveChunkDirectory(logStore.getDirectory());
 
@@ -389,6 +394,7 @@ protected void startUp() throws Exception {
 
     searchMetadataStore = new SearchMetadataStore(curatorFramework, zkConfig, false);
     snapshotMetadataStore = new SnapshotMetadataStore(curatorFramework, zkConfig);
+    fieldRedactionMetadataStore = new FieldRedactionMetadataStore(curatorFramework, zkConfig, true);
 
     stopIngestion = false;
   }
@@ -440,6 +446,7 @@ protected void shutDown() throws IOException {
 
     searchMetadataStore.close();
     snapshotMetadataStore.close();
+    fieldRedactionMetadataStore.close();
     LOG.info("Closed indexing chunk manager.");
   }
 

astra/src/main/java/com/slack/astra/chunkManager/RecoveryChunkManager.java

@@ -16,6 +16,7 @@
 import com.slack.astra.chunk.SearchContext;
 import com.slack.astra.chunkrollover.NeverRolloverChunkStrategy;
 import com.slack.astra.logstore.LogMessage;
+import com.slack.astra.metadata.fieldredaction.FieldRedactionMetadataStore;
 import com.slack.astra.metadata.search.SearchMetadataStore;
 import com.slack.astra.metadata.snapshot.SnapshotMetadataStore;
 import com.slack.astra.proto.config.AstraConfigs;
@@ -217,6 +218,7 @@ public static RecoveryChunkManager<LogMessage> fromConfig(
       MeterRegistry meterRegistry,
       SearchMetadataStore searchMetadataStore,
       SnapshotMetadataStore snapshotMetadataStore,
+      FieldRedactionMetadataStore fieldRedactionMetadataStore,
       AstraConfigs.IndexerConfig indexerConfig,
       BlobStore blobStore)
       throws Exception {
@@ -230,6 +232,7 @@ public static RecoveryChunkManager<LogMessage> fromConfig(
             meterRegistry,
             searchMetadataStore,
             snapshotMetadataStore,
+            fieldRedactionMetadataStore,
             searchContext);
 
     ChunkRolloverFactory chunkRolloverFactory =

astra/src/main/java/com/slack/astra/logstore/LuceneIndexStoreImpl.java

@@ -1,6 +1,8 @@
 package com.slack.astra.logstore;
 
 import com.slack.astra.logstore.schema.SchemaAwareLogDocumentBuilderImpl;
+import com.slack.astra.logstore.search.fieldRedaction.RedactionFilterDirectoryReader;
+import com.slack.astra.metadata.fieldredaction.FieldRedactionMetadataStore;
 import com.slack.astra.metadata.schema.LuceneFieldDef;
 import com.slack.astra.proto.config.AstraConfigs;
 import com.slack.astra.util.RuntimeHalterImpl;
@@ -21,6 +23,7 @@
 import org.apache.commons.io.FileUtils;
 import org.apache.lucene.analysis.Analyzer;
 import org.apache.lucene.analysis.standard.StandardAnalyzer;
+import org.apache.lucene.index.DirectoryReader;
 import org.apache.lucene.index.IndexCommit;
 import org.apache.lucene.index.IndexWriter;
 import org.apache.lucene.index.IndexWriterConfig;
@@ -81,15 +84,19 @@ public class LuceneIndexStoreImpl implements LogStore {
 
   // TODO: Set the policy via a lucene config file.
   public static LuceneIndexStoreImpl makeLogStore(
-      File dataDirectory, AstraConfigs.LuceneConfig luceneConfig, MeterRegistry metricsRegistry)
+      File dataDirectory,
+      AstraConfigs.LuceneConfig luceneConfig,
+      MeterRegistry metricsRegistry,
+      FieldRedactionMetadataStore fieldRedactionMetadataStore)
       throws IOException {
     return makeLogStore(
         dataDirectory,
         LuceneIndexStoreConfig.getCommitDuration(luceneConfig.getCommitDurationSecs()),
         LuceneIndexStoreConfig.getRefreshDuration(luceneConfig.getRefreshDurationSecs()),
         luceneConfig.getEnableFullTextSearch(),
         SchemaAwareLogDocumentBuilderImpl.FieldConflictPolicy.CONVERT_VALUE_AND_DUPLICATE_FIELD,
-        metricsRegistry);
+        metricsRegistry,
+        fieldRedactionMetadataStore);
   }
 
   public static LuceneIndexStoreImpl makeLogStore(
@@ -98,7 +105,8 @@ public static LuceneIndexStoreImpl makeLogStore(
       Duration refreshInterval,
       boolean enableFullTextSearch,
       SchemaAwareLogDocumentBuilderImpl.FieldConflictPolicy fieldConflictPolicy,
-      MeterRegistry metricsRegistry)
+      MeterRegistry metricsRegistry,
+      FieldRedactionMetadataStore fieldRedactionMetadataStore)
       throws IOException {
     // TODO: Move all these config values into chunk?
     // TODO: Chunk should create log store?
@@ -110,11 +118,15 @@ public static LuceneIndexStoreImpl makeLogStore(
         indexStoreCfg,
         SchemaAwareLogDocumentBuilderImpl.build(
             fieldConflictPolicy, enableFullTextSearch, metricsRegistry),
-        metricsRegistry);
+        metricsRegistry,
+        fieldRedactionMetadataStore);
   }
 
   public LuceneIndexStoreImpl(
-      LuceneIndexStoreConfig config, DocumentBuilder documentBuilder, MeterRegistry registry)
+      LuceneIndexStoreConfig config,
+      DocumentBuilder documentBuilder,
+      MeterRegistry registry,
+      FieldRedactionMetadataStore fieldRedactionMetadataStore)
       throws IOException {
 
     this.documentBuilder = documentBuilder;
@@ -126,7 +138,11 @@ public LuceneIndexStoreImpl(
         buildIndexWriterConfig(analyzer, this.snapshotDeletionPolicy, config, registry);
     indexDirectory = new MMapDirectory(config.indexFolder(id).toPath());
     indexWriter = Optional.of(new IndexWriter(indexDirectory, indexWriterConfig));
-    this.searcherManager = new SearcherManager(indexWriter.get(), false, false, null);
+
+    RedactionFilterDirectoryReader reader =
+        new RedactionFilterDirectoryReader(
+            DirectoryReader.open(indexWriter.get(), false, false), fieldRedactionMetadataStore);
+    this.searcherManager = new SearcherManager(reader, null);
 
     scheduledCommit.scheduleWithFixedDelay(
         () -> {

astra/src/main/java/com/slack/astra/logstore/search/LogIndexSearcherImpl.java

@@ -13,6 +13,8 @@
 import com.slack.astra.logstore.LogMessage.SystemField;
 import com.slack.astra.logstore.LogWireMessage;
 import com.slack.astra.logstore.opensearch.OpenSearchAdapter;
+import com.slack.astra.logstore.search.fieldRedaction.RedactionFilterDirectoryReader;
+import com.slack.astra.metadata.fieldredaction.FieldRedactionMetadataStore;
 import com.slack.astra.metadata.schema.LuceneFieldDef;
 import com.slack.astra.util.JsonUtil;
 import java.io.IOException;
@@ -24,6 +26,7 @@
 import java.util.concurrent.ConcurrentHashMap;
 import java.util.concurrent.TimeUnit;
 import java.util.stream.Collectors;
+import org.apache.lucene.index.DirectoryReader;
 import org.apache.lucene.search.CollectorManager;
 import org.apache.lucene.search.IndexSearcher;
 import org.apache.lucene.search.MultiCollectorManager;
@@ -58,16 +61,26 @@ public class LogIndexSearcherImpl implements LogIndexSearcher<LogMessage> {
   private final ReferenceManager.RefreshListener refreshListener;
 
   @VisibleForTesting
-  public static SearcherManager searcherManagerFromChunkId(String chunkId, BlobStore blobStore)
+  public static SearcherManager searcherManagerFromChunkId(
+      String chunkId, BlobStore blobStore, FieldRedactionMetadataStore fieldRedactionMetadataStore)
       throws IOException {
     Directory directory = new S3RemoteDirectory(chunkId, blobStore);
-    return new SearcherManager(directory, null);
+    DirectoryReader directoryReader = DirectoryReader.open(directory);
+
+    RedactionFilterDirectoryReader reader =
+        new RedactionFilterDirectoryReader(directoryReader, fieldRedactionMetadataStore);
+    return new SearcherManager(reader, null);
   }
 
   @VisibleForTesting
-  public static SearcherManager searcherManagerFromPath(Path path) throws IOException {
+  public static SearcherManager searcherManagerFromPath(
+      Path path, FieldRedactionMetadataStore fieldRedactionMetadataStore) throws IOException {
     MMapDirectory directory = new MMapDirectory(path);
-    return new SearcherManager(directory, null);
+    DirectoryReader directoryReader = DirectoryReader.open(directory);
+
+    RedactionFilterDirectoryReader reader =
+        new RedactionFilterDirectoryReader(directoryReader, fieldRedactionMetadataStore);
+    return new SearcherManager(reader, null);
   }
 
   public LogIndexSearcherImpl(
@@ -87,7 +100,6 @@ public void afterRefresh(boolean didRefresh) {
         };
     this.searcherManager = searcherManager;
     this.searcherManager.addListener(refreshListener);
-
     // initialize the adapter with whatever the default schema is
 
     openSearchAdapter.loadSchema();