refactor types for DB entity (#1548)

Author: epoberezkin

simplexmq.cabal

@@ -135,6 +135,7 @@ library
       Simplex.Messaging.Server.QueueStore.QueueInfo
       Simplex.Messaging.ServiceScheme
       Simplex.Messaging.Session
+      Simplex.Messaging.Agent.Store.Entity
       Simplex.Messaging.TMap
       Simplex.Messaging.Transport
       Simplex.Messaging.Transport.Buffer
@@ -308,6 +309,7 @@ library
     , network-transport ==0.5.6
     , network-udp ==0.0.*
     , random >=1.1 && <1.3
+    , scientific ==0.3.7.*
     , simple-logger ==0.1.*
     , socks ==0.6.*
     , stm ==2.5.*

src/Simplex/Messaging/Agent.hs

@@ -216,6 +216,7 @@ import Simplex.Messaging.Protocol
   )
 import qualified Simplex.Messaging.Protocol as SMP
 import Simplex.Messaging.ServiceScheme (ServiceScheme (..))
+import Simplex.Messaging.Agent.Store.Entity
 import qualified Simplex.Messaging.TMap as TM
 import Simplex.Messaging.Transport (SMPVersion)
 import Simplex.Messaging.Util
@@ -833,7 +834,7 @@ newConn c userId enableNtfs cMode userData_ clientData pqInitKeys subMode = do
     `catchE` \e -> withStore' c (`deleteConnRecord` connId) >> throwE e
 
 setContactShortLink' :: AgentClient -> ConnId -> ConnInfo -> Maybe CRClientData -> AM (ConnShortLink 'CMContact)
-setContactShortLink' c connId userData clientData = 
+setContactShortLink' c connId userData clientData =
   withConnLock c connId "setContactShortLink" $
     withStore c (`getConn` connId) >>= \case
       SomeConn _ (ContactConnection _ rq) -> do
@@ -934,7 +935,7 @@ newRcvConnSrv c userId connId enableNtfs cMode userData_ clientData pqInitKeys s
     createRcvQueue nonce_ qd e2eKeys = do
       AgentConfig {smpClientVRange = vr} <- asks config
       -- TODO [notifications] send correct NTF credentials here
-      -- let ntfCreds_ = Nothing 
+      -- let ntfCreds_ = Nothing
       (rq, qUri, tSess, sessId) <- newRcvQueue_ c userId connId srvWithAuth vr qd subMode nonce_ e2eKeys `catchAgentError` \e -> liftIO (print e) >> throwE e
       atomically $ incSMPServerStat c userId srv connCreated
       rq' <- withStore c $ \db -> updateNewConnRcv db connId rq
@@ -1122,7 +1123,7 @@ joinConnSrv c userId connId enableNtfs cReqUri@CRContactUri {} cInfo pqSup subMo
     Nothing -> throwE $ AGENT A_VERSION
 
 delInvSL :: AgentClient -> ConnId -> SMPServerWithAuth -> SMP.LinkId -> AM ()
-delInvSL c connId srv lnkId = 
+delInvSL c connId srv lnkId =
   withStore' c (\db -> deleteInvShortLink db (protoServer srv) lnkId) `catchE` \e ->
     liftIO $ nonBlockingWriteTBQueue (subQ c) ("", connId, AEvt SAEConn (ERR $ INTERNAL $ "error deleting short link " <> show e))
 
@@ -1293,7 +1294,7 @@ getConnectionMessages' c = mapM $ tryAgentError' . getConnectionMessage
       msg_ <- getQueueMessage c rq `catchAgentError` \e -> atomically (releaseGetLock c rq) >> throwError e
       when (isNothing msg_) $ do
         atomically $ releaseGetLock c rq
