Test stream-asset

Author: supersven

services/cargohold/src/CargoHold/API/Federation.hs

@@ -26,6 +26,7 @@ where
 import CargoHold.API.Error
 import CargoHold.API.V3
 import CargoHold.App
+import qualified CargoHold.S3 as S3
 import Control.Error
 import Imports
 import Servant.API
@@ -34,7 +35,6 @@ import Servant.Server hiding (Handler)
 import Servant.Server.Generic
 import Wire.API.Federation.API
 import qualified Wire.API.Federation.API.Cargohold as F
-import qualified CargoHold.S3 as S3
 import Wire.API.Routes.AssetBody
 
 type FederationAPI = "federation" :> ToServantApi (FedApi 'Cargohold)

services/cargohold/src/CargoHold/AWS.hs

@@ -34,16 +34,17 @@ module CargoHold.AWS
     send,
     sendCatch,
     exec,
+    execStream,
     execCatch,
   )
 where
 
 import CargoHold.CloudFront
 import CargoHold.Options
+import Conduit
 import Control.Lens hiding ((.=))
 import Control.Monad.Catch
 import qualified Control.Monad.Trans.AWS as AWST
-import Control.Monad.Trans.Resource
 import Control.Retry
 import Data.ByteString.Builder (toLazyByteString)
 import Imports
@@ -157,7 +158,7 @@ instance Exception Error
 --------------------------------------------------------------------------------
 -- Utilities
 
-sendCatch :: AWSRequest r => r -> Amazon (Either AWS.Error (Rs r))
+sendCatch :: (MonadCatch m, AWS.MonadAWS m, AWSRequest r) => r -> m (Either AWS.Error (Rs r))
 sendCatch = AWST.trying AWS._Error . AWS.send
 
 send :: AWSRequest r => r -> Amazon (Rs r)
@@ -176,7 +177,7 @@ exec env request = do
   resp <- execute env (sendCatch req)
   case resp of
     Left err -> do
-      Log.info $
+      Logger.info (view logger env) $
         Log.field "remote" (Log.val "S3")
           ~~ Log.msg (show err)
           ~~ Log.msg (show req)
@@ -185,6 +186,25 @@ exec env request = do
       throwM (GeneralError err)
     Right r -> return r
 
+execStream ::
+  (AWSRequest r, Show r) =>
+  Env ->
+  (Text -> r) ->
+  ResourceT IO (Rs r)
+execStream env request = do
+  let req = request (_s3Bucket env)
+  resp <- AWS.runAWS (view amazonkaEnv env) (sendCatch req)
+  case resp of
+    Left err -> do
+      Logger.info (view logger env) $
+        Log.field "remote" (Log.val "S3")
+          ~~ Log.msg (show err)
+          ~~ Log.msg (show req)
+      -- We just re-throw the error, but logging it here also gives us the request
+      -- that caused it.
+      throwM (GeneralError err)
+    Right r -> pure r
+
 execCatch ::
   (AWSRequest r, Show r, MonadLogger m, MonadIO m) =>
   Env ->

services/cargohold/src/CargoHold/S3.hs

@@ -129,15 +129,27 @@ uploadV3 prc (s3Key . mkKey -> key) originalHeaders@(V3.AssetHeaders _ cl) tok s
         & poContentType ?~ MIME.showType ct
         & poMetadata .~ metaHeaders tok prc
 
+-- | Turn a 'ResourceT IO' action into a pure @Conduit@.
+--
+-- This is possible because @Conduit@ itself is a monad transformer over
+-- 'ResourceT IO'. Removing the outer 'ResourceT IO' layer makes it possible to
+-- pass this @Conduit@ to resource-oblivious code.
+flattenResourceT ::
+  ResourceT IO (ConduitT () ByteString (ResourceT IO) ()) ->
+  ConduitT () ByteString (ResourceT IO) ()
+flattenResourceT = join . lift
+
 downloadV3 ::
   V3.AssetKey ->
   ExceptT Error App (ConduitM () ByteString (ResourceT IO) ())
 downloadV3 (s3Key . mkKey -> key) = do
-   _streamBody . view gorsBody <$> exec req 
+  env <- view aws
+  pure . flattenResourceT $ _streamBody . view gorsBody <$> AWS.execStream env req
   where
     req :: Text -> GetObject
-    req b = getObject (BucketName b) (ObjectKey key) 
-             & goResponseContentType ?~ MIME.showType octets
+    req b =
+      getObject (BucketName b) (ObjectKey key)
+        & goResponseContentType ?~ MIME.showType octets
 
 getMetadataV3 :: V3.AssetKey -> ExceptT Error App (Maybe S3AssetMeta)
 getMetadataV3 (s3Key . mkKey -> key) = do

services/cargohold/test/integration/API/Federation.hs

@@ -30,7 +30,9 @@ tests s =
         ],
       testGroup
         "stream-asset"
-        [test s "streaming large asset" testLargeAsset]
+        [ test s "streaming large asset" testLargeAsset,
+          test s "stream an asset" testStreamAsset
+        ]
     ]
 
 testGetAssetAvailable :: Bool -> TestSignature ()
@@ -140,3 +142,32 @@ testLargeAsset c = do
     -- check that the first chunk is received
     chunk <- responseBody resp
     print chunk
+
+testStreamAsset :: TestSignature ()
+testStreamAsset c = do
+  -- Initial upload
+  let bdy = (applicationOctetStream, "Hello World")
+      settings =
+        defAssetSettings
+          & set setAssetRetention (Just AssetVolatile)
+  uid <- liftIO $ Id <$> nextRandom
+  ast :: Asset <-
+    responseJsonError
+      =<< uploadSimple (c . path "/assets/v3") uid settings bdy
+      <!! const 201 === statusCode
+
+  -- Call get-asset federation API
+  let tok = view assetToken ast
+  let key = view assetKey ast
+  let ga =
+        GetAsset
+          { gaUser = uid,
+            gaToken = tok,
+            gaKey = qUnqualified key
+          }
+  respBody <- fmap responseBody $
+       post (c . path "/federation/stream-asset" . json ga)
+       <!! const 200 === statusCode
+
+
+  liftIO $ respBody @?= Just "Hello World"