Abstract Batch downloader from glean db

Summary: Add `BatchLocationParser` to db config to add the cas dependency at runtime.

Reviewed By: pepeiborra

Differential Revision: D70560447

fbshipit-source-id: 4f4e551efe157277e320578cfc3c8af2bab9135e

Author: iamirzhan

glean.cabal.in

@@ -470,6 +470,7 @@ library db
         Glean.Database.Backup.Backend
         Glean.Database.Backup.Locator
         Glean.Database.Backup.Mock
+        Glean.Database.BatchLocation
         Glean.Database.Catalog.Filter
         Glean.Database.Catalog
         Glean.Database.Catalog.Local.Files

glean/db/Glean/Database/BatchLocation.hs

@@ -0,0 +1,39 @@
+{-
+  Copyright (c) Meta Platforms, Inc. and affiliates.
+  All rights reserved.
+
+  This source code is licensed under the BSD-style license found in the
+  LICENSE file in the root directory of this source tree.
+-}
+
+module Glean.Database.BatchLocation
+  ( Parser(..)
+  , Location(..)
+  , DefaultParser(..)
+  ) where
+
+import Glean.Types as Thrift
+import Data.Text (Text)
+import Glean.Util.Some
+import Control.Exception
+
+-- | A batch's location parser
+class Parser a where
+  -- | Parse a string returning a location to get the batch from
+  fromString :: a -> Text -> Some Location
+
+instance Parser (Some Parser) where
+  fromString (Some parser) = fromString parser
+
+-- | A batch's location to download a batch from
+class Location a where
+  downloadBatch :: a -> Thrift.BatchFormat -> IO Thrift.Batch
+
+instance Location (Some Location) where
+  downloadBatch (Some location) format = downloadBatch location format
+
+data DefaultParser = DefaultParser
+
+instance Parser DefaultParser where
+  fromString _ _ = throw $
+    Thrift.Exception "Batch location's parser is not implemented"

glean/db/Glean/Database/Config.hs

@@ -63,6 +63,7 @@ import Util.Log (logInfo)
 import Glean.Angle.Types
 import qualified Glean.Database.Backup.Backend as Backup -- from glean/util
 import qualified Glean.Database.Backup.Mock as Backup.Mock
+import qualified Glean.Database.BatchLocation as BatchLocation
 import Glean.Database.Catalog (Catalog)
 import qualified Glean.Database.Catalog.Local.Files as Catalog
 import qualified Glean.Database.Catalog.Local.Memory as Catalog
@@ -155,6 +156,8 @@ data Config = Config
     -- ^ Logger for recording stats of databases produced
   , cfgBackupBackends :: HashMap Text (Some Backup.Backend)
     -- ^ Backup backends
+  , cfgBatchLocationParser :: Some BatchLocation.Parser
+    -- ^ Batch's location parser
   , cfgEnableRecursion :: Bool
     -- ^ Enable experimental support for recursion
   , cfgFilterAvailableDBs :: [Repo] -> IO [Repo]
@@ -202,6 +205,7 @@ instance Default Config where
     , cfgServerLogger = Some NullGleanServerLogger
     , cfgDatabaseLogger = Some NullGleanDatabaseLogger
     , cfgBackupBackends = HashMap.fromList [("mock", Backup.Mock.mock)]
+    , cfgBatchLocationParser = Some BatchLocation.DefaultParser
     , cfgEnableRecursion = False
     , cfgFilterAvailableDBs = const $ return []
     , cfgTracer = mempty
@@ -454,6 +458,7 @@ options = do
     , cfgServerLogger = cfgServerLogger def
     , cfgDatabaseLogger = cfgDatabaseLogger def
     , cfgBackupBackends = cfgBackupBackends def
+    , cfgBatchLocationParser = cfgBatchLocationParser def
     , cfgFilterAvailableDBs = const $ return []
     , cfgTracer = mempty
     , cfgSchemaId = Nothing

glean/db/Glean/Database/Env.hs

@@ -137,6 +137,7 @@ initEnv evb envStorage envCatalog shardManager cfg
       { envEventBase = evb
       , envServerLogger = cfgServerLogger cfg
       , envDatabaseLogger = cfgDatabaseLogger cfg
+      , envBatchLocationParser = cfgBatchLocationParser cfg
       , envReadOnly = cfgReadOnly cfg
       , envMockWrites = cfgMockWrites cfg
       , envListener = cfgListener cfg

glean/db/Glean/Database/Types.hs

@@ -37,6 +37,7 @@ import Util.STM
 
 import Glean.Angle.Types hiding (describe)
 import qualified Glean.Database.Backup.Backend as Backup
+import qualified Glean.Database.BatchLocation as BatchLocation
 import Glean.Database.Catalog (Catalog)
 import Glean.Database.Config
 import Glean.Database.Meta
@@ -224,6 +225,7 @@ data Env = forall storage. Storage storage => Env
   { envEventBase :: EventBaseDataplane
   , envServerLogger :: Some GleanServerLogger
   , envDatabaseLogger :: Some GleanDatabaseLogger
+  , envBatchLocationParser :: Some BatchLocation.Parser
   , envLoggerRateLimit :: RateLimiterMap Text
   , envCatalog :: Catalog
   , envStorage :: storage

glean/db/Glean/Database/Writes.hs

@@ -71,6 +71,7 @@ import Util.Defer
 import Util.Log
 import Util.STM
 
+import Glean.Database.BatchLocation as BatchLocation
 import Glean.Database.Exception
 import Glean.Database.Open
 import Glean.Database.Schema.Types
@@ -377,11 +378,11 @@ enqueueBatchDescriptor env repo enqueueBatch waitPolicy = do
   traceMsg (envTracer env)
     (GleanTraceEnqueue repo EnqueueBatchDescriptor 0) $ do
   handle <- UUID.toText <$> UUID.nextRandom
-  _ <- case enqueueBatch of
+  descriptor <- case enqueueBatch of
     Thrift.EnqueueBatch_descriptor descriptor -> return descriptor
     Thrift.EnqueueBatch_EMPTY -> throwIO $ Thrift.Exception "empty batch"
   write <- enqueueWrite env repo 0 Nothing False $
-    throwIO $ Exception "not implemented"
+    writeContentFromBatch <$> downloadBatchFromLocation env descriptor
   when (waitPolicy == Thrift.EnqueueBatchWaitPolicy_Remember)
     $ rememberWrite env handle write
   return $ def { enqueueBatchResponse_handle = handle }
@@ -457,6 +458,19 @@ deleteWriteQueues env OpenDB{odbWriting = Just Writing{..}} = do
   -- get removed by the next write thread to encounter it.
 deleteWriteQueues _ _ = return ()
 
+downloadBatchFromLocation
+  :: Env
+  -> Thrift.BatchDescriptor
+  -> IO Thrift.Batch
+downloadBatchFromLocation Env{..} batchDescriptor =
+  let
+    batchFormat = batchDescriptor_format batchDescriptor
+    locationString = batchDescriptor_location batchDescriptor
+    location = BatchLocation.fromString envBatchLocationParser locationString
+  in
+    BatchLocation.downloadBatch location batchFormat
+
+
 k :: Int
 k = 1024