Add partition id

Nexus-Mods · Oct 1, 2024 · 206fc11 · 206fc11
1 parent 0af5a32
commit 206fc11
Show file tree

Hide file tree

Showing 2 changed files with 20 additions and 7 deletions.
diff --git a/src/Abstractions/NexusMods.Abstractions.Resources.DB/PersistedResourceLoader.cs b/src/Abstractions/NexusMods.Abstractions.Resources.DB/PersistedResourceLoader.cs
@@ -25,6 +25,7 @@ public class PersistedResourceLoader<TResourceIdentifier, TData> : IResourceLoad
     private readonly IdentifierToHash _identifierToHash;
     private readonly IdentifierToEntityId _identifierToEntityId;
     private readonly AttributeId _referenceAttributeId;
+    private readonly Optional<PartitionId> _partitionId;
 
     public PersistedResourceLoader(
         IConnection connection,
@@ -33,6 +34,7 @@ public PersistedResourceLoader(
         DataToBytes dataToBytes,
         BytesToData bytesToData,
         IdentifierToEntityId identifierToEntityId,
+        Optional<PartitionId> partitionId,
         IResourceLoader<TResourceIdentifier, TData> innerLoader)
     {
         _connection = connection;
@@ -46,8 +48,10 @@ public PersistedResourceLoader(
 
         _referenceAttribute = referenceAttribute;
         _referenceAttributeId = _referenceAttribute.GetDbId(_connection.Registry.Id);
+        _partitionId = partitionId;
     }
 
+    /// <inheritdoc/>
     public ValueTask<Resource<TData>> LoadResourceAsync(TResourceIdentifier resourceIdentifier, CancellationToken cancellationToken)
     {
         var entityId = _identifierToEntityId(resourceIdentifier);
@@ -96,14 +100,16 @@ private async ValueTask<Resource<TData>> SaveResource(ValueTuple<EntityId, TReso
         var bytes = _dataToBytes(resource.Data);
 
         using var tx = _connection.BeginTransaction();
-        var persisted = new PersistedResource.New(tx)
+        var tmpId = _partitionId.HasValue ? tx.TempId(_partitionId.Value) : tx.TempId();
+
+        var persisted = new PersistedResource.New(tx, tmpId)
         {
             Data = bytes,
             ExpiresAt = resource.ExpiresAt,
             ResourceIdentifierHash = _identifierToHash(resourceIdentifier.Item2),
         };
 
-        _referenceAttribute.Add(tx, resourceIdentifier.Item1, persisted.Id);
+        _referenceAttribute.Add(tx, resourceIdentifier.Item1, persisted);
         await tx.Commit();
 
         return resource;
@@ -121,7 +127,8 @@ public static IResourceLoader<TResourceIdentifier, byte[]> Persist<TResourceIden
         IConnection connection,
         ReferenceAttribute<PersistedResource> referenceAttribute,
         PersistedResourceLoader<TResourceIdentifier, byte[]>.IdentifierToHash identifierToHash,
-        PersistedResourceLoader<TResourceIdentifier, byte[]>.IdentifierToEntityId identifierToEntityId)
+        PersistedResourceLoader<TResourceIdentifier, byte[]>.IdentifierToEntityId identifierToEntityId,
+        Optional<PartitionId> partitionId)
         where TResourceIdentifier : notnull
     {
         return inner.Persist(
@@ -130,7 +137,8 @@ public static IResourceLoader<TResourceIdentifier, byte[]> Persist<TResourceIden
             identifierToHash: identifierToHash,
             identifierToEntityId: identifierToEntityId,
             dataToBytes: static bytes => bytes,
-            bytesToData: static bytes => bytes
+            bytesToData: static bytes => bytes,
+            partitionId: partitionId
         );
     }
 
@@ -141,19 +149,21 @@ public static IResourceLoader<TResourceIdentifier, TData> Persist<TResourceIdent
         PersistedResourceLoader<TResourceIdentifier, TData>.IdentifierToHash identifierToHash,
         PersistedResourceLoader<TResourceIdentifier, TData>.IdentifierToEntityId identifierToEntityId,
         PersistedResourceLoader<TResourceIdentifier, TData>.DataToBytes dataToBytes,
-        PersistedResourceLoader<TResourceIdentifier, TData>.BytesToData bytesToData)
+        PersistedResourceLoader<TResourceIdentifier, TData>.BytesToData bytesToData,
+        Optional<PartitionId> partitionId)
         where TResourceIdentifier : notnull
         where TData : notnull
     {
         return inner.Then(
-            state: (connection, referenceAttribute, identifierToHash, identifierToEntityId, dataToBytes, bytesToData),
+            state: (connection, referenceAttribute, identifierToHash, identifierToEntityId, dataToBytes, bytesToData, partitionId),
             factory: static (input, inner) => new PersistedResourceLoader<TResourceIdentifier, TData>(
                 connection: input.connection,
                 referenceAttribute: input.referenceAttribute,
                 identifierToHash: input.identifierToHash,
                 identifierToEntityId: input.identifierToEntityId,
                 dataToBytes: input.dataToBytes,
                 bytesToData: input.bytesToData,
+                partitionId: input.partitionId,
                 innerLoader: inner
             )
         );

diff --git a/tests/NexusMods.UI.Tests/ImageLoaderTests.cs b/tests/NexusMods.UI.Tests/ImageLoaderTests.cs
@@ -50,6 +50,8 @@ public async Task Test()
 
     private IResourceLoader<EntityId, Lifetime<Bitmap>> CreatePipeline()
     {
+        const byte partitionId = 123;
+
         var pipeline = HttpLoader.CreateDefault()
             .ChangeIdentifier<ValueTuple<EntityId, Uri>, Uri, byte[]>(static tuple => tuple.Item2)
             .Decode(decoderType: DecoderType.Skia)
@@ -62,7 +64,8 @@ private IResourceLoader<EntityId, Lifetime<Bitmap>> CreatePipeline()
                 connection: Connection,
                 referenceAttribute: NexusModsModPageMetadata.ThumbnailResource,
                 identifierToHash: static tuple => tuple.Item2.ToString().XxHash64AsUtf8(),
-                identifierToEntityId: static tuple => tuple.Item1
+                identifierToEntityId: static tuple => tuple.Item1,
+                partitionId: PartitionId.User(partitionId)
             )
             .Decode(decoderType: DecoderType.Qoi)
             .ToAvaloniaBitmap()