{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE ViewPatterns #-}

-- This file is part of the Wire Server implementation.
--
-- Copyright (C) 2022 Wire Swiss GmbH <opensource@wire.com>
--
-- This program is free software: you can redistribute it and/or modify it under
-- the terms of the GNU Affero General Public License as published by the Free
-- Software Foundation, either version 3 of the License, or (at your option) any
-- later version.
--
-- This program is distributed in the hope that it will be useful, but WITHOUT
-- ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or FITNESS
-- FOR A PARTICULAR PURPOSE. See the GNU Affero General Public License for more
-- details.
--
-- You should have received a copy of the GNU Affero General Public License along
-- with this program. If not, see <https://www.gnu.org/licenses/>.

-- | This module exports types and functions from Database.CQL.IO, while adding a few wire specific functions.
module Cassandra.Settings
  ( module C,
    initialContactsDisco,
    initialContactsPlain,
    dcAwareRandomPolicy,
    dcFilterPolicyIfConfigured,
  )
where

import Control.Lens
import Data.Aeson.Key qualified as Key
import Data.Aeson.Lens
import Data.List.NonEmpty (NonEmpty (..))
import Data.Text (pack, stripSuffix, unpack)
import Database.CQL.IO as C hiding (values)
import Database.CQL.IO.Tinylog as C (mkLogger)
import Imports
import Network.Wreq
import System.Logger qualified as Log

-- | This function is likely only useful at Wire, as it is Wire-infra specific.
-- Given a server name and a url returning a wire-custom "disco" json (AWS describe-instances-like json), e.g.
-- { "roles" : { "server_name": [ {"privateIpAddress": "...", ...}, {...} ] } },
-- return a list of IP addresses.
initialContactsDisco :: (MonadIO m) => String -> String -> m (NonEmpty String)
initialContactsDisco :: forall (m :: * -> *).
MonadIO m =>
String -> String -> m (NonEmpty String)
initialContactsDisco (String -> Text
pack -> Text
srv) String
url = IO (NonEmpty String) -> m (NonEmpty String)
forall a. IO a -> m a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO (NonEmpty String) -> m (NonEmpty String))
-> IO (NonEmpty String) -> m (NonEmpty String)
forall a b. (a -> b) -> a -> b
$ do
  Response Value
rs <- Response ByteString -> IO (Response Value)
forall (m :: * -> *).
MonadThrow m =>
Response ByteString -> m (Response Value)
asValue (Response ByteString -> IO (Response Value))
-> IO (Response ByteString) -> IO (Response Value)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< String -> IO (Response ByteString)
get String
url
  let srvs :: [Key]
srvs = (Text -> Key) -> [Text] -> [Key]
forall a b. (a -> b) -> [a] -> [b]
map Text -> Key
Key.fromText ([Text] -> [Key]) -> [Text] -> [Key]
forall a b. (a -> b) -> a -> b
$
        case Text -> Text -> Maybe Text
stripSuffix Text
"_seed" Text
srv of
          Maybe Text
Nothing -> [Text
srv, Text
srv Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
"_seed"]
          Just Text
_ -> [Text
srv] -- requesting only seeds is a valid use-case
  let ip :: [String]
ip =
        Response Value
rs
          Response Value
-> Getting (Endo [Text]) (Response Value) Text -> [Text]
forall s a. s -> Getting (Endo [a]) s a -> [a]
^.. (Value -> Const (Endo [Text]) Value)
-> Response Value -> Const (Endo [Text]) (Response Value)
forall body0 body1 (f :: * -> *).
Functor f =>
(body0 -> f body1) -> Response body0 -> f (Response body1)
responseBody
            ((Value -> Const (Endo [Text]) Value)
 -> Response Value -> Const (Endo [Text]) (Response Value))
-> ((Text -> Const (Endo [Text]) Text)
    -> Value -> Const (Endo [Text]) Value)
-> Getting (Endo [Text]) (Response Value) Text
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Key -> Traversal' Value Value
forall t. AsValue t => Key -> Traversal' t Value
key Key
"roles"
            ((Value -> Const (Endo [Text]) Value)
 -> Value -> Const (Endo [Text]) Value)
-> ((Text -> Const (Endo [Text]) Text)
    -> Value -> Const (Endo [Text]) Value)
-> (Text -> Const (Endo [Text]) Text)
-> Value
-> Const (Endo [Text]) Value
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Indexed Key Value (Const (Endo [Text]) Value)
-> Value -> Const (Endo [Text]) Value
forall t. AsValue t => IndexedTraversal' Key t Value
IndexedTraversal' Key Value Value
members
            (Indexed Key Value (Const (Endo [Text]) Value)
 -> Value -> Const (Endo [Text]) Value)
-> ((Text -> Const (Endo [Text]) Text)
    -> Indexed Key Value (Const (Endo [Text]) Value))
-> (Text -> Const (Endo [Text]) Text)
-> Value
-> Const (Endo [Text]) Value
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Key -> Bool)
-> Optical' (->) (Indexed Key) (Const (Endo [Text])) Value Value
forall i (p :: * -> * -> *) (f :: * -> *) a.
(Indexable i p, Applicative f) =>
(i -> Bool) -> Optical' p (Indexed i) f a a
indices (Key -> [Key] -> Bool
forall a. Eq a => a -> [a] -> Bool
forall (t :: * -> *) a. (Foldable t, Eq a) => a -> t a -> Bool
`elem` [Key]
srvs)
            Optical' (->) (Indexed Key) (Const (Endo [Text])) Value Value
