{-# 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,
    mkLogger,
  )
where

import Control.Lens
import Data.Aeson.Key qualified as Key
import Data.Aeson.Lens
import Data.ByteString qualified as BS
import Data.ByteString.Builder
import Data.ByteString.Char8 qualified as BS8
import Data.ByteString.Lazy qualified as LBS
import Data.List.NonEmpty (NonEmpty (..))
import Data.List.NonEmpty qualified as NonEmpty
import Data.Text (pack, stripSuffix, unpack)
import Database.CQL.IO as C hiding (values)
import Database.CQL.IO.Tinylog qualified as CT
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
  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 = (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 =
        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 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 = NonEmpty String -> m (NonEmpty String)
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (NonEmpty String -> m (NonEmpty String))
-> (Text -> NonEmpty String) -> Text -> m (NonEmpty String)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. String -> NonEmpty String
forall a. a -> NonEmpty a
NonEmpty.singleton (String -> NonEmpty String)
-> (Text -> String) -> Text -> NonEmpty String
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Text -> String
unpack

-- | 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
  randomPolicy <- IO Policy
C.random
  pure $ 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

mkLogger :: Maybe Text -> Log.Logger -> Logger
mkLogger :: Maybe Text -> Logger -> Logger
mkLogger Maybe Text
mName Logger
logger = Logger
base {logMessage = suppressUseKeyspaceWarning}
  where
    base :: Logger
base = Logger -> Logger
CT.mkLogger (Logger -> (Text -> Logger) -> Maybe Text -> Logger
forall b a. b -> (a -> b) -> Maybe a -> b
maybe Logger
logger (\Text
n -> Maybe Text -> Logger -> Logger
Log.clone (Text -> Maybe Text
forall a. a -> Maybe a
Just Text
n) Logger
logger) Maybe Text
mName)
    isUseKeyspaceWarning :: ByteString -> Bool
isUseKeyspaceWarning ByteString
msg = (ByteString -> Bool) -> [ByteString] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
all (\ByteString
needle -> ByteString
needle ByteString -> ByteString -> Bool
`BS.isInfixOf` ByteString
msg) [ByteString]
useKeyspaceNeedles
    suppressUseKeyspaceWarning :: LogLevel -> Builder -> IO ()
suppressUseKeyspaceWarning LogLevel
level Builder
builder = do
      let msg :: ByteString
msg = ByteString -> ByteString
LBS.toStrict (ByteString -> ByteString) -> ByteString -> ByteString
forall a b. (a -> b) -> a -> b
$ Builder -> ByteString
toLazyByteString Builder
builder
      case (LogLevel
level, ByteString -> Bool
isUseKeyspaceWarning ByteString
msg) of
        (LogLevel
LogWarn, Bool
True) -> () -> IO ()
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
        (LogLevel, Bool)
_ -> Logger -> LogLevel -> Builder -> IO ()
logMessage Logger
base LogLevel
level Builder
builder

-- This is a top-level constant to avoid repeated `pack` allocation on every log line
useKeyspaceNeedles :: [BS.ByteString]
useKeyspaceNeedles :: [ByteString]
useKeyspaceNeedles =
  [ String -> ByteString
BS8.pack String
"non-qualified table names",
    String -> ByteString
BS8.pack String
"Server warning: `USE <keyspace>`"
  ]