RxDB plugin to allow storage of attachments in external storage

rxdb-external-attachments-storage.ts

import {
  BulkWriteRow,
  EventBulk,
  RxAttachmentData,
  RxAttachmentWriteData,
  RxDocumentData,
  RxDocumentDataById,
  RxJsonSchema,
  RxStorage,
  RxStorageBulkWriteResponse,
  RxStorageChangeEvent,
  RxStorageInstance,
  RxStorageInstanceCreationParams,
  RxStorageQueryResult,
  RxStorageStatics,
} from "rxdb";
import {
  CompositePrimaryKey,
  PrimaryKey,
  StringKeys,
} from "rxdb/dist/types/types";
import { Observable } from "rxjs";

export interface RxAttachmentProvider {
  getDocumentAttachments: (
    collectionName: string,
    documentId: string
  ) => Promise<{ [id: string]: RxAttachmentData }>;
  getAttachment: (
    collectionName: string,
    documentId: string,
    attachmentId: string
  ) => Promise<string>;
  putAttachment: (
    collectionName: string,
    documentId: string,
    attachmentId: string,
    attachment: RxAttachmentWriteData
  ) => Promise<void>;
}

type ExternalAttachmentsInternals<RxDocType> = {
  storageInstance: RxStorageInstance<RxDocType, any, any>;
};

export type ExternalAttachmentsSettings = {
  storage: RxStorage<any, any>;
  attachments: RxAttachmentProvider;
};

export class RxExternalAttachments
  implements
    RxStorage<ExternalAttachmentsInternals<any>, ExternalAttachmentsSettings>
{
  public name = "external-attachments";
  public statics: RxStorageStatics;

  private storage: RxStorage<any, any>;

  constructor(public readonly settings: ExternalAttachmentsSettings) {
    this.statics = settings.storage.statics;
    this.storage = settings.storage;
  }

  async createStorageInstance<RxDocType>(
    params: RxStorageInstanceCreationParams<RxDocType, any>
  ): Promise<RxStorageInstance<RxDocType, any, any>> {
    const storageInstance = await this.storage.createStorageInstance(params);
    const internals: ExternalAttachmentsInternals<any> = {
      storageInstance,
    };

    return new RxExternalAttachmentsStorageInstance(
      this,
      params.databaseName,
      params.collectionName,
      params.schema,
      internals,
      this.settings
    );
  }
}

export class RxExternalAttachmentsStorageInstance<RxDocType>
  implements
    RxStorageInstance<RxDocType, ExternalAttachmentsInternals<RxDocType>, any>
{
  private primaryKey: Extract<keyof RxDocType, string>;
  private provider: RxAttachmentProvider;

  constructor(
    public readonly storage: RxExternalAttachments,
    public readonly databaseName: string,
    public readonly collectionName: string,
    public readonly schema: Readonly<RxJsonSchema<RxDocumentData<RxDocType>>>,
    public readonly internals: ExternalAttachmentsInternals<RxDocType>,
    public readonly options: Readonly<ExternalAttachmentsSettings>
  ) {
    this.primaryKey = getPrimaryFieldOfPrimaryKey(schema.primaryKey) as any;
    this.provider = options.attachments;
  }

  async _resolveDocumentAttachments(
    documentId: string,
    doc: RxDocumentData<RxDocType>
  ): Promise<RxDocumentData<RxDocType>> {
    const attachments = await this.provider.getDocumentAttachments(
      this.collectionName,
      documentId
    );
    return { ...doc, _attachments: attachments };
  }

  async _resolveAttachments(
    docs: RxDocumentDataById<RxDocType>
  ): Promise<RxDocumentDataById<RxDocType>> {
    if (!this.schema.attachments) return docs;

    const promises = Object.entries(docs).map(async ([documentId, doc]) => [
      documentId,
      await this._resolveDocumentAttachments(documentId, doc),
    ]);

    const results = await Promise.all(promises);

    return Object.fromEntries(results);
  }

  async bulkWrite(
    documentWrites: BulkWriteRow<RxDocType>[]
  ): Promise<RxStorageBulkWriteResponse<RxDocType>> {
    for (const row of documentWrites) {
      for (const [attachmentId, attachment] of Object.entries(
        row.document._attachments
      )) {
        if (!("data" in attachment)) continue;

        await this.provider.putAttachment(
          this.collectionName,
          row.document[this.primaryKey] as any,
          attachmentId,
          attachment
        );
      }
    }

    const result = await this.internals.storageInstance.bulkWrite(
      documentWrites
    );

    if (result.success) {
      const docsWithAttachments = await this._resolveAttachments(
        result.success
      );

      return { ...result, success: docsWithAttachments };
    }

    return result;
  }

  async findDocumentsById(
    ids: string[],
    withDeleted: boolean
  ): Promise<RxDocumentDataById<RxDocType>> {
    const docs = await this.internals.storageInstance.findDocumentsById(
      ids,
      withDeleted
    );

    return await this._resolveAttachments(docs);
  }

  async query(preparedQuery: any): Promise<RxStorageQueryResult<RxDocType>> {
    const results = await this.internals.storageInstance.query(preparedQuery);

    if (!this.schema.attachments) return results;

    const documents = await Promise.all(
      results.documents.map(async (doc) => ({
        ...doc,
        _attachments: await this._resolveDocumentAttachments(
          doc[this.primaryKey] as any,
          doc
        ),
      }))
    );

    return { documents: documents as any };
  }

  getAttachmentData(documentId: string, attachmentId: string): Promise<string> {
    return this.provider.getAttachment(
      this.collectionName,
      documentId,
      attachmentId
    );
  }

  async getChangedDocumentsSince(
    limit: number,
    checkpoint?: any
  ): Promise<{ document: RxDocumentData<RxDocType>; checkpoint: any }[]> {
    const results =
      await this.internals.storageInstance.getChangedDocumentsSince(
        limit,
        checkpoint
      );

    // await Promise.all(
    //   results.map(async ({ document: doc }) => {
    //     doc._attachments = await this.options.getAttachmentsData(doc);
    //   })
    // );

    return results;
  }

  changeStream(): Observable<
    EventBulk<RxStorageChangeEvent<RxDocumentData<RxDocType>>>
  > {
    return this.internals.storageInstance.changeStream();
  }

  cleanup(minimumDeletedTime: number): Promise<boolean> {
    return this.internals.storageInstance.cleanup(minimumDeletedTime);
  }

  close(): Promise<void> {
    return this.internals.storageInstance.close();
  }

  remove(): Promise<void> {
    return this.internals.storageInstance.remove();
  }
}

export function getRxExternalAttachments(
  settings: ExternalAttachmentsSettings
): RxExternalAttachments {
  return new RxExternalAttachments(settings);
}

function getPrimaryFieldOfPrimaryKey<RxDocType>(
  primaryKey: PrimaryKey<RxDocType>
): StringKeys<RxDocType> {
  if (typeof primaryKey === "string") {
    return primaryKey as any;
  } else {
    return (primaryKey as CompositePrimaryKey<RxDocType>).key;
  }
}