-> ((Text -> Const (Endo [Text]) Text)
    -> Value -> Const (Endo [Text]) Value)
-> (Text -> Const (Endo [Text]) Text)
-> Indexed Key Value (Const (Endo [Text]) Value)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Value -> Const (Endo [Text]) Value)
-> Value -> Const (Endo [Text]) Value
forall t. AsValue t => IndexedTraversal' Int t Value
IndexedTraversal' Int Value Value
values
            ((Value -> Const (Endo [Text]) Value)
 -> Value -> Const (Endo [Text]) Value)
-> ((Text -> Const (Endo [Text]) Text)
    -> Value -> Const (Endo [Text]) Value)
-> (Text -> Const (Endo [Text]) Text)
-> Value
-> Const (Endo [Text]) Value
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Key -> Traversal' Value Value
forall t. AsValue t => Key -> Traversal' t Value
key Key
"privateIpAddress"
            ((Value -> Const (Endo [Text]) Value)
 -> Value -> Const (Endo [Text]) Value)
-> ((Text -> Const (Endo [Text]) Text)
    -> Value -> Const (Endo [Text]) Value)
-> (Text -> Const (Endo [Text]) Text)
-> Value
-> Const (Endo [Text]) Value
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Text -> Const (Endo [Text]) Text)
-> Value -> Const (Endo [Text]) Value
forall t. AsValue t => Prism' t Text
Prism' Value Text
_String
          [Text] -> ([Text] -> [String]) -> [String]
forall a b. a -> (a -> b) -> b
& (Text -> String) -> [Text] -> [String]
forall a b. (a -> b) -> [a] -> [b]
map Text -> String
unpack
  case [String]
ip of
    String
i : [String]
ii -> NonEmpty String -> IO (NonEmpty String)
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (String
i String -> [String] -> NonEmpty String
forall a. a -> [a] -> NonEmpty a
:| [String]
ii)
    [String]
_ -> String -> IO (NonEmpty String)
forall a. HasCallStack => String -> a
error String
"initial-contacts: no IP addresses found."

-- | Puts the address into a list using the same signature as the other initialContacts
initialContactsPlain :: (MonadIO m) => Text -> m (NonEmpty String)
initialContactsPlain :: forall (m :: * -> *). MonadIO m => Text -> m (NonEmpty String)
initialContactsPlain Text
address = NonEmpty String -> m (NonEmpty String)
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (NonEmpty String -> m (NonEmpty String))
-> NonEmpty String -> m (NonEmpty String)
forall a b. (a -> b) -> a -> b
$ Text -> String
unpack Text
address String -> [String] -> NonEmpty String
forall a. a -> [a] -> NonEmpty a
:| []

-- | Use dcAwareRandomPolicy if config option filterNodesByDatacentre is set,
-- otherwise use all available nodes with the default random policy.
--
-- This is only useful during a cassandra datacentre migration.
dcFilterPolicyIfConfigured :: Log.Logger -> Maybe Text -> IO Policy
dcFilterPolicyIfConfigured :: Logger -> Maybe Text -> IO Policy
dcFilterPolicyIfConfigured Logger
lgr Maybe Text
mDatacentre = do
  Logger -> (Msg -> Msg) -> IO ()
forall (m :: * -> *). MonadIO m => Logger -> (Msg -> Msg) -> m ()
Log.info Logger
lgr ((Msg -> Msg) -> IO ()) -> (Msg -> Msg) -> IO ()
forall a b. (a -> b) -> a -> b
$
    Text -> Msg -> Msg
forall a. ToBytes a => a -> Msg -> Msg
Log.msg (Text
"Using the following cassandra load balancing options ('Policy'):" :: Text)
      (Msg -> Msg) -> (Msg -> Msg) -> Msg -> Msg
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ByteString -> String -> Msg -> Msg
forall a. ToBytes a => ByteString -> a -> Msg -> Msg
Log.field ByteString
"filter_datacentre" (Maybe Text -> String
forall a. Show a => a -> String
show Maybe Text
mDatacentre)
  IO Policy -> (Text -> IO Policy) -> Maybe Text -> IO Policy
forall b a. b -> (a -> b) -> Maybe a -> b
maybe IO Policy
random Text -> IO Policy
dcAwareRandomPolicy Maybe Text
mDatacentre

-- | Return hosts in random order for a given DC.
--
-- This is only useful during a cassandra datacentre migration.
dcAwareRandomPolicy :: Text -> IO Policy
dcAwareRandomPolicy :: Text -> IO Policy
dcAwareRandomPolicy Text
dc = do
  Policy
randomPolicy <- IO Policy
C.random
  Policy -> IO Policy
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Policy -> IO Policy) -> Policy -> IO Policy
forall a b. (a -> b) -> a -> b
$ Policy
randomPolicy {acceptable = dcAcceptable}
  where
    dcAcceptable :: Host -> IO Bool
    dcAcceptable :: Host -> IO Bool
dcAcceptable Host
host = Bool -> IO Bool
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Bool -> IO Bool) -> Bool -> IO Bool
forall a b. (a -> b) -> a -> b
$ (Host
host Host -> Getting Text Host Text -> Text
forall s a. s -> Getting a s a -> a
^. Getting Text Host Text
Lens' Host Text
dataCentre) Text -> Text -> Bool
forall a. Eq a => a -> a -> Bool
== Text
dc