Merge pull request #33 from unisoncomputing/syncv2/ensure-serialized

Ensure Serialized Entities

Author: ChrisPenner

src/Share/BackgroundJobs/SerializedEntitiesMigration/Worker.hs

@@ -9,9 +9,11 @@ import Share.Codebase qualified as Codebase
 import Share.Codebase.Types (CodebaseEnv (..))
 import Share.Postgres
 import Share.Postgres qualified as PG
+import Share.Postgres.Causal.Queries qualified as CQ
 import Share.Postgres.Definitions.Queries qualified as DefnQ
 import Share.Postgres.Hashes.Queries qualified as HQ
 import Share.Postgres.IDs
+import Share.Postgres.Patches.Queries qualified as PQ
 import Share.Postgres.Sync.Queries qualified as SQ
 import Share.Prelude
 import Share.Web.Authorization qualified as AuthZ
@@ -81,12 +83,18 @@ processComponents !_authZReceipt = do
                 |]
           pure True
 
-saveUnsandboxedSerializedEntities :: (QueryM m) => Hash32 -> TempEntity -> m ()
+saveUnsandboxedSerializedEntities :: Hash32 -> TempEntity -> Codebase.CodebaseM e ()
 saveUnsandboxedSerializedEntities hash entity = do
   let serialised = SyncV2.serialiseCBORBytes entity
   case entity of
     Entity.TC {} -> error "Unexpected term component"
     Entity.DC {} -> error "Unexpected decl component"
-    Entity.P {} -> SQ.saveSerializedPatch hash serialised
-    Entity.C {} -> SQ.saveSerializedCausal hash serialised
-    Entity.N {} -> SQ.saveSerializedNamespace hash serialised
+    Entity.P {} -> do
+      patchId <- HQ.expectPatchIdsOf id (fromHash32 @PatchHash hash)
+      PQ.saveSerializedPatch patchId serialised
+    Entity.C {} -> do
+      cId <- CQ.expectCausalIdByHash (fromHash32 @CausalHash hash)
+      CQ.saveSerializedCausal cId serialised
+    Entity.N {} -> do
+      bhId <- HQ.expectBranchHashId (fromHash32 @BranchHash hash)
+      CQ.saveSerializedNamespace bhId serialised

src/Share/Codebase.hs

@@ -468,7 +468,7 @@ squashCausal Causal.Causal {valueHash = unsquashedBranchHash, value} = do
       let squashedBranchHead = branch {V2.children = snd <$> squashedChildren}
       (squashedBranchHashId, squashedBranchHash) <- CausalQ.saveV2BranchShallow squashedBranchHead
       let ancestors = mempty
-      (squashedCausalId, squashedCausalHash) <- CausalQ.saveCausal Nothing squashedBranchHashId ancestors
+      (squashedCausalId, squashedCausalHash) <- CausalQ.saveCausal Nothing Nothing squashedBranchHashId ancestors
       let squashedCausalBranch =
             Causal.Causal
               { causalHash = squashedCausalHash,

src/Share/Postgres/Causal/Queries.hs

@@ -25,38 +25,61 @@ module Share.Postgres.Causal.Queries
     hashCausal,
     bestCommonAncestor,
     isFastForward,
+
+    -- * Sync
+    expectCausalEntity,
+    expectNamespaceEntity,
+
+    -- * For migrations, can probably remove this export later.
+    saveSerializedCausal,
+    saveSerializedNamespace,
   )
 where
 
 import Control.Lens
+import Data.ByteString.Lazy.Char8 qualified as BL
 import Data.Map qualified as Map
 import Data.Set qualified as Set
+import Data.Vector qualified as Vector
 import Share.Codebase.Types (CodebaseM)
 import Share.Codebase.Types qualified as Codebase
 import Share.IDs (UserId)
 import Share.Postgres
 import Share.Postgres.Causal.Types
 import Share.Postgres.Definitions.Queries qualified as Defn
+import Share.Postgres.Definitions.Queries qualified as DefnQ
 import Share.Postgres.Definitions.Types
 import Share.Postgres.Hashes.Queries qualified as HashQ
 import Share.Postgres.IDs
 import Share.Postgres.Patches.Queries qualified as PatchQ
