Add enqueueBatchHash to Backend class

Summary: * Add `enqueueBatchHash` function to Backend class * Call new thrift method from Remote.hs backend Reviewed By: simonmar Differential Revision: D69524200 fbshipit-source-id: 5665df8b0daf2e115d32f10f45047eb2c8ffa037
facebookincubator · Feb 14, 2025 · f2107c9 · f2107c9
1 parent 00dc2d1
commit f2107c9
Show file tree

Hide file tree

Showing 8 changed files with 37 additions and 1 deletion.
diff --git a/glean/client/hs/Glean/Remote.hs b/glean/client/hs/Glean/Remote.hs
@@ -270,6 +270,8 @@ instance Backend ThriftBackend where
       GleanService.sendBatch cbatch
   enqueueJsonBatch t repo batch =
     withShard t repo $ GleanService.sendJsonBatch repo batch
+  enqueueBatchDescriptor t repo batch waitPolicy =
+    withShard t repo $ GleanService.enqueueBatch repo batch waitPolicy
   pollBatch t handle = withoutShard t $ GleanService.finishBatch handle
 
   displayBackend = show

diff --git a/glean/client/hs/local/Glean/LocalOrRemote.hs b/glean/client/hs/local/Glean/LocalOrRemote.hs
@@ -217,6 +217,7 @@ instance Backend (Some LocalOrRemote) where
 
   enqueueBatch (Some backend) = enqueueBatch backend
   enqueueJsonBatch (Some backend) = enqueueJsonBatch backend
+  enqueueBatchDescriptor (Some backend) = enqueueBatchDescriptor backend
   pollBatch (Some backend) = pollBatch backend
   displayBackend (Some backend) = displayBackend backend
   hasDatabase (Some backend) = hasDatabase backend

diff --git a/glean/db/Glean/Backend/Local.hs b/glean/db/Glean/Backend/Local.hs
@@ -175,6 +175,8 @@ instance Backend Database.Env where
 
   enqueueBatch env cbatch = Database.enqueueBatch env cbatch Nothing
   enqueueJsonBatch env cbatch = Database.enqueueJsonBatch env cbatch
+  enqueueBatchDescriptor env repo batch waitPolicy =
+    Database.enqueueBatchDescriptor env repo batch waitPolicy
   pollBatch env handle = Database.pollBatch env handle
 
   displayBackend _ = "(local backend)"

diff --git a/glean/db/Glean/Backend/Logging.hs b/glean/db/Glean/Backend/Logging.hs
@@ -149,6 +149,11 @@ instance Backend LoggingBackend where
   enqueueJsonBatch (LoggingBackend env) repo batch =
     loggingAction (runLogRepo "enqueueJsonBatch" env repo) (const mempty) $
       enqueueJsonBatch env repo batch
+  enqueueBatchDescriptor (LoggingBackend env) repo batch waitPolicy =
+    loggingAction
+      (runLogRepo "enqueueBatchDescriptor" env repo)
+      (const mempty) $
+        enqueueBatchDescriptor env repo batch waitPolicy
   pollBatch (LoggingBackend env) handle =
     loggingAction (runLogCmd "pollBatch" env) (const mempty) $
       pollBatch env handle

diff --git a/glean/db/Glean/Database/Writes.hs b/glean/db/Glean/Database/Writes.hs
@@ -32,6 +32,7 @@
 module Glean.Database.Writes
   ( enqueueBatch
   , enqueueJsonBatch
+  , enqueueBatchDescriptor
   , enqueueCheckpoint
   , pollBatch
   , reapWrites
@@ -363,6 +364,15 @@ enqueueJsonBatch env repo batch = do
   when (sendJsonBatch_remember batch) $ rememberWrite env handle write
   return $ def { sendJsonBatchResponse_handle = handle }
 
+enqueueBatchDescriptor
+  :: Env
+  -> Repo
+  -> EnqueueBatch
+  -> EnqueueBatchWaitPolicy
+  -> IO Thrift.EnqueueBatchResponse
+enqueueBatchDescriptor _ _ _ _ = do
+  error "not implemented"
+
 pollBatch :: Env -> Handle -> IO FinishResponse
 pollBatch env@Env{..} handle = do
   r <- HashMap.lookup handle <$> readTVarIO envWrites

diff --git a/glean/hs/Glean/Backend/Retry.hs b/glean/hs/Glean/Backend/Retry.hs
@@ -61,6 +61,10 @@ instance Backend RetryWritesBackend where
     retryChannelExceptions policy $ enqueueBatch backend batch
   enqueueJsonBatch (RetryWritesBackend policy backend) repo batch =
     retryChannelExceptions policy $ enqueueJsonBatch backend repo batch
+  enqueueBatchDescriptor (RetryWritesBackend policy backend)
+    repo batch waitPolicy =
+      retryChannelExceptions policy
+        $ enqueueBatchDescriptor backend repo batch waitPolicy
   pollBatch (RetryWritesBackend policy backend) handle =
     retryChannelExceptions policy $ pollBatch backend handle
 
@@ -115,6 +119,8 @@ instance Backend RetryReadsBackend where
 
   enqueueBatch (RetryReadsBackend _ backend) = enqueueBatch backend
   enqueueJsonBatch (RetryReadsBackend _ backend) = enqueueJsonBatch backend
+  enqueueBatchDescriptor (RetryReadsBackend _ backend) =
+    enqueueBatchDescriptor backend
   pollBatch (RetryReadsBackend _ backend) = pollBatch backend
 
   displayBackend (RetryReadsBackend _ backend) = displayBackend backend

diff --git a/glean/hs/Glean/Backend/Types.hs b/glean/hs/Glean/Backend/Types.hs
@@ -147,6 +147,14 @@ class Backend a where
     -> Thrift.SendJsonBatch
     -> IO Thrift.SendJsonBatchResponse
 
+  -- Enqueue a batch descriptor to be downloaded and written
+  enqueueBatchDescriptor
+    :: a
+    -> Thrift.Repo
+    -> Thrift.EnqueueBatch
+    -> Thrift.EnqueueBatchWaitPolicy
+    -> IO Thrift.EnqueueBatchResponse
+
   -- Poll the status of a write batch
   pollBatch :: a -> Thrift.Handle -> IO Thrift.FinishResponse
 
@@ -202,6 +210,7 @@ instance Backend (Some Backend) where
 
   enqueueBatch (Some backend) = enqueueBatch backend
   enqueueJsonBatch (Some backend) = enqueueJsonBatch backend
+  enqueueBatchDescriptor (Some backend) = enqueueBatchDescriptor backend
   pollBatch (Some backend) = pollBatch backend
   displayBackend (Some backend) = displayBackend backend
   hasDatabase (Some backend) = hasDatabase backend

diff --git a/glean/server/Glean/Handler.hs b/glean/server/Glean/Handler.hs
@@ -41,7 +41,8 @@ handler State{..} req =
 
     Service.SendBatch cbatch -> Backend.enqueueBatch backend cbatch
 
-    Service.EnqueueBatch {} -> error "not implemented"
+    Service.EnqueueBatch repo batch waitPolicy
+      -> Backend.enqueueBatchDescriptor backend repo batch waitPolicy
 
     Service.FinishBatch handle -> Backend.pollBatch backend handle