{-# LANGUAGE ConstraintKinds #-}
{-# LANGUAGE DataKinds #-}
{-# LANGUAGE GADTs #-}
{-# LANGUAGE TypeOperators #-}

{-|

Functions for managing images on a Kubernetes cluster.

-}

module Test.Sandwich.Contexts.Kubernetes.Images (
  -- * Introduce a set of images
  introduceImages

  -- * Query images
  , getLoadedImages
  , getLoadedImages'

  , clusterContainsImage
  , clusterContainsImage'

  -- * Load images
  , loadImage
  , loadImage'

  -- * Load images if not present
  , loadImageIfNecessary
  , loadImageIfNecessary'

  -- * Retry helpers
  , withImageLoadRetry
  , withImageLoadRetry'

  -- * Util
  , findAllImages
  , findAllImages'

  -- * Types
  , kubernetesClusterImages
  , HasKubernetesClusterImagesContext

  , ImageLoadSpec(..)
  , ImagePullPolicy(..)
  ) where

import Control.Monad.Catch (Handler(..), MonadMask)
import Control.Monad.IO.Unlift
import Control.Monad.Logger
import Control.Retry
import Data.String.Interpolate
import Data.Text as T
import Relude
import Test.Sandwich
import Test.Sandwich.Contexts.Kubernetes.FindImages
import qualified Test.Sandwich.Contexts.Kubernetes.KindCluster as Kind
import qualified Test.Sandwich.Contexts.Kubernetes.MinikubeCluster.Images as Minikube
import Test.Sandwich.Contexts.Kubernetes.Types
import Test.Sandwich.Contexts.Kubernetes.Util.Images


-- | Get the images loaded onto the cluster.
getLoadedImages :: (
  HasCallStack, KubernetesClusterBasic context m
  )
  -- | List of image names
  => m (Set Text)
getLoadedImages :: forall context (m :: * -> *).
(HasCallStack, KubernetesClusterBasic context m) =>
m (Set Text)
getLoadedImages = Label "kubernetesCluster" KubernetesClusterContext
-> m KubernetesClusterContext
forall context (l :: Symbol) a (m :: * -> *).
(HasLabel context l a, MonadReader context m) =>
Label l a -> m a
getContext Label "kubernetesCluster" KubernetesClusterContext
kubernetesCluster m KubernetesClusterContext
-> (KubernetesClusterContext -> m (Set Text)) -> m (Set Text)
forall a b. m a -> (a -> m b) -> m b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= KubernetesClusterContext -> m (Set Text)
forall context (m :: * -> *).
(HasCallStack, KubernetesBasic context m) =>
KubernetesClusterContext -> m (Set Text)
getLoadedImages'

-- | Same as 'getLoadedImages', but allows you to pass in the 'KubernetesClusterContext'.
getLoadedImages' :: (
  HasCallStack, KubernetesBasic context m
  )
  -- | Cluster context
  => KubernetesClusterContext
  -- | List of image names
  -> m (Set Text)
getLoadedImages' :: forall context (m :: * -> *).
(HasCallStack, KubernetesBasic context m) =>
KubernetesClusterContext -> m (Set Text)
getLoadedImages' kcc :: KubernetesClusterContext
kcc@(KubernetesClusterContext {KubernetesClusterType
kubernetesClusterType :: KubernetesClusterType
kubernetesClusterType :: KubernetesClusterContext -> KubernetesClusterType
kubernetesClusterType, Text
kubernetesClusterName :: Text
kubernetesClusterName :: KubernetesClusterContext -> Text
kubernetesClusterName}) = do
  Text -> m (Set Text) -> m (Set Text)
forall (m :: * -> *) context a.
(MonadUnliftIO m, HasBaseContextMonad context m,
 HasTestTimer context) =>
Text -> m a -> m a
timeAction [i|Getting loaded images|] (m (Set Text) -> m (Set Text)) -> m (Set Text) -> m (Set Text)
forall a b. (a -> b) -> a -> b
$ do
    case KubernetesClusterType
kubernetesClusterType of
      (KubernetesClusterKind {String
Maybe [(String, String)]
Text
kubernetesClusterTypeKindBinary :: String
kubernetesClusterTypeKindClusterName :: Text
kubernetesClusterTypeKindClusterDriver :: Text
kubernetesClusterTypeKindClusterEnvironment :: Maybe [(String, String)]
kubernetesClusterTypeKindClusterEnvironment :: KubernetesClusterType -> Maybe [(String, String)]
kubernetesClusterTypeKindClusterDriver :: KubernetesClusterType -> Text
kubernetesClusterTypeKindClusterName :: KubernetesClusterType -> Text
kubernetesClusterTypeKindBinary :: KubernetesClusterType -> String
..}) ->
        KubernetesClusterContext
-> Text -> String -> Maybe [(String, String)] -> m (Set Text)
forall (m :: * -> *).
(HasCallStack, MonadUnliftIO m, MonadLogger m) =>
KubernetesClusterContext
-> Text -> String -> Maybe [(String, String)] -> m (Set Text)
Kind.getLoadedImagesKind KubernetesClusterContext
kcc Text
kubernetesClusterTypeKindClusterDriver String
kubernetesClusterTypeKindBinary Maybe [(String, String)]
forall a. Maybe a
Nothing
        -- Kind.loadImage kindBinary kindClusterName image env
      (KubernetesClusterMinikube {String
[Text]
Text
kubernetesClusterTypeMinikubeBinary :: String
kubernetesClusterTypeMinikubeProfileName :: Text
kubernetesClusterTypeMinikubeFlags :: [Text]
kubernetesClusterTypeMinikubeFlags :: KubernetesClusterType -> [Text]
kubernetesClusterTypeMinikubeProfileName :: KubernetesClusterType -> Text
kubernetesClusterTypeMinikubeBinary :: KubernetesClusterType -> String
..}) ->
        -- Note: don't pass minikubeFlags here. These are pretty much intended for "minikube start" only.
        -- TODO: clarify the documentation and possibly add an extra field where extra options can be passed
        -- to "minikube image" commands.
        String -> Text -> [Text] -> m (Set Text)
forall (m :: * -> *).
(MonadUnliftIO m, MonadLogger m) =>
String -> Text -> [Text] -> m (Set Text)
Minikube.getLoadedImagesMinikube String
kubernetesClusterTypeMinikubeBinary Text
kubernetesClusterName []

-- | Test if a cluster has a given image loaded.
clusterContainsImage :: (
  HasCallStack, KubernetesClusterBasic context m
  )
  -- | Image
  => Text
  -> m Bool
clusterContainsImage :: forall context (m :: * -> *).
(HasCallStack, KubernetesClusterBasic context m) =>
Text -> m Bool
clusterContainsImage Text
image = do
  kcc <- Label "kubernetesCluster" KubernetesClusterContext
-> m KubernetesClusterContext
forall context (l :: Symbol) a (m :: * -> *).
(HasLabel context l a, MonadReader context m) =>
Label l a -> m a
getContext Label "kubernetesCluster" KubernetesClusterContext
kubernetesCluster
  clusterContainsImage' kcc image

-- | Same as 'clusterContainsImage', but allows you to pass in the 'KubernetesClusterContext'.
clusterContainsImage' :: (
  HasCallStack, MonadUnliftIO m, MonadLogger m
  )
  -- | Cluster context
  => KubernetesClusterContext
  -- | Image
  -> Text
  -> m Bool
clusterContainsImage' :: forall (m :: * -> *).
(HasCallStack, MonadUnliftIO m, MonadLogger m) =>
KubernetesClusterContext -> Text -> m Bool
clusterContainsImage' kcc :: KubernetesClusterContext
kcc@(KubernetesClusterContext {KubernetesClusterType
kubernetesClusterType :: KubernetesClusterContext -> KubernetesClusterType
kubernetesClusterType :: KubernetesClusterType
kubernetesClusterType, Text
kubernetesClusterName :: KubernetesClusterContext -> Text
kubernetesClusterName :: Text
kubernetesClusterName}) Text
image = do
  case KubernetesClusterType
kubernetesClusterType of
    KubernetesClusterKind {String
Maybe [(String, String)]
Text
kubernetesClusterTypeKindClusterEnvironment :: KubernetesClusterType -> Maybe [(String, String)]
kubernetesClusterTypeKindClusterDriver :: KubernetesClusterType -> Text
kubernetesClusterTypeKindClusterName :: KubernetesClusterType -> Text
kubernetesClusterTypeKindBinary :: KubernetesClusterType -> String
kubernetesClusterTypeKindBinary :: String
kubernetesClusterTypeKindClusterName :: Text
kubernetesClusterTypeKindClusterDriver :: Text
kubernetesClusterTypeKindClusterEnvironment :: Maybe [(String, String)]
..} ->
      KubernetesClusterContext
-> Text -> String -> Maybe [(String, String)] -> Text -> m Bool
forall (m :: * -> *).
(HasCallStack, MonadUnliftIO m, MonadLogger m) =>
KubernetesClusterContext
-> Text -> String -> Maybe [(String, String)] -> Text -> m Bool
Kind.clusterContainsImageKind KubernetesClusterContext
kcc Text
kubernetesClusterTypeKindClusterDriver String
kubernetesClusterTypeKindBinary Maybe [(String, String)]
kubernetesClusterTypeKindClusterEnvironment Text
image
    KubernetesClusterMinikube {String
[Text]
Text
kubernetesClusterTypeMinikubeFlags :: KubernetesClusterType -> [Text]
kubernetesClusterTypeMinikubeProfileName :: KubernetesClusterType -> Text
kubernetesClusterTypeMinikubeBinary :: KubernetesClusterType -> String
kubernetesClusterTypeMinikubeBinary :: String
kubernetesClusterTypeMinikubeProfileName :: Text
kubernetesClusterTypeMinikubeFlags :: [Text]
..} ->
      String -> Text -> [Text] -> Text -> m Bool
forall (m :: * -> *).
(MonadUnliftIO m, MonadLogger m) =>
String -> Text -> [Text] -> Text -> m Bool
Minikube.clusterContainsImageMinikube String
kubernetesClusterTypeMinikubeBinary Text
kubernetesClusterName [] Text
image

-- | Same as 'loadImage', but first checks if the given image is already present on the cluster.
loadImageIfNecessary :: (
  HasCallStack, MonadFail m, KubernetesClusterBasic context m
  )
  -- | Image load spec
  => ImageLoadSpec
  -> m ()
loadImageIfNecessary :: forall (m :: * -> *) context.
(HasCallStack, MonadFail m, KubernetesClusterBasic context m) =>
ImageLoadSpec -> m ()
loadImageIfNecessary ImageLoadSpec
image = do
  kcc <- Label "kubernetesCluster" KubernetesClusterContext
-> m KubernetesClusterContext
forall context (l :: Symbol) a (m :: * -> *).
(HasLabel context l a, MonadReader context m) =>
Label l a -> m a
getContext Label "kubernetesCluster" KubernetesClusterContext
kubernetesCluster
  loadImageIfNecessary' kcc image

-- | Same as 'loadImage', but allows you to pass in the 'KubernetesClusterContext'.
loadImageIfNecessary' :: (
  HasCallStack, MonadFail m, KubernetesBasic context m
  )
  -- | Cluster context
  => KubernetesClusterContext
  -- | Image load spec
  -> ImageLoadSpec
  -- | The transformed image name
  -> m ()
loadImageIfNecessary' :: forall (m :: * -> *) context.
(HasCallStack, MonadFail m, KubernetesBasic context m) =>
KubernetesClusterContext -> ImageLoadSpec -> m ()
loadImageIfNecessary' KubernetesClusterContext
kcc ImageLoadSpec
imageLoadSpec = do
  m Bool -> m () -> m ()
forall (m :: * -> *). Monad m => m Bool -> m () -> m ()
unlessM (ImageLoadSpec -> m Text
forall (m :: * -> *).
(MonadUnliftIO m, MonadLogger m) =>
ImageLoadSpec -> m Text
imageLoadSpecToImageName ImageLoadSpec
imageLoadSpec m Text -> (Text -> m Bool) -> m Bool
forall a b. m a -> (a -> m b) -> m b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= KubernetesClusterContext -> Text -> m Bool
forall (m :: * -> *).
(HasCallStack, MonadUnliftIO m, MonadLogger m) =>
KubernetesClusterContext -> Text -> m Bool
clusterContainsImage' KubernetesClusterContext
kcc) (m () -> m ()) -> m () -> m ()
forall a b. (a -> b) -> a -> b
$
    m Text -> m ()
forall (f :: * -> *) a. Functor f => f a -> f ()
void (m Text -> m ()) -> m Text -> m ()
forall a b. (a -> b) -> a -> b
$ KubernetesClusterContext -> ImageLoadSpec -> m Text
forall (m :: * -> *) context.
(HasCallStack, MonadFail m, KubernetesBasic context m) =>
KubernetesClusterContext -> ImageLoadSpec -> m Text
loadImage' KubernetesClusterContext
kcc ImageLoadSpec
imageLoadSpec

-- | Load an image into a Kubernetes cluster. This will load the image onto the cluster
-- and return the modified image name (i.e. the name by which the cluster knows the image).
loadImage :: (
  HasCallStack, MonadFail m, KubernetesClusterBasic context m
  )
  -- | Image load spec
  => ImageLoadSpec
  -- | The loaded image name
  -> m Text
loadImage :: forall (m :: * -> *) context.
(HasCallStack, MonadFail m, KubernetesClusterBasic context m) =>
ImageLoadSpec -> m Text
loadImage ImageLoadSpec
image = do
  kcc <- Label "kubernetesCluster" KubernetesClusterContext
-> m KubernetesClusterContext
forall context (l :: Symbol) a (m :: * -> *).
(HasLabel context l a, MonadReader context m) =>
Label l a -> m a
getContext Label "kubernetesCluster" KubernetesClusterContext
kubernetesCluster
  loadImage' kcc image

-- | Same as 'loadImage', but allows you to pass in the 'KubernetesClusterContext'.
loadImage' :: (
  HasCallStack, MonadFail m, KubernetesBasic context m
  )
  -- | Cluster context
  => KubernetesClusterContext
  -- | Image load spec
  -> ImageLoadSpec
  -- | The loaded image name
  -> m Text