+import Share.Postgres.Serialization qualified as S
+import Share.Postgres.Sync.Conversions qualified as Cv
 import Share.Prelude
 import Share.Utils.Postgres (OrdBy, ordered)
 import Share.Web.Errors (MissingExpectedEntity (MissingExpectedEntity))
 import U.Codebase.Branch hiding (NamespaceStats, nonEmptyChildren)
 import U.Codebase.Branch qualified as V2 hiding (NamespaceStats)
 import U.Codebase.Causal qualified as Causal
+import U.Codebase.Causal qualified as U
 import U.Codebase.Reference
 import U.Codebase.Referent
 import U.Codebase.Referent qualified as Referent
+import U.Codebase.Sqlite.Branch.Format (LocalBranchBytes (..))
+import U.Codebase.Sqlite.Branch.Format qualified as BranchFormat
 import U.Codebase.Sqlite.Branch.Full qualified as BranchFull
+import U.Codebase.Sqlite.LocalizeObject qualified as Localize
+import U.Codebase.Sqlite.TempEntity (TempEntity)
 import Unison.Codebase.Path qualified as Path
 import Unison.Hash (Hash)
 import Unison.Hash32 (Hash32)
+import Unison.Hash32 qualified as Hash32
 import Unison.Hashing.V2 qualified as H
 import Unison.NameSegment.Internal as NameSegment
 import Unison.Reference qualified as Reference
+import Unison.Sync.Common qualified as SyncCommon
+import Unison.Sync.Types qualified as Sync
+import Unison.SyncV2.Types (CBORBytes (..))
+import Unison.SyncV2.Types qualified as SyncV2
 import Unison.Util.Map qualified as Map
 
 expectCausalNamespace :: (HasCallStack, QueryM m) => CausalId -> m (CausalNamespace m)
@@ -407,18 +430,22 @@ loadCausalNamespaceAtPath causalId path = do
 -- | Given a namespace whose dependencies have all been pre-saved, save it to the database under the given hash.
 savePgNamespace ::
   (HasCallStack) =>
+  -- | The pre-serialized namespace, if available. If Nothing it will be re-generated, which is slower.
+  Maybe TempEntity ->
   -- Normally we'd prefer to always hash it ourselves, but there are some bad hashes in the wild
   -- that we need to support saving, if we're passed a hash to save a branch at we will save
   -- it at that hash regardless of what the _actual_ hash is.
   Maybe BranchHash ->
   PgNamespace ->
   CodebaseM e (BranchHashId, BranchHash)
