Relay Snapshots

App.tsx

// create relay environment
import { RelaySnapshot } from './RelaySnapshot'
import { RecordSource } from './RecordSource'
import { Store } from './Store'

export const environment = new Environment({
  network: Network.create(fetchQuery),
  store: new Store(new RecordSource()),
})

// some component
const App: React.FC = () => {
  const environment = useRelayEnvironment()
  const snapshot = React.useRef<RelaySnapshot>()

  const onTakeSnapshot = () => {
    environment.getStore().createSnapshot()
  }

  const onRestoreSnapshot = () => {
    environment.getStore().restoreSnapshot(snapshot.current)
  }

  // ...
}

RecordSource.ts

import { DataID, RecordState } from 'relay-runtime'
import {
  MutableRecordSource,
  Record,
} from 'relay-runtime/lib/store/RelayStoreTypes'
import { RelaySnapshot } from './RelaySnapshot'

type RecordMap = {
  [key: string]: any
}

enum RelayRecordState {
  EXISTENT = 'EXISTENT',
  NONEXISTENT = 'NONEXISTENT',
  UNKNOWN = 'UNKNOWN',
}

export class RecordSource implements MutableRecordSource {
  _records: RecordMap
  _snapshot: RelaySnapshot

  constructor(records?: RecordMap) {
    this._records = records || {}
  }

  clear(): void {
    this._records = {}
  }

  delete(dataID: DataID): void {
    this._records[dataID] = null
  }

  get(dataID: DataID): Record | null | undefined {
    return this._records[dataID]
  }

  getRecordIDs(): Array<DataID> {
    return Object.keys(this._records)
  }

  getStatus(dataID: DataID): RecordState {
    if (!this._records.hasOwnProperty(dataID)) {
      return RelayRecordState.UNKNOWN
    }

    return this._records[dataID] == null
      ? RelayRecordState.NONEXISTENT
      : RelayRecordState.EXISTENT
  }

  has(dataID: DataID): boolean {
    return this._records.hasOwnProperty(dataID)
  }

  remove(dataID: DataID): void {
    delete this._records[dataID]
  }

  set(dataID: DataID, record: Record): void {
    if (this._snapshot) {
      this._snapshot.addUpdatedRecordID(dataID)
    }

    this._records[dataID] = record
  }

  size(): number {
    return Object.keys(this._records).length
  }

  toJSON(): RecordMap {
    return this._records
  }

  snapshot(): RelaySnapshot {
    this._snapshot = new RelaySnapshot(this._records)

    return this._snapshot
  }

  restore(snapshot: RelaySnapshot): void {
    this._records = { ...snapshot.getRecords() }
  }
}

RelaySnapshot.ts

import { RecordMap } from './RelayTypes'
import { DataID } from 'react-relay'

export class RelaySnapshot {
  private _records: RecordMap
  private _updatedRecordIDs: Set<DataID>

  constructor(records: RecordMap) {
    this._records = { ...records }
    this._updatedRecordIDs = new Set()
  }

  addUpdatedRecordID(dataID: DataID): void {
    this._updatedRecordIDs.add(dataID)
  }

  getUpdatedRecordIDs(): Set<DataID> {
    return this._updatedRecordIDs
  }

  getRecords(): RecordMap {
    return this._records
  }
}

Store.ts

import { Store as RelayStore } from 'relay-runtime'
import { RelaySnapshot } from './RelaySnapshot'
import { RelayRecordSource } from './RelayRecordSource'

export class Store extends RelayStore {
  createSnapshot(): RelaySnapshot {
    const recordSource = this.getSource() as RelayRecordSource

    return recordSource.snapshot()
  }

  restoreSnapshot(snapshot: RelaySnapshot): void {
    const self = this as any
    const recordSource = this.getSource() as RelayRecordSource
    const updatedRecordIDs = [...snapshot.getUpdatedRecordIDs()]

    recordSource.restore(snapshot)

    self._updatedRecordIDs = updatedRecordIDs.reduce((map, id) => {
      map[id] = true

      return map
    }, {})

    self.publish(recordSource)
    self.notify()
  }
}