{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE DisambiguateRecordFields #-}

-- | Integration with the Casa server.

module Pantry.Casa where

import qualified Casa.Client as Casa
import qualified Casa.Types as Casa
import           Conduit
import qualified Data.HashMap.Strict as HM
import qualified Pantry.SHA256 as SHA256
import           Pantry.Storage
import           Pantry.Types as P
import           RIO
import qualified RIO.ByteString as B

-- | Lookup a tree.
casaLookupTree ::
     (HasPantryConfig env, HasLogFunc env)
  => TreeKey
  -> RIO env (Maybe (TreeKey, P.Tree))
casaLookupTree :: TreeKey -> RIO env (Maybe (TreeKey, Tree))
casaLookupTree (P.TreeKey BlobKey
key) =
  (SomeException -> RIO env (Maybe (TreeKey, Tree)))
-> RIO env (Maybe (TreeKey, Tree))
-> RIO env (Maybe (TreeKey, Tree))
forall (m :: * -> *) a.
MonadUnliftIO m =>
(SomeException -> m a) -> m a -> m a
handleAny (RIO env (Maybe (TreeKey, Tree))
-> SomeException -> RIO env (Maybe (TreeKey, Tree))
forall a b. a -> b -> a
const (Maybe (TreeKey, Tree) -> RIO env (Maybe (TreeKey, Tree))
forall (f :: * -> *) a. Applicative f => a -> f a
pure Maybe (TreeKey, Tree)
forall a. Maybe a
Nothing))
    (ReaderT SqlBackend (RIO env) (Maybe (TreeKey, Tree))
-> RIO env (Maybe (TreeKey, Tree))
forall env a.
(HasPantryConfig env, HasLogFunc env) =>
ReaderT SqlBackend (RIO env) a -> RIO env a
withStorage
      (ConduitT
  ()
  Void
  (ResourceT (ReaderT SqlBackend (RIO env)))
  (Maybe (TreeKey, Tree))
-> ReaderT SqlBackend (RIO env) (Maybe (TreeKey, Tree))
forall (m :: * -> *) r.
MonadUnliftIO m =>
ConduitT () Void (ResourceT m) r -> m r
runConduitRes (Identity BlobKey
-> ConduitT
     ()
     (BlobKey, ByteString)
     (ResourceT (ReaderT SqlBackend (RIO env)))
     ()
forall (f :: * -> *) env i.
(Foldable f, HasPantryConfig env, HasLogFunc env) =>
f BlobKey
-> ConduitT
     i
     (BlobKey, ByteString)
     (ResourceT (ReaderT SqlBackend (RIO env)))
     ()
casaBlobSource (BlobKey -> Identity BlobKey
forall a. a -> Identity a
Identity BlobKey
key) ConduitT
  ()
  (BlobKey, ByteString)
  (ResourceT (ReaderT SqlBackend (RIO env)))
  ()
-> ConduitM
     (BlobKey, ByteString)
     Void
     (ResourceT (ReaderT SqlBackend (RIO env)))
     (Maybe (TreeKey, Tree))
-> ConduitT
     ()
     Void
     (ResourceT (ReaderT SqlBackend (RIO env)))
     (Maybe (TreeKey, Tree))
forall (m :: * -> *) a b c r.
Monad m =>
ConduitM a b m () -> ConduitM b c m r -> ConduitM a c m r
.| ((BlobKey, ByteString)
 -> ResourceT (ReaderT SqlBackend (RIO env)) (TreeKey, Tree))
-> ConduitT
     (BlobKey, ByteString)
     (TreeKey, Tree)
     (ResourceT (ReaderT SqlBackend (RIO env)))
     ()
forall (m :: * -> *) a b.
Monad m =>
(a -> m b) -> ConduitT a b m ()
mapMC (BlobKey, ByteString)
-> ResourceT (ReaderT SqlBackend (RIO env)) (TreeKey, Tree)
forall (m :: * -> *).
MonadThrow m =>
(BlobKey, ByteString) -> m (TreeKey, Tree)
parseTreeM ConduitT
  (BlobKey, ByteString)
  (TreeKey, Tree)
  (ResourceT (ReaderT SqlBackend (RIO env)))
  ()
-> ConduitM
     (TreeKey, Tree)
     Void
     (ResourceT (ReaderT SqlBackend (RIO env)))
     (Maybe (TreeKey, Tree))
-> ConduitM
     (BlobKey, ByteString)
     Void
     (ResourceT (ReaderT SqlBackend (RIO env)))
     (Maybe (TreeKey, Tree))
forall (m :: * -> *) a b c r.
Monad m =>
ConduitM a b m () -> ConduitM b c m r -> ConduitM a c m r
.| ConduitM
  (TreeKey, Tree)
  Void
  (ResourceT (ReaderT SqlBackend (RIO env)))
  (Maybe (TreeKey, Tree))
forall (m :: * -> *) i. Monad m => Consumer i m (Maybe i)
await)))

-- | Lookup a single blob. If possible, prefer 'casaBlobSource', and
-- query a group of keys at once, rather than one at a time. This will
-- have better network performance.
casaLookupKey ::
     (HasPantryConfig env, HasLogFunc env)
  => BlobKey
  -> RIO env (Maybe ByteString)
casaLookupKey :: BlobKey -> RIO env (Maybe ByteString)
casaLookupKey BlobKey
key =
  (SomeException -> RIO env (Maybe ByteString))
-> RIO env (Maybe ByteString) -> RIO env (Maybe ByteString)
forall (m :: * -> *) a.
MonadUnliftIO m =>
(SomeException -> m a) -> m a -> m a
handleAny (RIO env (Maybe ByteString)
-> SomeException -> RIO env (Maybe ByteString)
forall a b. a -> b -> a
const (Maybe ByteString -> RIO env (Maybe ByteString)
forall (f :: * -> *) a. Applicative f => a -> f a
pure Maybe ByteString
forall a. Maybe a
Nothing))
  ((Maybe (BlobKey, ByteString) -> Maybe ByteString)
-> RIO env (Maybe (BlobKey, ByteString))
-> RIO env (Maybe ByteString)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap
    (((BlobKey, ByteString) -> ByteString)
-> Maybe (BlobKey, ByteString) -> Maybe ByteString
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (BlobKey, ByteString) -> ByteString
forall a b. (a, b) -> b
snd)
    (ReaderT SqlBackend (RIO env) (Maybe (BlobKey, ByteString))
-> RIO env (Maybe (BlobKey, ByteString))
forall env a.
(HasPantryConfig env, HasLogFunc env) =>
ReaderT SqlBackend (RIO env) a -> RIO env a
withStorage (ConduitT
  ()
  Void
  (ResourceT (ReaderT SqlBackend (RIO env)))
  (Maybe (BlobKey, ByteString))
-> ReaderT SqlBackend (RIO env) (Maybe (BlobKey, ByteString))
forall (m :: * -> *) r.
MonadUnliftIO m =>
ConduitT () Void (ResourceT m) r -> m r
runConduitRes (Identity BlobKey
-> ConduitT
     ()
     (BlobKey, ByteString)
     (ResourceT (ReaderT SqlBackend (RIO env)))
     ()
forall (f :: * -> *) env i.
(Foldable f, HasPantryConfig env, HasLogFunc env) =>
f BlobKey
-> ConduitT
     i
     (BlobKey, ByteString)
     (ResourceT (ReaderT SqlBackend (RIO env)))
     ()
casaBlobSource (BlobKey -> Identity BlobKey
forall a. a -> Identity a
Identity BlobKey
key) ConduitT
  ()
  (BlobKey, ByteString)
  (ResourceT (ReaderT SqlBackend (RIO env)))
  ()
-> ConduitM
     (BlobKey, ByteString)
     Void
     (ResourceT (ReaderT SqlBackend (RIO env)))
     (Maybe (BlobKey, ByteString))
-> ConduitT
     ()
     Void
     (ResourceT (ReaderT SqlBackend (RIO env)))
     (Maybe (BlobKey, ByteString))
forall (m :: * -> *) a b c r.
Monad m =>
ConduitM a b m () -> ConduitM b c m r -> ConduitM a c m r
.| ConduitM
  (BlobKey, ByteString)
  Void
  (ResourceT (ReaderT SqlBackend (RIO env)))
  (Maybe (BlobKey, ByteString))
forall (m :: * -> *) i. Monad m => Consumer i m (Maybe i)
await))))

-- | A source of blobs given a set of keys. All blobs are
-- automatically stored in the local pantry database.
casaBlobSource ::
     (Foldable f, HasPantryConfig env, HasLogFunc env)
  => f BlobKey
  -> ConduitT i (BlobKey, ByteString) (ResourceT (ReaderT SqlBackend (RIO env))) ()
casaBlobSource :: f BlobKey
-> ConduitT
     i
     (BlobKey, ByteString)
     (ResourceT (ReaderT SqlBackend (RIO env)))
     ()
casaBlobSource f BlobKey
keys = ConduitT
  i
  (BlobKey, ByteString)
  (ResourceT (ReaderT SqlBackend (RIO env)))
  ()
forall i.
ConduitT
  i
  (BlobKey, ByteString)
  (ResourceT (ReaderT SqlBackend (RIO env)))
  ()
source ConduitT
  i
  (BlobKey, ByteString)
  (ResourceT (ReaderT SqlBackend (RIO env)))
  ()
-> ConduitM
     (BlobKey, ByteString)
     (BlobKey, ByteString)
     (ResourceT (ReaderT SqlBackend (RIO env)))
     ()
-> ConduitT
     i
     (BlobKey, ByteString)
     (ResourceT (ReaderT SqlBackend (RIO env)))
     ()
forall (m :: * -> *) a b c r.
Monad m =>
ConduitM a b m () -> ConduitM b c m r -> ConduitM a c m r
.| ConduitM
  (BlobKey, ByteString)
  (BlobKey, ByteString)
  (ResourceT (ReaderT SqlBackend (RIO env)))
  ()
convert ConduitM
  (BlobKey, ByteString)
  (BlobKey, ByteString)
  (ResourceT (ReaderT SqlBackend (RIO env)))
  ()
-> ConduitM
     (BlobKey, ByteString)
     (BlobKey, ByteString)
     (ResourceT (ReaderT SqlBackend (RIO env)))
     ()
-> ConduitM
     (BlobKey, ByteString)
     (BlobKey, ByteString)
     (ResourceT (ReaderT SqlBackend (RIO env)))
     ()
forall (m :: * -> *) a b c r.
Monad m =>
ConduitM a b m () -> ConduitM b c m r -> ConduitM a c m r
.| ConduitM
  (BlobKey, ByteString)
  (BlobKey, ByteString)
  (ResourceT (ReaderT SqlBackend (RIO env)))
  ()
forall a.
ConduitT
  (a, ByteString)
  (a, ByteString)
  (ResourceT (ReaderT SqlBackend (RIO env)))
  ()
store
  where
    source :: ConduitT
  i
  (BlobKey, ByteString)
  (ResourceT (ReaderT SqlBackend (RIO env)))
  ()
source = do
      CasaRepoPrefix
pullUrl <- ResourceT (ReaderT SqlBackend (RIO env)) CasaRepoPrefix
-> ConduitT
     i
     (BlobKey, ByteString)
     (ResourceT (ReaderT SqlBackend (RIO env)))
     CasaRepoPrefix
forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift (ResourceT (ReaderT SqlBackend (RIO env)) CasaRepoPrefix
 -> ConduitT
      i
      (BlobKey, ByteString)
      (ResourceT (ReaderT SqlBackend (RIO env)))
      CasaRepoPrefix)
-> ResourceT (ReaderT SqlBackend (RIO env)) CasaRepoPrefix
-> ConduitT
     i
     (BlobKey, ByteString)
     (ResourceT (ReaderT SqlBackend (RIO env)))
     CasaRepoPrefix
forall a b. (a -> b) -> a -> b
$ ReaderT SqlBackend (RIO env) CasaRepoPrefix
-> ResourceT (ReaderT SqlBackend (RIO env)) CasaRepoPrefix
forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift (ReaderT SqlBackend (RIO env) CasaRepoPrefix
 -> ResourceT (ReaderT SqlBackend (RIO env)) CasaRepoPrefix)
-> ReaderT SqlBackend (RIO env) CasaRepoPrefix
-> ResourceT (ReaderT SqlBackend (RIO env)) CasaRepoPrefix
forall a b. (a -> b) -> a -> b
$ RIO env CasaRepoPrefix
-> ReaderT SqlBackend (RIO env) CasaRepoPrefix
forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift (RIO env CasaRepoPrefix
 -> ReaderT SqlBackend (RIO env) CasaRepoPrefix)
-> RIO env CasaRepoPrefix
-> ReaderT SqlBackend (RIO env) CasaRepoPrefix
forall a b. (a -> b) -> a -> b
$ Getting CasaRepoPrefix env CasaRepoPrefix -> RIO env CasaRepoPrefix
forall s (m :: * -> *) a. MonadReader s m => Getting a s a -> m a
view (Getting CasaRepoPrefix env CasaRepoPrefix
 -> RIO env CasaRepoPrefix)
-> Getting CasaRepoPrefix env CasaRepoPrefix
-> RIO env CasaRepoPrefix
forall a b. (a -> b) -> a -> b
$ (PantryConfig -> Const CasaRepoPrefix PantryConfig)
-> env -> Const CasaRepoPrefix env
forall env. HasPantryConfig env => Lens' env PantryConfig
pantryConfigL ((PantryConfig -> Const CasaRepoPrefix PantryConfig)
 -> env -> Const CasaRepoPrefix env)
-> ((CasaRepoPrefix -> Const CasaRepoPrefix CasaRepoPrefix)
    -> PantryConfig -> Const CasaRepoPrefix PantryConfig)
-> Getting CasaRepoPrefix env CasaRepoPrefix
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (PantryConfig -> CasaRepoPrefix)
-> SimpleGetter PantryConfig CasaRepoPrefix
forall s a. (s -> a) -> SimpleGetter s a
to PantryConfig -> CasaRepoPrefix
pcCasaRepoPrefix
      Int
maxPerRequest <- ResourceT (ReaderT SqlBackend (RIO env)) Int
-> ConduitT
     i
     (BlobKey, ByteString)
     (ResourceT (ReaderT SqlBackend (RIO env)))
     Int
forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift (ResourceT (ReaderT SqlBackend (RIO env)) Int
 -> ConduitT
      i
      (BlobKey, ByteString)
      (ResourceT (ReaderT SqlBackend (RIO env)))
      Int)
-> ResourceT (ReaderT SqlBackend (RIO env)) Int
-> ConduitT
     i
     (BlobKey, ByteString)
     (ResourceT (ReaderT SqlBackend (RIO env)))
     Int
forall a b. (a -> b) -> a -> b
$ ReaderT SqlBackend (RIO env) Int
-> ResourceT (ReaderT SqlBackend (RIO env)) Int
forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift (ReaderT SqlBackend (RIO env) Int
 -> ResourceT (ReaderT SqlBackend (RIO env)) Int)
-> ReaderT SqlBackend (RIO env) Int
-> ResourceT (ReaderT SqlBackend (RIO env)) Int
forall a b. (a -> b) -> a -> b
$ RIO env Int -> ReaderT SqlBackend (RIO env) Int
forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift (RIO env Int -> ReaderT SqlBackend (RIO env) Int)
-> RIO env Int -> ReaderT SqlBackend (RIO env) Int
forall a b. (a -> b) -> a -> b
$ Getting Int env Int -> RIO env Int
forall s (m :: * -> *) a. MonadReader s m => Getting a s a -> m a
view (Getting Int env Int -> RIO env Int)
-> Getting Int env Int -> RIO env Int
forall a b. (a -> b) -> a -> b
$ (PantryConfig -> Const Int PantryConfig) -> env -> Const Int env
forall env. HasPantryConfig env => Lens' env PantryConfig
pantryConfigL ((PantryConfig -> Const Int PantryConfig) -> env -> Const Int env)
-> ((Int -> Const Int Int)
    -> PantryConfig -> Const Int PantryConfig)
-> Getting Int env Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (PantryConfig -> Int) -> SimpleGetter PantryConfig Int
forall s a. (s -> a) -> SimpleGetter s a
to PantryConfig -> Int
pcCasaMaxPerRequest
      SourceConfig
-> ConduitT
     i
     (BlobKey, ByteString)
     (ResourceT (ReaderT SqlBackend (RIO env)))
     ()
forall (m :: * -> *) i.
(MonadThrow m, MonadResource m, MonadIO m) =>
SourceConfig -> ConduitT i (BlobKey, ByteString) m ()
Casa.blobsSource
        (SourceConfig :: CasaRepoPrefix -> HashMap BlobKey Int -> Int -> SourceConfig
Casa.SourceConfig
           { sourceConfigUrl :: CasaRepoPrefix
sourceConfigUrl = CasaRepoPrefix
pullUrl
           , sourceConfigBlobs :: HashMap BlobKey Int
sourceConfigBlobs = f BlobKey -> HashMap BlobKey Int
forall (f :: * -> *).
Foldable f =>
f BlobKey -> HashMap BlobKey Int
toBlobKeyMap f BlobKey
keys
           , sourceConfigMaxBlobsPerRequest :: Int
sourceConfigMaxBlobsPerRequest = Int
maxPerRequest
           })
      where
        toBlobKeyMap :: Foldable f => f BlobKey -> HashMap Casa.BlobKey Int
        toBlobKeyMap :: f BlobKey -> HashMap BlobKey Int
toBlobKeyMap = [(BlobKey, Int)] -> HashMap BlobKey Int
forall k v. (Eq k, Hashable k) => [(k, v)] -> HashMap k v
HM.fromList ([(BlobKey, Int)] -> HashMap BlobKey Int)
-> (f BlobKey -> [(BlobKey, Int)])
-> f BlobKey
-> HashMap BlobKey Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (BlobKey -> (BlobKey, Int)) -> [BlobKey] -> [(BlobKey, Int)]
forall a b. (a -> b) -> [a] -> [b]
map BlobKey -> (BlobKey, Int)
forall b. Num b => BlobKey -> (BlobKey, b)
unpackBlobKey ([BlobKey] -> [(BlobKey, Int)])
-> (f BlobKey -> [BlobKey]) -> f BlobKey -> [(BlobKey, Int)]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. f BlobKey -> [BlobKey]
forall (t :: * -> *) a. Foldable t => t a -> [a]
toList
        unpackBlobKey :: BlobKey -> (BlobKey, b)
unpackBlobKey (P.BlobKey SHA256
sha256 (FileSize Word
fileSize)) =
          (ByteString -> BlobKey
Casa.BlobKey (SHA256 -> ByteString
SHA256.toRaw SHA256
sha256), Word -> b
forall a b. (Integral a, Num b) => a -> b
fromIntegral Word
fileSize)
    convert :: ConduitM
  (BlobKey, ByteString)
  (BlobKey, ByteString)
  (ResourceT (ReaderT SqlBackend (RIO env)))
  ()
convert = ((BlobKey, ByteString)
 -> ResourceT (ReaderT SqlBackend (RIO env)) (BlobKey, ByteString))
-> ConduitM
     (BlobKey, ByteString)
     (BlobKey, ByteString)
     (ResourceT (ReaderT SqlBackend (RIO env)))
     ()
forall (m :: * -> *) a b.
Monad m =>
(a -> m b) -> ConduitT a b m ()
mapMC (BlobKey, ByteString)
-> ResourceT (ReaderT SqlBackend (RIO env)) (BlobKey, ByteString)
forall (m :: * -> *).
MonadThrow m =>
(BlobKey, ByteString) -> m (BlobKey, ByteString)
toBlobKeyAndBlob
      where
        toBlobKeyAndBlob ::
             MonadThrow m
          => (Casa.BlobKey, ByteString)
          -> m (BlobKey, ByteString)
        toBlobKeyAndBlob :: (BlobKey, ByteString) -> m (BlobKey, ByteString)
toBlobKeyAndBlob (Casa.BlobKey ByteString
keyBytes, ByteString
blob) = do
          SHA256
sha256 <-
            case ByteString -> Either SHA256Exception SHA256
SHA256.fromRaw ByteString
keyBytes of
              Left SHA256Exception
e -> SHA256Exception -> m SHA256
forall (m :: * -> *) e a. (MonadThrow m, Exception e) => e -> m a
throwM SHA256Exception
e
              Right SHA256
sha -> SHA256 -> m SHA256
forall (f :: * -> *) a. Applicative f => a -> f a
pure SHA256
sha
          (BlobKey, ByteString) -> m (BlobKey, ByteString)
forall (f :: * -> *) a. Applicative f => a -> f a
pure (SHA256 -> FileSize -> BlobKey
BlobKey SHA256
sha256 (Word -> FileSize
FileSize (Int -> Word
forall a b. (Integral a, Num b) => a -> b
fromIntegral (ByteString -> Int
B.length ByteString
blob))), ByteString
blob)
    store :: ConduitT
  (a, ByteString)
  (a, ByteString)
  (ResourceT (ReaderT SqlBackend (RIO env)))
  ()
store = ((a, ByteString)
 -> ResourceT (ReaderT SqlBackend (RIO env)) (a, ByteString))
-> ConduitT
     (a, ByteString)
     (a, ByteString)
     (ResourceT (ReaderT SqlBackend (RIO env)))
     ()
forall (m :: * -> *) a b.
Monad m =>
(a -> m b) -> ConduitT a b m ()
mapMC (a, ByteString)
-> ResourceT (ReaderT SqlBackend (RIO env)) (a, ByteString)
forall (t :: (* -> *) -> * -> *) env a.
(Monad (t (ReaderT SqlBackend (RIO env))), MonadTrans t) =>
(a, ByteString) -> t (ReaderT SqlBackend (RIO env)) (a, ByteString)
insertBlob
      where
        insertBlob :: (a, ByteString) -> t (ReaderT SqlBackend (RIO env)) (a, ByteString)
insertBlob original :: (a, ByteString)
original@(a
_key, ByteString
binary) = do
          (BlobId, BlobKey)
_ <- ReaderT SqlBackend (RIO env) (BlobId, BlobKey)
-> t (ReaderT SqlBackend (RIO env)) (BlobId, BlobKey)
forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift (ByteString -> ReaderT SqlBackend (RIO env) (BlobId, BlobKey)
forall env.
ByteString -> ReaderT SqlBackend (RIO env) (BlobId, BlobKey)
storeBlob ByteString
binary)
          (a, ByteString) -> t (ReaderT SqlBackend (RIO env)) (a, ByteString)
forall (f :: * -> *) a. Applicative f => a -> f a
pure (a, ByteString)
original