chore: adding write delete stale logic (#98)

fixes: #94

Author: mnorbury

lib/src/main/java/io/cloudquery/internal/servers/plugin/v3/PluginServer.java

@@ -2,6 +2,7 @@
 
 import com.google.protobuf.ByteString;
 import io.cloudquery.helper.ArrowHelper;
+import io.cloudquery.messages.WriteDeleteStale;
 import io.cloudquery.messages.WriteInsert;
 import io.cloudquery.messages.WriteMessage;
 import io.cloudquery.messages.WriteMigrateTable;
@@ -16,6 +17,7 @@
 import io.grpc.stub.StreamObserver;
 import java.io.IOException;
 import java.util.ArrayList;
+import java.util.Date;
 import java.util.List;
 
 public class PluginServer extends PluginImplBase {
@@ -125,6 +127,10 @@ public void onNext(Write.Request request) {
             plugin.write(processMigrateTableRequest(request));
           } else if (messageCase == Write.Request.MessageCase.INSERT) {
             plugin.write(processInsertRequest(request));
+          } else if (messageCase == Write.Request.MessageCase.DELETE) {
+            plugin.write(processDeleteStaleRequest(request));
+          } else {
+            throw new IllegalArgumentException("Unknown message type: " + messageCase);
           }
         } catch (IOException | ValidationException ex) {
           onError(ex);
@@ -166,4 +172,13 @@ private WriteMessage processInsertRequest(Write.Request request)
     ByteString record = insert.getRecord();
     return new WriteInsert(ArrowHelper.decodeResource(record));
   }
+
+  private WriteMessage processDeleteStaleRequest(Write.Request request)
+      throws IOException, ValidationException {
+    Write.MessageDeleteStale messageDeleteStale = request.getDelete();
+    return new WriteDeleteStale(
+        messageDeleteStale.getTableName(),
+        messageDeleteStale.getSourceName(),
+        new Date(messageDeleteStale.getSyncTime().getSeconds() * 1000));
+  }
 }

lib/src/main/java/io/cloudquery/memdb/MemDBClient.java

@@ -1,24 +1,30 @@
 package io.cloudquery.memdb;
 
+import io.cloudquery.messages.WriteDeleteStale;
 import io.cloudquery.messages.WriteInsert;
 import io.cloudquery.messages.WriteMessage;
 import io.cloudquery.messages.WriteMigrateTable;
+import io.cloudquery.scalar.Timestamp;
 import io.cloudquery.schema.ClientMeta;
+import io.cloudquery.schema.Column;
 import io.cloudquery.schema.Resource;
 import io.cloudquery.schema.Table;
 import io.cloudquery.schema.TableColumnChange;
 import java.util.ArrayList;