loadImage' :: forall (m :: * -> *) context.
(HasCallStack, MonadFail m, KubernetesBasic context m) =>
KubernetesClusterContext -> ImageLoadSpec -> m Text
loadImage' (KubernetesClusterContext {KubernetesClusterType
kubernetesClusterType :: KubernetesClusterContext -> KubernetesClusterType
kubernetesClusterType :: KubernetesClusterType
kubernetesClusterType, Text
kubernetesClusterName :: KubernetesClusterContext -> Text
kubernetesClusterName :: Text
kubernetesClusterName}) ImageLoadSpec
imageLoadSpec = do
  Text -> m ()
forall (m :: * -> *). (HasCallStack, MonadLogger m) => Text -> m ()
debug [i|Loading container image '#{imageLoadSpec}'|]
  Text -> m Text -> m Text
forall (m :: * -> *) context a.
(MonadUnliftIO m, HasBaseContextMonad context m,
 HasTestTimer context) =>
Text -> m a -> m a
timeAction [i|Loading container image '#{imageLoadSpec}'|] (m Text -> m Text) -> m Text -> m Text
forall a b. (a -> b) -> a -> b
$ do
    case KubernetesClusterType
kubernetesClusterType of
      (KubernetesClusterKind {String
Maybe [(String, String)]
Text
kubernetesClusterTypeKindClusterEnvironment :: KubernetesClusterType -> Maybe [(String, String)]
kubernetesClusterTypeKindClusterDriver :: KubernetesClusterType -> Text
kubernetesClusterTypeKindClusterName :: KubernetesClusterType -> Text
kubernetesClusterTypeKindBinary :: KubernetesClusterType -> String
kubernetesClusterTypeKindBinary :: String
kubernetesClusterTypeKindClusterName :: Text
kubernetesClusterTypeKindClusterDriver :: Text
kubernetesClusterTypeKindClusterEnvironment :: Maybe [(String, String)]
..}) ->
        String