-        forM_ msgTs_ $ \msgTs -> withStore' c $ \db -> setLastBrokerTs db connId (DBQueueId dbQueueId) msgTs
+        forM_ msgTs_ $ \msgTs -> withStore' c $ \db -> setLastBrokerTs db connId (DBEntityId dbQueueId) msgTs
       pure msg_
 {-# INLINE getConnectionMessages' #-}
 
@@ -1910,7 +1911,7 @@ switchConnection' c connId =
       _ -> throwE $ CMD PROHIBITED "switchConnection: not duplex"
 
 switchDuplexConnection :: AgentClient -> Connection 'CDuplex -> RcvQueue -> AM ConnectionStats
-switchDuplexConnection c (DuplexConnection cData@ConnData {connId, userId} rqs sqs) rq@RcvQueue {server, dbQueueId = DBQueueId dbQueueId, sndId} = do
+switchDuplexConnection c (DuplexConnection cData@ConnData {connId, userId} rqs sqs) rq@RcvQueue {server, dbQueueId = DBEntityId dbQueueId, sndId} = do
   checkRQSwchStatus rq RSSwitchStarted
   clientVRange <- asks $ smpClientVRange . config
   -- try to get the server that is different from all queues, or at least from the primary rcv queue
@@ -2940,7 +2941,7 @@ processSMPTransmissions c@AgentClient {subQ} (tSess@(userId, srv, _), _v, sessId
               Just qInfo@(Compatible sqInfo@SMPQueueInfo {queueAddress}) ->
                 case (findQ (qAddress sqInfo) sqs, findQ addr sqs) of
                   (Just _, _) -> qError "QADD: queue address is already used in connection"
-                  (_, Just sq@SndQueue {dbQueueId = DBQueueId dbQueueId}) -> do
+                  (_, Just sq@SndQueue {dbQueueId = DBEntityId dbQueueId}) -> do
                     let (delSqs, keepSqs) = L.partition ((Just dbQueueId ==) . dbReplaceQId) sqs
                     case L.nonEmpty keepSqs of
                       Just sqs' -> do
@@ -3278,7 +3279,7 @@ newSndQueue userId connId (Compatible (SMPQueueInfo smpClientVersion SMPQueueAdd
             e2ePubKey = Just e2ePubKey,
             -- setting status to Secured prevents SKEY when queue was already secured with LKEY
             status = if isJust sndKeys_ then Secured else New,
-            dbQueueId = DBNewQueue,
+            dbQueueId = DBNewEntity,
             primary = True,
             dbReplaceQueueId = Nothing,
             sndSwchStatus = Nothing,

src/Simplex/Messaging/Agent/Client.hs

@@ -278,6 +278,7 @@ import Simplex.Messaging.Protocol
 import qualified Simplex.Messaging.Protocol as SMP
 import Simplex.Messaging.Server.QueueStore.QueueInfo
 import Simplex.Messaging.Session
+import Simplex.Messaging.Agent.Store.Entity
 import Simplex.Messaging.TMap (TMap)
 import qualified Simplex.Messaging.TMap as TM
 import Simplex.Messaging.Transport (SMPVersion, SessionId, THandleParams (sessionId, thVersion), TransportError (..), TransportPeer (..), sndAuthKeySMPVersion, shortLinksSMPVersion)
@@ -1083,7 +1084,7 @@ sendOrProxySMPCommand ::
   UserId ->
   SMPServer ->
   ConnId -> -- session entity ID, for short links LinkId is used
-  ByteString -> 
+  ByteString ->
   SMP.EntityId -> -- sender or link ID
   (SMPClient -> ProxiedRelay -> ExceptT SMPClientError IO (Either ProxyClientError a)) ->
   (SMPClient -> ExceptT SMPClientError IO a) ->
@@ -1395,7 +1396,7 @@ newRcvQueue_ c userId connId (ProtoServerWithAuth srv auth) vRange cqrd subMode
             queueMode,
             shortLink,
             status = New,
-            dbQueueId = DBNewQueue,
+            dbQueueId = DBNewEntity,
             primary = True,
             dbReplaceQueueId = Nothing,
             rcvSwchStatus = Nothing,
@@ -1408,7 +1409,7 @@ newRcvQueue_ c userId connId (ProtoServerWithAuth srv auth) vRange cqrd subMode
   where
     mkShortLinkCreds :: (THandleParams SMPVersion 'TClient, QueueIdsKeys) -> AM (Maybe ShortLinkCreds)
     mkShortLinkCreds (thParams', QIK {sndId, queueMode, linkId}) = case (cqrd, queueMode) of
-      (CQRMessaging ld, Just QMMessaging) -> 
+      (CQRMessaging ld, Just QMMessaging) ->
         withLinkData ld $ \lnkId CQRData {linkKey, privSigKey, srvReq = (sndId', d)} ->
           if sndId == sndId'
             then pure $ Just $ ShortLinkCreds lnkId linkKey privSigKey (fst d)

src/Simplex/Messaging/Agent/Store.hs

@@ -52,30 +52,19 @@ import Simplex.Messaging.Protocol
     VersionSMPC,
   )
 import qualified Simplex.Messaging.Protocol as SMP
+import Simplex.Messaging.Agent.Store.Entity
 
 createStore :: DBOpts -> MigrationConfirmation -> IO (Either MigrationError DBStore)
 createStore dbOpts = createDBStore dbOpts appMigrations
 
 -- * Queue types
 
-data QueueStored = QSStored | QSNew
+type RcvQueue = StoredRcvQueue 'DBStored
 
-data SQueueStored (q :: QueueStored) where
-  SQSStored :: SQueueStored 'QSStored
-  SQSNew :: SQueueStored 'QSNew
-
-data DBQueueId (q :: QueueStored) where
-  DBQueueId :: Int64 -> DBQueueId 'QSStored
-  DBNewQueue :: DBQueueId 'QSNew
-
-deriving instance Show (DBQueueId q)
-
-type RcvQueue = StoredRcvQueue 'QSStored
-
-type NewRcvQueue = StoredRcvQueue 'QSNew
+type NewRcvQueue = StoredRcvQueue 'DBNew
 
 -- | A receive queue. SMP queue through which the agent receives messages from a sender.
-data StoredRcvQueue (q :: QueueStored) = RcvQueue
+data StoredRcvQueue (q :: DBStored) = RcvQueue
   { userId :: UserId,
     connId :: ConnId,
     server :: SMPServer,
@@ -98,7 +87,7 @@ data StoredRcvQueue (q :: QueueStored) = RcvQueue
     -- | queue status
     status :: QueueStatus,
     -- | database queue ID (within connection)
-    dbQueueId :: DBQueueId q,
+    dbQueueId :: DBEntityId' q,
     -- | True for a primary or a next primary queue of the connection (next if dbReplaceQueueId is set)
     primary :: Bool,
     -- | database queue ID to replace, Nothing if this queue is not replacing another, `Just Nothing` is used for replacing old queues
@@ -160,12 +149,12 @@ data InvShortLink = InvShortLink
   }
   deriving (Show)
 
-type SndQueue = StoredSndQueue 'QSStored
+type SndQueue = StoredSndQueue 'DBStored
 
-type NewSndQueue = StoredSndQueue 'QSNew
+type NewSndQueue = StoredSndQueue 'DBNew
 
 -- | A send queue. SMP queue through which the agent sends messages to a recipient.
-data StoredSndQueue (q :: QueueStored) = SndQueue
+data StoredSndQueue (q :: DBStored) = SndQueue
   { userId :: UserId,
     connId :: ConnId,
     server :: SMPServer,
@@ -184,7 +173,7 @@ data StoredSndQueue (q :: QueueStored) = SndQueue
     -- | queue status
     status :: QueueStatus,
     -- | database queue ID (within connection)
-    dbQueueId :: DBQueueId q,
+    dbQueueId :: DBEntityId' q,
     -- | True for a primary or a next primary queue of the connection (next if dbReplaceQueueId is set)
     primary :: Bool,
     -- | ID of the queue this one is replacing
@@ -257,7 +246,7 @@ instance SMPQueueRec RcvQueue where
   {-# INLINE qUserId #-}
   qConnId RcvQueue {connId} = connId
   {-# INLINE qConnId #-}
-  dbQId RcvQueue {dbQueueId = DBQueueId qId} = qId
+  dbQId RcvQueue {dbQueueId = DBEntityId qId} = qId
   {-# INLINE dbQId #-}
   dbReplaceQId RcvQueue {dbReplaceQueueId} = dbReplaceQueueId
   {-# INLINE dbReplaceQId #-}
@@ -267,7 +256,7 @@ instance SMPQueueRec SndQueue where
   {-# INLINE qUserId #-}
   qConnId SndQueue {connId} = connId
   {-# INLINE qConnId #-}
-  dbQId SndQueue {dbQueueId = DBQueueId qId} = qId
+  dbQId SndQueue {dbQueueId = DBEntityId qId} = qId
   {-# INLINE dbQId #-}
   dbReplaceQId SndQueue {dbReplaceQueueId} = dbReplaceQueueId
   {-# INLINE dbReplaceQId #-}

src/Simplex/Messaging/Agent/Store/AgentStore.hs

@@ -283,6 +283,7 @@ import Simplex.Messaging.Notifications.Types
 import Simplex.Messaging.Parsers (parseAll)
 import Simplex.Messaging.Protocol
 import qualified Simplex.Messaging.Protocol as SMP
+import Simplex.Messaging.Agent.Store.Entity
 import Simplex.Messaging.Transport.Client (TransportHost)
 import Simplex.Messaging.Util (bshow, catchAllErrors, eitherToMaybe, firstRow, firstRow', ifM, maybeFirstRow, tshow, ($>>=), (<$$>))
 import Simplex.Messaging.Version.Internal
@@ -858,7 +859,7 @@ createRcvMsg db connId rq@RcvQueue {dbQueueId} rcvMsgData@RcvMsgData {msgMeta =
   updateRcvMsgHash db connId sndMsgId internalRcvId internalHash
   setLastBrokerTs db connId dbQueueId brokerTs
 
-setLastBrokerTs :: DB.Connection -> ConnId -> DBQueueId 'QSStored -> UTCTime -> IO ()
+setLastBrokerTs :: DB.Connection -> ConnId -> DBEntityId -> UTCTime -> IO ()
 setLastBrokerTs db connId dbQueueId brokerTs =
   DB.execute db "UPDATE rcv_queues SET last_broker_ts = ? WHERE conn_id = ? AND rcv_queue_id = ? AND (last_broker_ts IS NULL OR last_broker_ts < ?)" (brokerTs, connId, dbQueueId, brokerTs)
 
@@ -1212,7 +1213,7 @@ getSndRatchet db connId v =
     DB.query db "SELECT ratchet_state, x3dh_pub_key_1, x3dh_pub_key_2, pq_pub_kem FROM ratchets WHERE conn_id = ?" (Only connId)
   where
     result = \case
-      (Just ratchetState, Just k1, Just k2, pKem_) -> 
+      (Just ratchetState, Just k1, Just k2, pKem_) ->
         let params = case pKem_ of
               Nothing -> CR.AE2ERatchetParams CR.SRKSProposed (CR.E2ERatchetParams v k1 k2 Nothing)
               Just (CR.ARKP s pKem) -> CR.AE2ERatchetParams s (CR.E2ERatchetParams v k1 k2 (Just pKem))
@@ -1811,15 +1812,6 @@ instance ToField QueueStatus where toField = toField . serializeQueueStatus
 
 instance FromField QueueStatus where fromField = fromTextField_ queueStatusT
 
-instance ToField (DBQueueId 'QSStored) where toField (DBQueueId qId) = toField qId
-
-instance FromField (DBQueueId 'QSStored) where 
-#if defined(dbPostgres)
-  fromField x dat = DBQueueId <$> fromField x dat
-#else
-  fromField x = DBQueueId <$> fromField x
-#endif
-
 instance ToField InternalRcvId where toField (InternalRcvId x) = toField x
 
 deriving newtype instance FromField InternalRcvId
@@ -2018,13 +2010,13 @@ insertSndQueue_ db connId' sq@SndQueue {..} serverKeyHash_ = do
         smp_client_version=EXCLUDED.smp_client_version,
         server_key_hash=EXCLUDED.server_key_hash
     |]
-    ((host server, port server, sndId, queueMode, connId', sndPublicKey, sndPrivateKey, e2ePubKey, e2eDhSecret) 
+    ((host server, port server, sndId, queueMode, connId', sndPublicKey, sndPrivateKey, e2ePubKey, e2eDhSecret)
     :. (status, qId, BI primary, dbReplaceQueueId, smpClientVersion, serverKeyHash_))
   pure (sq :: NewSndQueue) {connId = connId', dbQueueId = qId}
 
-newQueueId_ :: [Only Int64] -> DBQueueId 'QSStored
-newQueueId_ [] = DBQueueId 1
-newQueueId_ (Only maxId : _) = DBQueueId (maxId + 1)
+newQueueId_ :: [Only Int64] -> DBEntityId
+newQueueId_ [] = DBEntityId 1
+newQueueId_ (Only maxId : _) = DBEntityId (maxId + 1)
 
 -- * getConn helpers
 
@@ -2160,7 +2152,7 @@ rcvQueueQuery =
 
 toRcvQueue ::
   (UserId, C.KeyHash, ConnId, NonEmpty TransportHost, ServiceName, SMP.RecipientId, SMP.RcvPrivateAuthKey, SMP.RcvDhSecret, C.PrivateKeyX25519, Maybe C.DhSecretX25519, SMP.SenderId, Maybe QueueMode)
-    :. (QueueStatus, DBQueueId 'QSStored, BoolInt, Maybe Int64, Maybe RcvSwitchStatus, Maybe VersionSMPC, Int)
+    :. (QueueStatus, DBEntityId, BoolInt, Maybe Int64, Maybe RcvSwitchStatus, Maybe VersionSMPC, Int)
     :. (Maybe SMP.NtfPublicAuthKey, Maybe SMP.NtfPrivateAuthKey, Maybe SMP.NotifierId, Maybe RcvNtfDhSecret)
     :. (Maybe SMP.LinkId, Maybe LinkKey, Maybe C.PrivateKeyEd25519, Maybe EncDataBytes) ->
   RcvQueue
@@ -2210,7 +2202,7 @@ sndQueueQuery =
 toSndQueue ::
   (UserId, C.KeyHash, ConnId, NonEmpty TransportHost, ServiceName, SenderId, Maybe QueueMode)
     :. (Maybe SndPublicAuthKey, SndPrivateAuthKey, Maybe C.PublicKeyX25519, C.DhSecretX25519, QueueStatus)
-    :. (DBQueueId 'QSStored, BoolInt, Maybe Int64, Maybe SndSwitchStatus, VersionSMPC) ->
+    :. (DBEntityId, BoolInt, Maybe Int64, Maybe SndSwitchStatus, VersionSMPC) ->
   SndQueue
 toSndQueue
   ( (userId, keyHash, connId, host, port, sndId, queueMode)

src/Simplex/Messaging/Agent/Store/Entity.hs

@@ -0,0 +1,72 @@
+{-# LANGUAGE CPP #-}
+{-# LANGUAGE DataKinds #-}
+{-# LANGUAGE FlexibleInstances #-}
+{-# LANGUAGE GADTs #-}
+{-# LANGUAGE KindSignatures #-}
+{-# LANGUAGE LambdaCase #-}
+{-# LANGUAGE ScopedTypeVariables #-}
+{-# LANGUAGE StandaloneDeriving #-}
+{-# LANGUAGE TypeApplications #-}
+
+module Simplex.Messaging.Agent.Store.Entity where
+
+import Data.Aeson (FromJSON (..), ToJSON (..))
+import qualified Data.Aeson as J
+import qualified Data.Aeson.Encoding as JE
+import Data.Int (Int64)
+import Data.Scientific (floatingOrInteger)
+import Simplex.Messaging.Agent.Store.DB (FromField (..), ToField (..))
+
+data DBStored = DBStored | DBNew
+
+data SDBStored (s :: DBStored) where
+  SDBStored :: SDBStored 'DBStored
+  SDBNew :: SDBStored 'DBNew
+
+deriving instance Show (SDBStored s)
+
+class DBStoredI s where sdbStored :: SDBStored s
+
+instance DBStoredI 'DBStored where sdbStored = SDBStored
+
+instance DBStoredI 'DBNew where sdbStored = SDBNew
+
+data DBEntityId' (s :: DBStored) where
+  DBEntityId :: Int64 -> DBEntityId' 'DBStored
+  DBNewEntity :: DBEntityId' 'DBNew
+
+deriving instance Show (DBEntityId' s)
+
+deriving instance Eq (DBEntityId' s)
+
+type DBEntityId = DBEntityId' 'DBStored
+
+type DBNewEntity = DBEntityId' 'DBNew
+
+instance ToJSON (DBEntityId' s) where
+  toEncoding = \case
+    DBEntityId i -> toEncoding i
+    DBNewEntity -> JE.null_
+  toJSON = \case
+    DBEntityId i -> toJSON i
+    DBNewEntity -> J.Null
+
+instance DBStoredI s => FromJSON (DBEntityId' s) where
+  parseJSON v = case (v, sdbStored @s) of
+    (J.Null, SDBNew) -> pure DBNewEntity
+    (J.Number n, SDBStored) -> case floatingOrInteger n of
+      Left (_ :: Double) -> fail "bad DBEntityId"
+      Right i -> pure $ DBEntityId (fromInteger i)
+    _ -> fail "bad DBEntityId"
+  omittedField = case sdbStored @s of
+    SDBStored -> Nothing
+    SDBNew -> Just DBNewEntity
+
+instance FromField DBEntityId where
+#if defined(dbPostgres)
+  fromField x dat = DBEntityId <$> fromField x dat
+#else
+  fromField x = DBEntityId <$> fromField x
+#endif
+
+instance ToField DBEntityId where toField (DBEntityId i) = toField i