spike: fully qualify cassandra table names

libs/cassandra-util/src/Cassandra/Util.hs

@@ -19,6 +19,9 @@ module Cassandra.Util
   ( defInitCassandra,
     initCassandraForService,
     initCassandra,
+    registerClientKeyspace,
+    lookupClientKeyspace,
+    requireClientKeyspace,
     Writetime (..),
     writetimeToInt64,
   )
@@ -39,7 +42,39 @@ import Database.CQL.IO
 import Database.CQL.IO.Tinylog qualified as CT
 import Imports hiding (init)
 import OpenSSL.Session qualified as OpenSSL
+import System.IO.Unsafe (unsafePerformIO)
 import System.Logger qualified as Log
+import System.Mem.StableName (StableName, eqStableName, makeStableName)
+
+{-# NOINLINE clientKeyspaces #-}
+clientKeyspaces :: IORef [(StableName ClientState, Keyspace)]
+clientKeyspaces = unsafePerformIO $ newIORef []
+
+registerClientKeyspace :: ClientState -> Keyspace -> IO ()
+registerClientKeyspace clientState keyspace = do
+  stable <- makeStableName clientState
+  atomicModifyIORef' clientKeyspaces $ \entries ->
+    ((stable, keyspace) : entries, ())
+
+lookupClientKeyspace :: ClientState -> IO (Maybe Keyspace)
+lookupClientKeyspace clientState = do
+  stable <- makeStableName clientState
+  entries <- readIORef clientKeyspaces
+  pure $ go stable entries
+  where
+    go _ [] = Nothing
+    go needle ((candidate, keyspace) : rest) =
+      if eqStableName needle candidate
+        then Just keyspace
+        else go needle rest
+
+requireClientKeyspace :: (HasCallStack) => ClientState -> IO Keyspace
+requireClientKeyspace clientState =
+  lookupClientKeyspace clientState >>= \case
+    Just keyspace -> pure keyspace
+    Nothing ->
+      error
+        "Missing keyspace for Cassandra ClientState. Initialize via Cassandra.Util.defInitCassandra/initCassandraForService or call registerClientKeyspace."
 
 defInitCassandra :: CassandraOpts -> Log.Logger -> IO ClientState
 defInitCassandra opts logger = do
@@ -50,7 +85,9 @@ defInitCassandra opts logger = do
           . setKeyspace (Keyspace opts.keyspace)
           . setProtocolVersion V4
           $ defSettings
-  initCassandra basicCasSettings opts.tlsCa logger
+  clientState <- initCassandra basicCasSettings opts.tlsCa logger
+  registerClientKeyspace clientState (Keyspace opts.keyspace)
+  pure clientState
 
 -- | Create Cassandra `ClientState` ("connection") for a service
 initCassandraForService ::
@@ -79,6 +116,7 @@ initCassandraForService opts serviceName discoUrl mbSchemaVersion logger = do
           . setPolicy (dcFilterPolicyIfConfigured logger opts.filterNodesByDatacentre)
           $ defSettings
   p <- initCassandra basicCasSettings opts.tlsCa logger
+  registerClientKeyspace p (Keyspace opts.keyspace)
   maybe (pure ()) (\v -> runClient p $ (versionCheck v)) mbSchemaVersion
   pure p
 

libs/wire-subsystems/src/Wire/ActivationCodeStore/Cassandra.hs

@@ -18,6 +18,7 @@
 module Wire.ActivationCodeStore.Cassandra (interpretActivationCodeStoreToCassandra) where
 
 import Cassandra
+import Cassandra.Util (requireClientKeyspace)
 import Data.Id
 import Data.Text (pack)
 import Data.Text.Ascii qualified as Ascii
@@ -33,40 +34,44 @@ import Wire.API.User.Activation
 import Wire.API.User.EmailAddress
 import Wire.ActivationCodeStore
 import Wire.UserKeyStore
+import Wire.Util (qualifiedTableName)
 
 interpretActivationCodeStoreToCassandra :: (Member (Embed IO) r) => ClientState -> InterpreterFor ActivationCodeStore r
 interpretActivationCodeStoreToCassandra casClient =
   interpret $
-    runEmbedded (runClient casClient) . embed . \case
+    \case
       LookupActivationCode ek -> do
-        liftIO (mkActivationKey ek)
-          >>= retry x1 . query1 cql . params LocalQuorum . Identity
-      NewActivationCode ek timeout uid -> newActivationCodeImpl ek timeout uid
+        keyspace <- embed @IO $ requireClientKeyspace casClient
+        runEmbedded (runClient casClient) . embed $ do
+          liftIO (mkActivationKey ek)
+            >>= retry x1 . query1 (lookupActivationCodeQuery keyspace) . params LocalQuorum . Identity
+      NewActivationCode ek timeout uid -> do
+        keyspace <- embed @IO $ requireClientKeyspace casClient
+        runEmbedded (runClient casClient) . embed $ newActivationCodeImpl keyspace ek timeout uid
   where
-    cql :: PrepQuery R (Identity ActivationKey) (Maybe UserId, ActivationCode)
-    cql =
-      [sql|
-      SELECT user, code FROM activation_keys WHERE key = ?
-      |]
+    lookupActivationCodeQuery :: Keyspace -> PrepQuery R (Identity ActivationKey) (Maybe UserId, ActivationCode)
+    lookupActivationCodeQuery keyspace =
+      fromString $ "SELECT user, code FROM " <> table keyspace "activation_keys" <> " WHERE key = ?"
 
 -- | Create a new pending activation for a given 'EmailKey'.
 newActivationCodeImpl ::
   (MonadClient m) =>
+  Keyspace ->
   EmailKey ->
   -- | The timeout for the activation code.
   Timeout ->
   -- | The user with whom to associate the activation code.
   Maybe UserId ->
   m Activation
-newActivationCodeImpl uk timeout u = do
+newActivationCodeImpl keyspace uk timeout u = do
   let typ = "email"
       key = fromEmail (emailKeyOrig uk)
   code <- liftIO $ genCode
   insert typ key code
   where
     insert t k c = do
       key <- liftIO $ mkActivationKey uk
-      retry x5 . write keyInsert $ params LocalQuorum (key, t, k, c, u, maxAttempts, round timeout)
+      retry x5 . write (keyInsert keyspace) $ params LocalQuorum (key, t, k, c, u, maxAttempts, round timeout)
       pure $ Activation key c
     genCode =
       ActivationCode . Ascii.unsafeFromText . pack . printf "%06d"
@@ -85,12 +90,16 @@ mkActivationKey k = do
       . T.encodeUtf8
       $ emailKeyUniq k
 
-keyInsert :: PrepQuery W (ActivationKey, Text, Text, ActivationCode, Maybe UserId, Int32, Int32) ()
-keyInsert =
-  "INSERT INTO activation_keys \
-  \(key, key_type, key_text, code, user, retries) VALUES \
-  \(?  , ?       , ?       , ?   , ?   , ?      ) USING TTL ?"
+keyInsert :: Keyspace -> PrepQuery W (ActivationKey, Text, Text, ActivationCode, Maybe UserId, Int32, Int32) ()
+keyInsert keyspace =
+  fromString $
+    "INSERT INTO "
+      <> table keyspace "activation_keys"
+      <> " (key, key_type, key_text, code, user, retries) VALUES (?  , ?       , ?       , ?   , ?   , ?      ) USING TTL ?"
 
 -- | Max. number of activation attempts per 'ActivationKey'.
 maxAttempts :: Int32
 maxAttempts = 3
+
+table :: Keyspace -> String -> String
+table = qualifiedTableName

libs/wire-subsystems/src/Wire/BlockListStore/Cassandra.hs

@@ -21,10 +21,12 @@ module Wire.BlockListStore.Cassandra
 where
 
 import Cassandra
+import Cassandra.Util (requireClientKeyspace)
 import Imports
 import Polysemy
 import Wire.BlockListStore (BlockListStore (..))
 import Wire.UserKeyStore
+import Wire.Util (qualifiedTableName)
 
 interpretBlockListStoreToCassandra ::
   forall r.
@@ -33,30 +35,39 @@ interpretBlockListStoreToCassandra ::
   InterpreterFor BlockListStore r
 interpretBlockListStoreToCassandra casClient =
   interpret $
-    embed @IO . runClient casClient . \case
-      Insert uk -> insert uk
-      Exists uk -> exists uk
-      Delete uk -> delete uk
+    \case
+      Insert uk -> do
+        keyspace <- embed @IO $ requireClientKeyspace casClient
+        embed @IO $ runClient casClient (insert keyspace uk)
+      Exists uk -> do
+        keyspace <- embed @IO $ requireClientKeyspace casClient
+        embed @IO $ runClient casClient (exists keyspace uk)
+      Delete uk -> do
+        keyspace <- embed @IO $ requireClientKeyspace casClient
+        embed @IO $ runClient casClient (delete keyspace uk)
 
 --------------------------------------------------------------------------------
 -- UserKey block listing
 
-insert :: (MonadClient m) => EmailKey -> m ()
-insert uk = retry x5 $ write keyInsert (params LocalQuorum (Identity $ emailKeyUniq uk))
+insert :: (MonadClient m) => Keyspace -> EmailKey -> m ()
+insert keyspace uk = retry x5 $ write (keyInsert keyspace) (params LocalQuorum (Identity $ emailKeyUniq uk))
 
-exists :: (MonadClient m) => EmailKey -> m Bool
-exists uk =
+exists :: (MonadClient m) => Keyspace -> EmailKey -> m Bool
+exists keyspace uk =
   (pure . isJust) . fmap runIdentity
-    =<< retry x1 (query1 keySelect (params LocalQuorum (Identity $ emailKeyUniq uk)))
+    =<< retry x1 (query1 (keySelect keyspace) (params LocalQuorum (Identity $ emailKeyUniq uk)))
 
-delete :: (MonadClient m) => EmailKey -> m ()
-delete uk = retry x5 $ write keyDelete (params LocalQuorum (Identity $ emailKeyUniq uk))
+delete :: (MonadClient m) => Keyspace -> EmailKey -> m ()
+delete keyspace uk = retry x5 $ write (keyDelete keyspace) (params LocalQuorum (Identity $ emailKeyUniq uk))
 
-keyInsert :: PrepQuery W (Identity Text) ()
-keyInsert = "INSERT INTO blacklist (key) VALUES (?)"
+keyInsert :: Keyspace -> PrepQuery W (Identity Text) ()
+keyInsert keyspace = fromString $ "INSERT INTO " <> table keyspace "blacklist" <> " (key) VALUES (?)"
 
-keySelect :: PrepQuery R (Identity Text) (Identity Text)
-keySelect = "SELECT key FROM blacklist WHERE key = ?"
+keySelect :: Keyspace -> PrepQuery R (Identity Text) (Identity Text)
+keySelect keyspace = fromString $ "SELECT key FROM " <> table keyspace "blacklist" <> " WHERE key = ?"
 
-keyDelete :: PrepQuery W (Identity Text) ()
-keyDelete = "DELETE FROM blacklist WHERE key = ?"
+keyDelete :: Keyspace -> PrepQuery W (Identity Text) ()
+keyDelete keyspace = fromString $ "DELETE FROM " <> table keyspace "blacklist" <> " WHERE key = ?"
+
+table :: Keyspace -> String -> String
+table = qualifiedTableName