-savePgNamespace mayBh b@(BranchFull.Branch {terms, types, patches, children}) = do
+savePgNamespace maySerialized mayBh b@(BranchFull.Branch {terms, types, patches, children}) = do
   codebaseOwnerUserId <- asks Codebase.codebaseOwner
   bh <- whenNothing mayBh $ hashPgNamespace b
   bhId <- HashQ.ensureBranchHashId bh
   queryExpect1Col [sql| SELECT EXISTS (SELECT FROM namespaces WHERE namespace_hash_id = #{bhId}) |] >>= \case
-    False -> doSave bhId
+    False -> do
+      doSave bhId
+      doSaveSerialized bhId
     True -> pure ()
   execute_
     [sql| INSERT INTO namespace_ownership (namespace_hash_id, user_id)
@@ -427,6 +454,14 @@ savePgNamespace mayBh b@(BranchFull.Branch {terms, types, patches, children}) =
     |]
   pure (bhId, bh)
   where
+    doSaveSerialized :: BranchHashId -> CodebaseM e ()
+    doSaveSerialized bhId = do
+      nsEntity <- case maySerialized of
+        Just serialized -> pure serialized
+        Nothing -> SyncCommon.entityToTempEntity id . Sync.N <$> expectNamespaceEntity bhId
+      let serializedNamespace = SyncV2.serialiseCBORBytes nsEntity
+      saveSerializedNamespace bhId serializedNamespace
+
     doSave :: BranchHashId -> CodebaseM e ()
     doSave bhId = do
       -- Expand all term mappings into a list
@@ -608,6 +643,34 @@ savePgNamespace mayBh b@(BranchFull.Branch {terms, types, patches, children}) =
       -- Note: this must be run AFTER inserting the namespace and all its children.
       execute_ [sql| SELECT save_namespace(#{bhId}) |]
 
+saveSerializedNamespace :: (QueryM m) => BranchHashId -> CBORBytes TempEntity -> m ()
+saveSerializedNamespace bhId (CBORBytes bytes) = do
+  bytesId <- DefnQ.ensureBytesIdsOf id (BL.toStrict bytes)
+  execute_
+    [sql|
+      INSERT INTO serialized_namespaces (namespace_hash_id, bytes_id)
+        VALUES (#{bhId}, #{bytesId})
+      ON CONFLICT DO NOTHING
+    |]
+
+expectNamespaceEntity :: BranchHashId -> CodebaseM e (Sync.Namespace Text Hash32)
+expectNamespaceEntity bhId = do
+  v2Branch <- expectNamespace bhId
+  second Hash32.fromHash <$> branchToEntity v2Branch
+  where
+    branchToEntity branch = do
+      branchFull <- Cv.branchV2ToBF branch
+      let (BranchFormat.LocalIds {branchTextLookup, branchDefnLookup, branchPatchLookup, branchChildLookup}, localBranch) = Localize.localizeBranchG branchFull
+      let bytes = LocalBranchBytes $ S.encodeNamespace localBranch
+      pure $
+        Sync.Namespace
+          { textLookup = Vector.toList branchTextLookup,
+            defnLookup = Vector.toList branchDefnLookup,
+            patchLookup = Vector.toList branchPatchLookup,
+            childLookup = Vector.toList branchChildLookup,
+            bytes = bytes
+          }
+
 -- | Hash a namespace into a BranchHash
 hashPgNamespace :: forall m. (QueryM m) => PgNamespace -> m BranchHash
 hashPgNamespace b = do
@@ -671,14 +734,23 @@ hashCausal branchHashId ancestorIds = do
   let hCausal = H.Causal {branchHash = unBranchHash branchHash, parents = ancestors}
   pure . CausalHash . H.contentHash $ hCausal
 
-saveCausal :: Maybe CausalHash -> BranchHashId -> Set CausalId -> CodebaseM e (CausalId, CausalHash)
-saveCausal mayCh bhId ancestorIds = do
+saveCausal ::
+  -- | The pre-serialized causal, if available. If Nothing it will be re-generated, which is slower.
+  Maybe TempEntity ->
+  Maybe CausalHash ->
+  BranchHashId ->
+  Set CausalId ->
+  CodebaseM e (CausalId, CausalHash)
+saveCausal maySerializedCausal mayCh bhId ancestorIds = do
   ch <- maybe (hashCausal bhId ancestorIds) pure mayCh
   codebaseOwnerUserId <- asks Codebase.codebaseOwner
   cId <-
     query1Col [sql| SELECT id FROM causals WHERE hash = #{ch} |] >>= \case
       Just cId -> pure cId
-      Nothing -> doSave ch
+      Nothing -> do
+        cId <- doSave ch
+        doSaveSerialized cId
+        pure cId
   execute_
     [sql|
     INSERT INTO causal_ownership (user_id, causal_id)
@@ -687,6 +759,14 @@ saveCausal mayCh bhId ancestorIds = do
     |]
   pure (cId, ch)
   where
+    doSaveSerialized cId = do
+      causalEntity <- case maySerializedCausal of
+        Just serializedCausal -> pure serializedCausal
+        Nothing -> do
+          SyncCommon.entityToTempEntity id . Sync.C <$> expectCausalEntity cId
+      let serializedCausal = SyncV2.serialiseCBORBytes causalEntity
+      saveSerializedCausal cId serializedCausal
+
     doSave ch = do
       cId <-
         queryExpect1Col
@@ -707,6 +787,26 @@ saveCausal mayCh bhId ancestorIds = do
       |]
       pure cId
 
+saveSerializedCausal :: (QueryM m) => CausalId -> CBORBytes TempEntity -> m ()
+saveSerializedCausal causalId (CBORBytes bytes) = do
+  bytesId <- DefnQ.ensureBytesIdsOf id (BL.toStrict bytes)
+  execute_
+    [sql|
+      INSERT INTO serialized_causals (causal_id, bytes_id)
+        VALUES (#{causalId}, #{bytesId})
+      ON CONFLICT DO NOTHING
+    |]
+
+expectCausalEntity :: (HasCallStack) => CausalId -> CodebaseM e (Sync.Causal Hash32)
+expectCausalEntity causalId = do
+  U.Causal {valueHash, parents} <- expectCausalNamespace causalId
+  pure $
+    ( Sync.Causal
+        { namespaceHash = Hash32.fromHash $ unBranchHash valueHash,
+          parents = Set.map (Hash32.fromHash . unCausalHash) . Map.keysSet $ parents
+        }
+    )
+
 -- | Get the ref to the result of squashing if we've squashed that ref in the past.
 -- Also adds the squash result to current codebase if we find it.
 tryGetCachedSquashResult :: BranchHash -> CodebaseM e (Maybe CausalId)
@@ -744,7 +844,7 @@ saveSquashResult unsquashedBranchHash squashedCausalHashId = do
 saveV2BranchShallow :: V2.Branch (CodebaseM e) -> CodebaseM e (BranchHashId, BranchHash)
 saveV2BranchShallow v2Branch = do
   pgNamespace <- expectV2BranchDependencies v2Branch
-  savePgNamespace Nothing pgNamespace
+  savePgNamespace Nothing Nothing pgNamespace
   where
     expectV2BranchDependencies :: V2.Branch (CodebaseM e) -> CodebaseM e PgNamespace
     expectV2BranchDependencies V2.Branch {terms, types, patches, children} = do

src/Share/Postgres/Definitions/Queries.hs

@@ -28,11 +28,15 @@ module Share.Postgres.Definitions.Queries
     ensureBytesIdsOf,
     expectTextsOf,
     saveTypeComponent,
+
+    -- * For Migrations
+    saveSerializedComponent,
   )
 where
 
 import Control.Lens
 import Data.ByteString qualified as BS
+import Data.ByteString.Lazy.Char8 qualified as BL
 import Data.List.NonEmpty qualified as NE
 import Data.List.NonEmpty qualified as NonEmpty
 import Data.Set qualified as Set
@@ -41,6 +45,7 @@ import Data.Vector (Vector)
 import Data.Vector qualified as Vector
 import Servant (err500)
 import Share.Codebase.Types (CodebaseEnv (..), CodebaseM)
+import Share.Codebase.Types qualified as Codebase
 import Share.IDs
 import Share.Postgres
 import Share.Postgres qualified as PG
@@ -64,6 +69,7 @@ import U.Codebase.Sqlite.HashHandle (HashHandle (..))
 import U.Codebase.Sqlite.LocalIds qualified as LocalIds
 import U.Codebase.Sqlite.Queries qualified as Util
 import U.Codebase.Sqlite.Symbol (Symbol)
+import U.Codebase.Sqlite.TempEntity (TempEntity)
 import U.Codebase.Sqlite.Term.Format qualified as TermFormat
 import U.Codebase.Sqlite.V2.HashHandle (v2HashHandle)
 import U.Codebase.Term qualified as V2
@@ -81,7 +87,10 @@ import Unison.Reference qualified as V1Reference
 import Unison.Referent qualified as V1Referent
 import Unison.Runtime.IOSource qualified as Decls
 import Unison.Server.Types qualified as Tags
+import Unison.Sync.Common qualified as SyncCommon
 import Unison.Sync.Types qualified as Share
+import Unison.SyncV2.Types (CBORBytes (..))
+import Unison.SyncV2.Types qualified as SyncV2
 
 type ResolvedLocalIds = LocalIds.LocalIds' Text ComponentHash
 
@@ -731,12 +740,12 @@ saveTermComponent componentHash elements = do
   let encodedElements =
         elements <&> \(localIds, trm, typ) ->
           (localIds, TermComponentElementBytes $ termComponentElementToByteString (TermComponentElement trm typ), typ)
-  saveEncodedTermComponent componentHash encodedElements
+  saveEncodedTermComponent componentHash Nothing encodedElements
 
 -- | Save an already-encoded term component to the database. This is more efficient than
 -- 'saveTermComponent' in cases where you've already got a serialized term (like during sync).
-saveEncodedTermComponent :: ComponentHash -> [(PgLocalIds, TermComponentElementBytes, TermFormat.Type)] -> CodebaseM e ()
-saveEncodedTermComponent componentHash elements = do
+saveEncodedTermComponent :: ComponentHash -> Maybe TempEntity -> [(PgLocalIds, TermComponentElementBytes, TermFormat.Type)] -> CodebaseM e ()
+saveEncodedTermComponent componentHash maySerialized elements = do
   codebaseOwnerUserId <- asks codebaseOwner
   componentHashId <- HashQ.ensureComponentHashId componentHash
   let elementsTable = elements & imap \i _ -> pgComponentIndex $ fromIntegral @Int i
@@ -775,7 +784,16 @@ saveEncodedTermComponent componentHash elements = do
         SELECT #{codebaseOwnerUserId}, element.term_id, element.bytes_id
           FROM elements element
     |]
+  doSaveSerialized componentHashId
   where
+    doSaveSerialized chId = do
+      componentEntity <- case maySerialized of
+        Just serialized -> pure serialized
+        Nothing -> do
+          SyncCommon.entityToTempEntity id . Share.TC <$> expectShareTermComponent chId
+      let serializedEntity = SyncV2.serialiseCBORBytes componentEntity
+      saveSerializedComponent chId serializedEntity
+
     saveSharedTermAndLocalMappings :: ComponentHashId -> CodebaseM e (NE.NonEmpty TermId)
     saveSharedTermAndLocalMappings componentHashId = do
       let HashHandle {toReference = hashType} = v2HashHandle
@@ -847,8 +865,8 @@ saveEncodedTermComponent componentHash elements = do
         |]
       pure termIds
 
-saveTypeComponent :: ComponentHash -> [(PgLocalIds, DeclFormat.Decl Symbol)] -> CodebaseM e ()
-saveTypeComponent componentHash elements = do
+saveTypeComponent :: ComponentHash -> Maybe TempEntity -> [(PgLocalIds, DeclFormat.Decl Symbol)] -> CodebaseM e ()
+saveTypeComponent componentHash maySerialized elements = do
   codebaseOwnerUserId <- asks codebaseOwner
   componentHashId <- HashQ.ensureComponentHashId componentHash
   let elementsTable = elements & imap \i _ -> fromIntegral @Int @Int32 i
@@ -886,7 +904,16 @@ saveTypeComponent componentHash elements = do
         SELECT #{codebaseOwnerUserId}, element.type_id, element.bytes_id
           FROM elements element
     |]
+  doSaveSerialized componentHashId
   where
+    doSaveSerialized chId = do
+      componentEntity <- case maySerialized of
+        Just serialized -> pure serialized
+        Nothing -> do
+          SyncCommon.entityToTempEntity id . Share.DC <$> expectShareTypeComponent chId
+      let serializedEntity = SyncV2.serialiseCBORBytes componentEntity
+      saveSerializedComponent chId serializedEntity
+
     saveConstructors :: [(TypeId, (PgLocalIds, DeclFormat.Decl Symbol))] -> CodebaseM e ()
     saveConstructors typeElements = do
       typeElementsWithResolvedLocals <- lift $ resolveLocalIdsOf (traversed . _2 . _1) typeElements
@@ -1121,3 +1148,14 @@ typeTagsByReferencesOf trav s = do
     tagFromDeclKind = \case
       DefnTypes.Ability -> Tags.Ability
       DefnTypes.Data -> Tags.Data
+
+saveSerializedComponent :: ComponentHashId -> CBORBytes TempEntity -> CodebaseM e ()
+saveSerializedComponent chId (CBORBytes bytes) = do
+  codebaseOwnerUserId <- asks Codebase.codebaseOwner
+  bytesId <- ensureBytesIdsOf id (BL.toStrict bytes)
+  execute_
+    [sql|
+      INSERT INTO serialized_components (user_id, component_hash_id, bytes_id)
+        VALUES (#{codebaseOwnerUserId}, #{chId}, #{bytesId})
+      ON CONFLICT DO NOTHING
+    |]