-> Text -> ImageLoadSpec -> Maybe [(String, String)] -> m Text
forall (m :: * -> *).
(HasCallStack, MonadUnliftIO m, MonadLoggerIO m) =>
String
-> Text -> ImageLoadSpec -> Maybe [(String, String)] -> m Text
Kind.loadImageKind String
kubernetesClusterTypeKindBinary Text
kubernetesClusterTypeKindClusterName ImageLoadSpec
imageLoadSpec Maybe [(String, String)]
kubernetesClusterTypeKindClusterEnvironment
      (KubernetesClusterMinikube {String
[Text]
Text
kubernetesClusterTypeMinikubeFlags :: KubernetesClusterType -> [Text]
kubernetesClusterTypeMinikubeProfileName :: KubernetesClusterType -> Text
kubernetesClusterTypeMinikubeBinary :: KubernetesClusterType -> String
kubernetesClusterTypeMinikubeBinary :: String
kubernetesClusterTypeMinikubeProfileName :: Text
kubernetesClusterTypeMinikubeFlags :: [Text]
..}) ->
        -- Don't pass minikubeFlags; see comment above.
        String -> Text -> [Text] -> ImageLoadSpec -> m Text
forall (m :: * -> *).
(HasCallStack, MonadUnliftIO m, MonadLoggerIO m, MonadFail m) =>
String -> Text -> [Text] -> ImageLoadSpec -> m Text
Minikube.loadImageMinikube String
kubernetesClusterTypeMinikubeBinary Text
kubernetesClusterName [] ImageLoadSpec
imageLoadSpec

        -- Because of the possible silent failure in "minikube image load", confirm that this
        -- image made it onto the cluster.
        -- At the moment this approach doesn't work, because if you do
        -- "minikube image load busybox:1.36.1-musl"
        -- followed by
        -- "minikube image ls",
        -- the result contains "docker.io/library/busybox:1.36.1-musl".
        -- Where did the docker.io/library/ come from? Need to understand this before we can
        -- check this properly.
        --
        -- image' <- Minikube.loadImage minikubeBinary kubernetesClusterName minikubeFlags image
        -- loadedImages <- Set.toList <$> getLoadedImages' kcc
        -- loadedImages `shouldContain` [image']
        -- return image'

-- | A combinator to wrap around your 'loadImage' or 'loadImageIfNecessary' calls to retry
-- on failure. Image loads sometimes fail on Minikube (version 1.33.0 at time of writing).
withImageLoadRetry :: (MonadLoggerIO m, MonadMask m) => ImageLoadSpec -> m a -> m a
withImageLoadRetry :: forall (m :: * -> *) a.
(MonadLoggerIO m, MonadMask m) =>
ImageLoadSpec -> m a -> m a
withImageLoadRetry = RetryPolicyM m -> ImageLoadSpec -> m a -> m a
forall (m :: * -> *) a.
(MonadLoggerIO m, MonadMask m) =>
RetryPolicyM m -> ImageLoadSpec -> m a -> m a
withImageLoadRetry' (Int -> RetryPolicyM m
forall (m :: * -> *). Monad m => Int -> RetryPolicyM m
exponentialBackoff Int
50000 RetryPolicyM m -> RetryPolicyM m -> RetryPolicyM m
forall a. Semigroup a => a -> a -> a
<> Int -> RetryPolicy
limitRetries Int
5)

