Merge pull request #101 from Nexus-Mods/import-export

Import export routines (backing up and restoring data)

src/NexusMods.MnemonicDB.Abstractions/Datom.cs

@@ -9,7 +9,7 @@ namespace NexusMods.MnemonicDB.Abstractions.DatomIterators;
 /// Represents a raw (unparsed) datom from an index. Most of the time this datom is only valid for the
 /// lifetime of the current iteration. It is not safe to store this datom for later use.
 /// </summary>
-public readonly struct Datom
+public readonly struct Datom : IEquatable<Datom>
 {
     private readonly KeyPrefix _prefix;
     private readonly ReadOnlyMemory<byte> _valueBlob;
@@ -42,6 +42,7 @@ public byte[] ToArray()
         _valueBlob.Span.CopyTo(array.AsSpan(KeyPrefix.Size));
         return array;
     }
+
 
     /// <summary>
     /// The KeyPrefix of the datom
@@ -143,4 +144,22 @@ public Datom Retract()
     {
         return new Datom(_prefix with {IsRetract = true, T = TxId.Tmp}, _valueBlob);
     }
+
+    /// <inheritdoc />
+    public bool Equals(Datom other)
+    {
+        return _prefix.Equals(other._prefix) && _valueBlob.Span.SequenceEqual(other._valueBlob.Span);
+    }
+
+    /// <inheritdoc />
+    public override bool Equals(object? obj)
+    {
+        return obj is Datom other && Equals(other);
+    }
+
+    /// <inheritdoc />
+    public override int GetHashCode()
+    {
+        return HashCode.Combine(_prefix.GetHashCode());
+    }
 }

src/NexusMods.MnemonicDB.Abstractions/IConnection.cs

@@ -55,6 +55,11 @@ public interface IConnection
     /// A service provider that entities can use to resolve their values
     /// </summary>
     public IServiceProvider ServiceProvider { get; }
+
+    /// <summary>
+    /// Gets the datom store that this connection uses to store data.
+    /// </summary>
+    public IDatomStore DatomStore { get; }
 
     /// <summary>
     /// Returns a snapshot of the database as of the given transaction id.

src/NexusMods.MnemonicDB.Abstractions/IDatomStore.cs

@@ -1,5 +1,6 @@
 using System;
 using System.Collections.Generic;
+using System.IO;
 using System.Threading.Tasks;
 using NexusMods.MnemonicDB.Abstractions.IndexSegments;
 using NexusMods.MnemonicDB.Abstractions.TxFunctions;
@@ -27,6 +28,17 @@ public interface IDatomStore : IDisposable
     /// </summary>
     AttributeCache AttributeCache { get; }
 
+    /// <summary>
+    /// Exports the database (including all indexes) to the given stream
+    /// </summary>
+    public Task ExportAsync(Stream stream);
+
+    /// <summary>
+    /// Imports the database (including all indexes) from the given stream.
+    /// Any existing data will be deleted before importing.
+    /// </summary>
+    public Task ImportAsync(Stream stream);
+
     /// <summary>
     ///     Transacts (adds) the given datoms into the store.
     /// </summary>

src/NexusMods.MnemonicDB.Abstractions/IndexSegments/IndexSegment.cs

@@ -47,7 +47,22 @@ public IndexSegment(ReadOnlySpan<byte> data, ReadOnlySpan<int> offsets, Attribut
 
         ReprocessData(_rowCount, data, offsets, memory.Span);
     }
+
+    /// <summary>
+    /// Create an index segment from raw data
+    /// </summary>
+    public IndexSegment(int rowCount, ReadOnlyMemory<byte> data, AttributeCache attributeCache)
+    {
+        _attributeCache = attributeCache;
+        _data = data;
+        _rowCount = rowCount;
+    }
 
+    /// <summary>
+    /// Gets read-only access to the data in this segment
+    /// </summary>
+    public ReadOnlyMemory<byte> Data => _data;
+
     /// <summary>
     /// All the upper values
     /// </summary>

src/NexusMods.MnemonicDB.Abstractions/Query/SliceDescriptor.cs