+import java.util.Date;
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
+import java.util.Objects;
+import java.util.Optional;
 import java.util.concurrent.locks.ReentrantReadWriteLock;
 
 public class MemDBClient implements ClientMeta {
   private static final String id = "memdb";
 
-  private ReentrantReadWriteLock lock = new ReentrantReadWriteLock();
-  private Map<String, Table> tables = new HashMap<>();
-  private Map<String, List<Resource>> memDB = new HashMap<>();
+  private final ReentrantReadWriteLock lock = new ReentrantReadWriteLock();
+  private final Map<String, Table> tables = new HashMap<>();
+  private final Map<String, List<Resource>> memDB = new HashMap<>();
 
   public MemDBClient() {}
 
@@ -37,6 +43,9 @@ public void write(WriteMessage message) {
       if (message instanceof WriteInsert insert) {
         insert(insert);
       }
+      if (message instanceof WriteDeleteStale deleteStale) {
+        deleteStale(deleteStale);
+      }
     } finally {
       lock.writeLock().unlock();
     }
@@ -74,6 +83,40 @@ private void overwrite(Table table, Resource resource) {
     memDB.get(tableName).add(resource);
   }
 
+  private void deleteStale(WriteDeleteStale deleteStale) {
+    String tableName = deleteStale.getTableName();
+
+    List<Resource> filteredList = new ArrayList<>();
+
+    for (int i = 0; i < memDB.get(tableName).size(); i++) {
+      Resource row = memDB.get(tableName).get(i);
+      Optional<Column> sourceColumn = row.getTable().getColumn(Column.CQ_SOURCE_NAME);
+      if (sourceColumn.isEmpty()) {
+        continue;
+      }
+      Optional<Column> syncColumn = row.getTable().getColumn(Column.CQ_SYNC_TIME);
+      if (syncColumn.isEmpty()) {
+        continue;
+      }
+
+      String sourceName = "";
+      if (row.get(Column.CQ_SOURCE_NAME) != null) {
+        sourceName = row.get(Column.CQ_SOURCE_NAME).toString();
+      }
+
+      if (Objects.equals(sourceName, deleteStale.getSourceName())) {
+        Date rowSyncTime = new Date(0);
+        if (row.get(Column.CQ_SYNC_TIME) != null) {
+          rowSyncTime = new Date(((Timestamp) row.get(Column.CQ_SYNC_TIME)).get());
+        }
+        if (!rowSyncTime.before(deleteStale.getTimestamp())) {
+          filteredList.add(row);
+        }
+      }
+    }
+    memDB.put(tableName, filteredList);
+  }
+
   public void close() {
     // do nothing
   }

lib/src/main/java/io/cloudquery/messages/WriteDeleteStale.java

@@ -0,0 +1,13 @@
+package io.cloudquery.messages;
+
+import java.util.Date;
+import lombok.AllArgsConstructor;
+import lombok.Getter;
+
+@AllArgsConstructor
+@Getter
+public class WriteDeleteStale extends WriteMessage {
+  private String tableName;
+  private String sourceName;
+  private Date timestamp;
+}

lib/src/main/java/io/cloudquery/schema/Column.java

@@ -8,6 +8,8 @@
 @Builder(toBuilder = true)
 @Getter
 public class Column {
+  public static final String CQ_SOURCE_NAME = "_cq_source_name";
+  public static final String CQ_SYNC_TIME = "_cq_sync_time";
   private String name;
   private String description;
   private ArrowType type;

lib/src/test/java/io/cloudquery/internal/servers/plugin/v3/PluginServerTest.java

@@ -5,6 +5,7 @@
 
 import com.google.protobuf.ByteString;
 import io.cloudquery.helper.ArrowHelper;
+import io.cloudquery.messages.WriteDeleteStale;
 import io.cloudquery.messages.WriteInsert;
 import io.cloudquery.messages.WriteMigrateTable;
 import io.cloudquery.plugin.Plugin;
@@ -78,6 +79,18 @@ public void shouldSendWriteInsertMessage() throws Exception {
     verify(plugin).write(any(WriteInsert.class));
   }
 
+  @Test
+  public void shouldSendWriteDeleteStaleMessage() throws Exception {
+    NullResponseStream<Write.Response> responseObserver = new NullResponseStream<>();
+
+    StreamObserver<Write.Request> writeService = pluginStub.write(responseObserver);
+    writeService.onNext(generateDeleteStaleMessage());
+    writeService.onCompleted();
+    responseObserver.await();
+
+    verify(plugin).write(any(WriteDeleteStale.class));
+  }
+
   private static Write.Request generateMigrateTableMessage() throws IOException {
     Table table = Table.builder().name("test").build();
     return Write.Request.newBuilder()
@@ -96,6 +109,11 @@ private Write.Request generateInsertMessage() throws IOException, ValidationExce
     return Write.Request.newBuilder().setInsert(messageInsert).build();
   }
 
+  private Write.Request generateDeleteStaleMessage() {
+    Write.MessageDeleteStale messageDeleteStale = Write.MessageDeleteStale.newBuilder().build();
+    return Write.Request.newBuilder().setDelete(messageDeleteStale).build();
+  }
+
   private static class NullResponseStream<T> implements StreamObserver<T> {
     private final CountDownLatch countDownLatch = new CountDownLatch(1);