@web5/agent/src/store-data.ts

import type { Jwk } from '@web5/crypto';

import ms from 'ms';
import { Convert, NodeStream, TtlCache } from '@web5/common';

import type { Web5PlatformAgent } from './types/agent.js';

import { TENANT_SEPARATOR } from './utils-internal.js';
import { getDataStoreTenant } from './utils-internal.js';
import { DwnInterface, DwnMessageParams } from './types/dwn.js';
import { ProtocolDefinition, RecordsReadReplyEntry } from '@tbd54566975/dwn-sdk-js';

export type DataStoreTenantParams = {
  agent: Web5PlatformAgent;
  tenant?: string;
}

export type DataStoreListParams = DataStoreTenantParams;

export type DataStoreGetParams = DataStoreTenantParams & {
  id: string;
  useCache?: boolean;
}

export type DataStoreSetParams<TStoreObject> = DataStoreTenantParams & {
  id: string;
  data: TStoreObject;
  preventDuplicates?: boolean;
  updateExisting?: boolean;
  useCache?: boolean;
}

export type DataStoreDeleteParams = DataStoreTenantParams & {
  id: string;
}

export interface AgentDataStore<TStoreObject> {
  delete(params: DataStoreDeleteParams): Promise<boolean>;

  get(params: DataStoreGetParams): Promise<TStoreObject | undefined>;

  list(params: DataStoreTenantParams): Promise<TStoreObject[]>;

  set(params: DataStoreSetParams<TStoreObject>): Promise<void>;
}

export class DwnDataStore<TStoreObject extends Record<string, any> = Jwk> implements AgentDataStore<TStoreObject> {
  protected name = 'DwnDataStore';

  /**
     * Cache of Store Objects referenced by DWN record ID to Store Objects.
     *
     * Up to 100 entries are retained for 15 minutes.
     */
  protected _cache = new TtlCache<string, TStoreObject>({ ttl: ms('15 minutes'), max: 100 });

  /**
   * Index for mappings from Store Identifier to DWN record ID.
   * Since these values don't change, we can use a long TTL.
   *
   * Up to 1,000 entries are retained for 21 days.
   * NOTE: The maximum number for the ttl is 2^31 - 1 milliseconds (24.8 days), setting to 21 days to be safe.
   */
  protected _index = new TtlCache<string, string>({ ttl: ms('21 days'), max: 1000 });

  /**
   * Cache of tenant DIDs that have been initialized with the protocol.
   * This is used to avoid redundant protocol initialization requests.
   *
   * Since these are default protocols and unlikely to change, we can use a long TTL.
   */
  protected _protocolInitializedCache: TtlCache<string, boolean> = new TtlCache({ ttl: ms('21 days'), max: 1000 });

  /**
   * The protocol assigned to this storage instance.
   */
  protected _recordProtocolDefinition!: ProtocolDefinition;

  /**
   * Properties to use when writing and querying records with the DWN store.
   */
  protected _recordProperties = {
    dataFormat: 'application/json',
  };

  public async delete({ id, agent, tenant }: DataStoreDeleteParams): Promise<boolean> {
    // Determine the tenant identifier (DID) for the delete operation.
    const tenantDid = await getDataStoreTenant({ agent, tenant, didUri: id });

    // Look up the DWN record ID of the object in the store with the given `id`.
    let matchingRecordId = await this.lookupRecordId({ id, tenantDid, agent });

    // Return false if the given ID was not found in the store.
    if (!matchingRecordId) return false;

    // If a record for the given ID was found, attempt to delete it.
    const { reply: { status } } = await agent.dwn.processRequest({
      author        : tenantDid,
      target        : tenantDid,
      messageType   : DwnInterface.RecordsDelete,
      messageParams : { recordId: matchingRecordId }
    });

    // If the record was successfully deleted, update the index/cache and return true;
    if (status.code === 202) {
      this._index.delete(`${tenantDid}${TENANT_SEPARATOR}${id}`);
      this._cache.delete(matchingRecordId);
      return true;
    }

    // If the Delete operation failed, throw an error.
    throw new Error(`${this.name}: Failed to delete '${id}' from store: (${status.code}) ${status.detail}`);
  }

  public async get({ id, agent, tenant, useCache = false }:
    DataStoreGetParams
  ): Promise<TStoreObject | undefined> {
    // Determine the tenant identifier (DID) for the list operation.
    const tenantDid = await getDataStoreTenant({ agent, tenant, didUri: id });

    // Look up the DWN record ID of the object in the store with the given `id`.
    let matchingRecordId = await this.lookupRecordId({ id, tenantDid, agent });

    // Return undefined if no matches were found.
    if (!matchingRecordId) return undefined;

    // Retrieve and return the stored object.
    return await this.getRecord({ recordId: matchingRecordId, tenantDid, agent, useCache });
  }