-- | Same as 'withImageLoadRetry', but allows you to specify the retry policy.
withImageLoadRetry' :: (MonadLoggerIO m, MonadMask m) => RetryPolicyM m -> ImageLoadSpec -> m a -> m a
withImageLoadRetry' :: forall (m :: * -> *) a.
(MonadLoggerIO m, MonadMask m) =>
RetryPolicyM m -> ImageLoadSpec -> m a -> m a
withImageLoadRetry' RetryPolicyM m
policy ImageLoadSpec
ils m a
action =
  RetryPolicyM m
-> [RetryStatus -> Handler m Bool] -> (RetryStatus -> m a) -> m a
forall (m :: * -> *) a.
(MonadIO m, MonadMask m) =>
RetryPolicyM m
-> [RetryStatus -> Handler m Bool] -> (RetryStatus -> m a) -> m a
recovering RetryPolicyM m
policy [\RetryStatus
_status -> (FailureReason -> m Bool) -> Handler m Bool
forall (m :: * -> *) a e. Exception e => (e -> m a) -> Handler m a
Handler (\(FailureReason
e :: FailureReason) -> do
                                             Text -> m ()
forall (m :: * -> *). (HasCallStack, MonadLogger m) => Text -> m ()
warn [i|#{ils}: retrying load due to exception: #{e}|]
                                             Bool -> m Bool
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return Bool
True)] ((RetryStatus -> m a) -> m a) -> (RetryStatus -> m a) -> m a
forall a b. (a -> b) -> a -> b
$ \RetryStatus
_ ->
    m a
action


-- | Introduce a list of images into a Kubernetes cluster.
-- Stores the list of transformed image names under the "kubernetesClusterImages" label.
introduceImages :: (
  HasCallStack, KubernetesClusterBasicWithoutReader context m
  )
  -- | Images to load
  => [ImageLoadSpec]
  -> SpecFree (LabelValue "kubernetesClusterImages" [Text] :> context) m ()
  -> SpecFree context m ()
introduceImages :: forall context (m :: * -> *).
(HasCallStack, KubernetesClusterBasicWithoutReader context m) =>
[ImageLoadSpec]
-> SpecFree
     (LabelValue "kubernetesClusterImages" [Text] :> context) m ()
-> SpecFree context m ()
introduceImages [ImageLoadSpec]
images = String
-> Label "kubernetesClusterImages" [Text]
-> ((HasCallStack => [Text] -> ExampleT context m [Result])
    -> ExampleT context m ())
-> SpecFree
     (LabelValue "kubernetesClusterImages" [Text] :> context) m ()
-> SpecFree context m ()
forall (l :: Symbol) intro context (m :: * -> *).
HasCallStack =>
String
-> Label l intro
-> ((HasCallStack => intro -> ExampleT context m [Result])
    -> ExampleT context m ())
-> SpecFree (LabelValue l intro :> context) m ()
-> SpecFree context m ()
introduceWith String
"Introduce cluster images" Label "kubernetesClusterImages" [Text]
kubernetesClusterImages (((HasCallStack => [Text] -> ExampleT context m [Result])
  -> ExampleT context m ())
 -> SpecFree
      (LabelValue "kubernetesClusterImages" [Text] :> context) m ()
 -> SpecFree context m ())
-> ((HasCallStack => [Text] -> ExampleT context m [Result])
    -> ExampleT context m ())
-> SpecFree
     (LabelValue "kubernetesClusterImages" [Text] :> context) m ()
-> SpecFree context m ()
forall a b. (a -> b) -> a -> b
$ \HasCallStack => [Text] -> ExampleT context m [Result]
action ->
  [ImageLoadSpec]
-> (ImageLoadSpec -> ExampleT context m Text)
-> ExampleT context m [Text]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
t a -> (a -> m b) -> m (t b)
forM [ImageLoadSpec]
images (\ImageLoadSpec
x -> ImageLoadSpec -> ExampleT context m Text
forall (m :: * -> *) context.
(HasCallStack, MonadFail m, KubernetesClusterBasic context m) =>
ImageLoadSpec -> m Text
loadImage ImageLoadSpec
x) ExampleT context m [Text]
-> ([Text] -> ExampleT context m ()) -> ExampleT context m ()
forall a b.
ExampleT context m a
-> (a -> ExampleT context m b) -> ExampleT context m b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= (ExampleT context m [Result] -> ExampleT context m ()
forall (f :: * -> *) a. Functor f => f a -> f ()
void (ExampleT context m [Result] -> ExampleT context m ())
-> ([Text] -> ExampleT context m [Result])
-> [Text]
-> ExampleT context m ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. HasCallStack => [Text] -> ExampleT context m [Result]
[Text] -> ExampleT context m [Result]
action)