@@ -35,11 +35,25 @@ public readonly struct SliceDescriptor
     /// </summary>
     public bool IsReverse => From.Compare(To, Index) > 0;
 
+    /// <summary>
+    /// Returns this descriptor with a reversed iteration order.
+    /// </summary>
+    public SliceDescriptor Reversed()
+    {
+        return new SliceDescriptor
+        {
+            Index = Index,
+            From = To,
+            To = From
+        };
+    }
+
     /// <summary>
     /// Returns true if the datom is within the slice, false otherwise.
     /// </summary>
     public bool Includes(in Datom datom)
     {
+
         return Index switch
         {
             IndexType.TxLog => DatomComparators.TxLogComparator.Compare(From, datom) <= 0 &&
@@ -50,6 +64,9 @@ public bool Includes(in Datom datom)
             IndexType.AEVTCurrent or IndexType.AEVTHistory =>
                 DatomComparators.AEVTComparator.Compare(From, datom) <= 0 &&
                 DatomComparators.AEVTComparator.Compare(datom, To) < 0,
+            IndexType.AVETCurrent or IndexType.AVETHistory =>
+                DatomComparators.AVETComparator.Compare(From, datom) <= 0 &&
+                DatomComparators.AVETComparator.Compare(datom, To) < 0,
             IndexType.VAETCurrent or IndexType.VAETHistory =>
                 DatomComparators.VAETComparator.Compare(From, datom) <= 0 &&
                 DatomComparators.VAETComparator.Compare(datom, To) < 0,

src/NexusMods.MnemonicDB/Connection.cs

@@ -74,6 +74,9 @@ private IObservable<IDb> ProcessUpdates(IObservable<IDb> dbStream)
     /// <inheritdoc />
     public IServiceProvider ServiceProvider { get; set; }
 
+    /// <inheritdoc />
+    public IDatomStore DatomStore => _store;
+
     /// <inheritdoc />
     public IDb Db
     {

src/NexusMods.MnemonicDB/NexusMods.MnemonicDB.csproj

@@ -18,6 +18,9 @@
     </ItemGroup>
     <ItemGroup>
         <Compile Remove="SourceGenerationContext.cs"/>
+        <Compile Update="Storage\ImportExport.cs">
+          <DependentUpon>DatomStore.cs</DependentUpon>
+        </Compile>
     </ItemGroup>
     <Import Project="$([MSBuild]::GetPathOfFileAbove('NuGet.Build.props', '$(MSBuildThisFileDirectory)../'))"/>
 </Project>

src/NexusMods.MnemonicDB/Storage/DatomStore.cs

@@ -3,6 +3,7 @@
 using System.Collections.Frozen;
 using System.Collections.Generic;
 using System.Diagnostics;
+using System.IO;
 using System.Linq;
 using System.Runtime.InteropServices;
 using System.Threading;
@@ -22,7 +23,7 @@
 
 namespace NexusMods.MnemonicDB.Storage;
 
-public class DatomStore : IDatomStore
+public partial class DatomStore : IDatomStore
 {
     private readonly IIndex _aevtCurrent;
     private readonly IIndex _aevtHistory;
@@ -78,7 +79,7 @@ public class DatomStore : IDatomStore
     /// <summary>
     /// DI constructor
     /// </summary>
-    public DatomStore(ILogger<DatomStore> logger, DatomStoreSettings settings, IStoreBackend backend)
+    public DatomStore(ILogger<DatomStore> logger, DatomStoreSettings settings, IStoreBackend backend, bool bootstrap = true)
     {
         _remapFunc = Remap;
         _dbStream = new DbStream();
@@ -116,7 +117,8 @@ public DatomStore(ILogger<DatomStore> logger, DatomStoreSettings settings, IStor
         _avetCurrent = _backend.GetIndex(IndexType.AVETCurrent);
         _avetHistory = _backend.GetIndex(IndexType.AVETHistory);
 
-        Bootstrap();
+        if (bootstrap) 
+            Bootstrap();
     }
 
     /// <inheritdoc />
@@ -620,4 +622,5 @@ private unsafe PrevState GetPreviousState(bool isRemapped, AttributeId attrId, I
 
     #endregion
 
+
 }

src/NexusMods.MnemonicDB/Storage/ImportExport.cs

@@ -0,0 +1,112 @@
+using System;
+using System.IO;
+using System.Linq;
+using System.Threading.Tasks;
+using Microsoft.Extensions.Logging;
+using NexusMods.MnemonicDB.Abstractions;
+using NexusMods.MnemonicDB.Abstractions.IndexSegments;
+using NexusMods.MnemonicDB.Abstractions.Query;
+
+namespace NexusMods.MnemonicDB.Storage;
+
+public partial class DatomStore
+{
+    // File format:
+    // FOURCC: "MDBX"
+    // ushort: version
+    // one or more chunks
+    // 
+    // chunk:
+    // byte: IndexType
+    // uint: datomCount (number of datoms in the chunk)
+    // uint: chunkSize (in bytes)
+    // datomBlob
+
+    private static readonly byte[] FourCC = "MDBX"u8.ToArray();
+    private const int ChunkSize = 1024 * 16;
+
+    /// <summary>
+    /// Exports the database to the given stream
+    /// </summary>
+    public async Task ExportAsync(Stream stream)
+    {
+        var exportedDatoms = 0;
+        var binaryWriter = new BinaryWriter(stream);
+        binaryWriter.Write(FourCC);
+        binaryWriter.Write((ushort)1);
+
+        var snapshot = _backend.GetSnapshot();
+
+        foreach (var indexType in Enum.GetValues<IndexType>())
+        {
+            var slice = SliceDescriptor.Create(indexType);
+            var chunks = snapshot.DatomsChunked(slice, ChunkSize);
+
+            foreach (var chunk in chunks)
+            {
+                var data = chunk.Data;
+                binaryWriter.Write((byte)indexType);
+                binaryWriter.Write((uint)chunk.Count);
+                binaryWriter.Write((uint)data.Length);
+                binaryWriter.Write(data.Span);
+                exportedDatoms += chunk.Count;
+            }
+        }
+        _logger.LogInformation("Exported {0} datoms", exportedDatoms);
+    }
+
+    public async Task ImportAsync(Stream stream)
+    {
+        CleanStore();
+        var importedCount = 0;
+        var binaryReader = new BinaryReader(stream);
+        var fourCC = binaryReader.ReadBytes(4);
+        if (!fourCC.SequenceEqual(FourCC))
+            throw new InvalidDataException("Invalid file format");
+
+        var version = binaryReader.ReadUInt16();
+        if (version != 1)
+            throw new InvalidDataException("Invalid file version");
+
+        while (stream.Position < stream.Length)
+        {
+            var indexType = (IndexType)binaryReader.ReadByte();
+            var datomCount = binaryReader.ReadUInt32();
+            var chunkSize = binaryReader.ReadUInt32();
+            var data = binaryReader.ReadBytes((int)chunkSize);
+            var segment = new IndexSegment((int)datomCount, data.AsMemory(), _backend.AttributeCache);
+
+            using var batch = _backend.CreateBatch();
+            var index = _backend.GetIndex(indexType);
+
+            foreach (var datom in segment) 
+                index.Put(batch, datom);
+
+            batch.Commit();
+            importedCount += (int)datomCount;
+        }
+
+        _logger.LogInformation("Imported {0} datoms", importedCount);
+        _nextIdCache.ResetCaches();
+        Bootstrap();
+    }
+
+    private void CleanStore()
+    { 
+        int datomCount = 0;
+        var snapshot = _backend.GetSnapshot();
+        using var batch = _backend.CreateBatch();
+        foreach (var index in Enum.GetValues<IndexType>())
+        {
+            var slice = SliceDescriptor.Create(index);
+            var datoms = snapshot.Datoms(slice);
+            foreach (var datom in datoms)
+            {
+                _backend.GetIndex(index).Delete(batch, datom);
+                datomCount++;
+            }
+        }
+        batch.Commit();
+        _logger.LogInformation("Cleaned {0} datoms", datomCount);
+    }
+}