  public async list({ agent, tenant}: DataStoreListParams): Promise<TStoreObject[]> {
    // Determine the tenant identifier (DID) for the list operation.
    const tenantDid = await getDataStoreTenant({ tenant, agent });

    // Query the DWN for all stored record objects.
    const storedRecords = await this.getAllRecords({ agent, tenantDid });

    return storedRecords;
  }

  public async set({ id, data, tenant, agent, preventDuplicates = true, updateExisting = false, useCache = false }:
    DataStoreSetParams<TStoreObject>
  ): Promise<void> {
    // Determine the tenant identifier (DID) for the set operation.
    const tenantDid = await getDataStoreTenant({ agent, tenant, didUri: id });

    // initialize the storage protocol if not already done
    await this.initialize({ tenant: tenantDid, agent });

    const messageParams: DwnMessageParams[DwnInterface.RecordsWrite] = { ...this._recordProperties };

    if (updateExisting) {
      // Look up the DWN record ID of the object in the store with the given `id`.
      const matchingRecordEntry = await this.getExistingRecordEntry({ id, tenantDid, agent });
      if (!matchingRecordEntry) {
        throw new Error(`${this.name}: Update failed due to missing entry for: ${id}`);
      }

      // set the recordId in the messageParams to update the existing record
      // set the dateCreated to the existing dateCreated as this is an immutable property
      messageParams.recordId = matchingRecordEntry.recordsWrite!.recordId;
      messageParams.dateCreated = matchingRecordEntry.recordsWrite!.descriptor.dateCreated;
    } else if (preventDuplicates) {
      // Look up the DWN record ID of the object in the store with the given `id`.
      const matchingRecordId = await this.lookupRecordId({ id, tenantDid, agent });
      if (matchingRecordId) {
        throw new Error(`${this.name}: Import failed due to duplicate entry for: ${id}`);
      }
    }


    // Convert the store object to a byte array, which will be the data payload of the DWN record.
    const dataBytes = Convert.object(data).toUint8Array();

    // Store the record in the DWN.
    const { message, reply: { status } } = await agent.dwn.processRequest({
      author        : tenantDid,
      target        : tenantDid,
      messageType   : DwnInterface.RecordsWrite,
      messageParams : { ...this._recordProperties, ...messageParams },
      dataStream    : new Blob([dataBytes], { type: 'application/json' })
    });

    // If the write fails, throw an error.
    if (!(message && status.code === 202)) {
      throw new Error(`${this.name}: Failed to write data to store for ${id}: ${status.detail}`);
    }

    // Add the ID of the newly created record to the index.
    this._index.set(`${tenantDid}${TENANT_SEPARATOR}${id}`, message.recordId);

    // If caching is enabled, add the store object to the cache.
    if (useCache) {
      this._cache.set(message.recordId, data);
    }
  }

  /**
   * Initialize the relevant protocol for the given tenant.
   * This confirms that the storage protocol is configured, otherwise it will be installed.
   */
  public async initialize({ tenant, agent }: DataStoreTenantParams) {
    const tenantDid = await getDataStoreTenant({ agent, tenant });
    if (this._protocolInitializedCache.has(tenantDid)) {
      return;
    }

    const { reply: { status, entries }} = await agent.dwn.processRequest({
      author        : tenantDid,
      target        : tenantDid,
      messageType   : DwnInterface.ProtocolsQuery,
      messageParams : {
        filter: {
          protocol: this._recordProtocolDefinition.protocol
        }
      },
    });

    if (status.code !== 200) {
      throw new Error(`Failed to query for protocols: ${status.code} - ${status.detail}`);
    }

    if (entries?.length === 0) {
      // protocol is not installed, install it
      await this.installProtocol(tenantDid, agent);
    }

    this._protocolInitializedCache.set(tenantDid, true);
  }

  protected async getAllRecords(_params: {
    agent: Web5PlatformAgent;
    tenantDid: string;
  }): Promise<TStoreObject[]> {
    throw new Error('Not implemented: Classes extending DwnDataStore must implement getAllRecords()');
  }

  private async getRecord({ recordId, tenantDid, agent, useCache }: {
    recordId: string;
    tenantDid: string;
    agent: Web5PlatformAgent;
    useCache: boolean;
  }): Promise<TStoreObject | undefined> {
    // If caching is enabled, check the cache for the record ID.
    if (useCache) {
      const record = this._cache.get(recordId);
      // If the record ID was present in the cache, return the associated store object.
      if (record) return record;
      // Otherwise, continue to read from the store.
    }

    // Read the record from the store.
    const { reply: readReply } = await agent.dwn.processRequest({
      author        : tenantDid,
      target        : tenantDid,
      messageType   : DwnInterface.RecordsRead,
      messageParams : { filter: { recordId } }
    });

    if (!readReply.entry?.data) {
      throw new Error(`${this.name}: Failed to read data from DWN for: ${recordId}`);
    }

    // If the record was found, convert back to store object format.
    const storeObject = await NodeStream.consumeToJson({ readable: readReply.entry.data }) as TStoreObject;

    // If caching is enabled, add the store object to the cache.
    if (useCache) {
      this._cache.set(recordId, storeObject);
    }

    return storeObject;
  }

  /**
   * Install the protocol for the given tenant using a `ProtocolsConfigure` message.
   */
  private async installProtocol(tenant: string, agent: Web5PlatformAgent) {
    const { reply : { status } } = await agent.dwn.processRequest({
      author        : tenant,
      target        : tenant,
      messageType   : DwnInterface.ProtocolsConfigure,
      messageParams : {
        definition: this._recordProtocolDefinition
      },
    });

    if (status.code !== 202) {
      throw new Error(`Failed to install protocol: ${status.code} - ${status.detail}`);
    }
  }

  private async lookupRecordId({ id, tenantDid, agent }: {
    id: string;
    tenantDid: string;
    agent: Web5PlatformAgent;
  }): Promise<string | undefined> {
    // Check the index for a matching ID and extend the index TTL.
    let recordId = this._index.get(`${tenantDid}${TENANT_SEPARATOR}${id}`, { updateAgeOnGet: true });

    // If no matching record ID was found in the index...
    if (!recordId) {
      // Query the DWN for all stored objects, which rebuilds the index.
      await this.getAllRecords({ agent, tenantDid });

      // Check the index again for a matching ID.
      recordId = this._index.get(`${tenantDid}${TENANT_SEPARATOR}${id}`);
    }

    return recordId;
  }

  private async getExistingRecordEntry({ id, tenantDid, agent }: {
    id: string;
    tenantDid: string;
    agent: Web5PlatformAgent;
  }): Promise<RecordsReadReplyEntry | undefined> {
    // Look up the DWN record ID of the object in the store with the given `id`.
    const recordId = await this.lookupRecordId({ id, tenantDid, agent });
    if (recordId) {
      // Read the record from the store.
      const { reply: readReply } = await agent.dwn.processRequest({
        author        : tenantDid,
        target        : tenantDid,
        messageType   : DwnInterface.RecordsRead,
        messageParams : { filter: { recordId } }
      });

      return readReply.entry;
    }
  }
}

export class InMemoryDataStore<TStoreObject extends Record<string, any> = Jwk> implements AgentDataStore<TStoreObject> {
  protected name = 'InMemoryDataStore';

  /**
   * A private field that contains the Map used as the in-memory data store.
   */
  private store: Map<string, TStoreObject> = new Map();

  public async delete({ id, agent, tenant }: DataStoreDeleteParams): Promise<boolean> {
    // Determine the tenant identifier (DID) for the delete operation.
    const tenantDid = await getDataStoreTenant({ agent, tenant, didUri: id });

    if (this.store.has(`${tenantDid}${TENANT_SEPARATOR}${id}`)) {
      // Record with given identifier exists so proceed with delete.
      this.store.delete(`${tenantDid}${TENANT_SEPARATOR}${id}`);
      return true;
    }

    // Record with given identifier not present so delete operation not possible.
    return false;
  }

  public async get({ id, agent, tenant }: DataStoreGetParams): Promise<TStoreObject | undefined> {
    // Determine the tenant identifier (DID) for the get operation.
    const tenantDid = await getDataStoreTenant({ agent, tenant, didUri: id });

    return this.store.get(`${tenantDid}${TENANT_SEPARATOR}${id}`);
  }

  public async list({ agent, tenant}: DataStoreListParams): Promise<TStoreObject[]> {
    // Determine the tenant identifier (DID) for the list operation.
    const tenantDid = await getDataStoreTenant({ tenant, agent });

    const result: TStoreObject[] = [];
    for (const [key, storedRecord] of this.store.entries()) {
      if (key.startsWith(`${tenantDid}${TENANT_SEPARATOR}`)) {
        result.push(storedRecord);
      }
    }

    return result;
  }

  public async set({ id, data, tenant, agent, preventDuplicates, updateExisting }: DataStoreSetParams<TStoreObject>): Promise<void> {
    // Determine the tenant identifier (DID) for the set operation.
    const tenantDid = await getDataStoreTenant({ agent, tenant, didUri: id });

    // If enabled, check if a record with the given `id` is already present in the store.
    if (updateExisting) {
      // Look up the DWN record ID of the object in the store with the given `id`.
      if (!this.store.has(`${tenantDid}${TENANT_SEPARATOR}${id}`)) {
        throw new Error(`${this.name}: Update failed due to missing entry for: ${id}`);
      }

      // set the recordId in the messageParams to update the existing record
    } else if (preventDuplicates) {
      const duplicateFound = this.store.has(`${tenantDid}${TENANT_SEPARATOR}${id}`);
      if (duplicateFound) {
        throw new Error(`${this.name}: Import failed due to duplicate entry for: ${id}`);
      }
    }

    // Make a deep copy so that the object stored does not share the same references as the input.
    const clonedData = structuredClone(data);
    this.store.set(`${tenantDid}${TENANT_SEPARATOR}${id}`, clonedData);
  }
}