TalkingBot/packages/core/src/storage.js

const fs = require('fs/promises');
const path = require('path');
const {
  PutObjectCommand,
  GetObjectCommand,
  HeadObjectCommand,
  ListObjectsV2Command,
  DeleteObjectCommand,
} = require('@aws-sdk/client-s3');

/**
 * Skip a step when its artifact is already stored.
 * @param {{ exists: (key: string) => Promise<boolean> }} storage
 * @param {string} key
 * @param {boolean} force
 */
async function shouldSkip(storage, key, force) {
  if (force) return false;
  return storage.exists(key);
}

class LocalFsStorage {
  /**
   * @param {string} root
   */
  constructor(root) {
    this.root = root;
  }

  resolve(key) {
    const full = path.resolve(this.root, key);
    const root = path.resolve(this.root);
    if (!full.startsWith(root + path.sep) && full !== root) {
      throw new Error(`Refusing to write outside storage root: ${key}`);
    }
    return full;
  }

  async putBytes(key, body) {
    const full = this.resolve(key);
    await fs.mkdir(path.dirname(full), { recursive: true });
    await fs.writeFile(full, body);
  }

  async getBytes(key) {
    return fs.readFile(this.resolve(key));
  }

  async putJson(key, value) {
    await this.putText(key, `${JSON.stringify(value, null, 2)}\n`);
  }

  async getJson(key) {
    const text = await this.getText(key);
    return JSON.parse(text);
  }

  async putText(key, text) {
    await this.putBytes(key, Buffer.from(text, 'utf8'));
  }

  async getText(key) {
    return fs.readFile(this.resolve(key), 'utf8');
  }

  async exists(key) {
    try {
      await fs.access(this.resolve(key));
      return true;
    } catch (error) {
      if (error.code === 'ENOENT') return false;
      throw error;
    }
  }

  async delete(key) {
    await fs.rm(this.resolve(key), { force: true });
  }

  async list(prefix) {
    const root = path.resolve(this.root);
    const keys = [];
    async function walk(dir) {
      let entries;
      try {
        entries = await fs.readdir(dir, { withFileTypes: true });
      } catch (error) {
        if (error.code === 'ENOENT') return;
        throw error;
      }
      for (const entry of entries) {
        const full = path.join(dir, entry.name);
        if (entry.isDirectory()) await walk(full);
        else {
          const key = path.relative(root, full).split(path.sep).join('/');
          if (key.startsWith(prefix)) keys.push(key);
        }
      }
    }
    await walk(root);
    return keys.sort();
  }
}

class S3Storage {
  /**
   * @param {import('@aws-sdk/client-s3').S3Client} client
   * @param {string} bucket
   */
  constructor(client, bucket) {
    this.client = client;
    this.bucket = bucket;
  }

  async putBytes(key, body, contentType = 'application/octet-stream') {
    await this.client.send(
      new PutObjectCommand({
        Bucket: this.bucket,
        Key: key,
        Body: body,
        ContentType: contentType,
      }),
    );
  }

  async getBytes(key) {
    const response = await this.client.send(new GetObjectCommand({ Bucket: this.bucket, Key: key }));
    const parts = [];
    for await (const part of response.Body) parts.push(part);
    return Buffer.concat(parts);
  }

  async putJson(key, value) {
    await this.putBytes(key, Buffer.from(`${JSON.stringify(value, null, 2)}\n`), 'application/json');
  }

  async getJson(key) {
    return JSON.parse(await this.getText(key));
  }

  async putText(key, text) {
    await this.putBytes(key, Buffer.from(text, 'utf8'), 'text/plain; charset=utf-8');
  }

  async getText(key) {
    return (await this.getBytes(key)).toString('utf8');
  }

  async exists(key) {
    try {
      await this.client.send(new HeadObjectCommand({ Bucket: this.bucket, Key: key }));
      return true;
    } catch (error) {
      const status = error.$metadata && error.$metadata.httpStatusCode;
      if (status === 404 || error.name === 'NotFound' || error.name === 'NoSuchKey') return false;
      throw error;
    }
  }

  async delete(key) {
    await this.client.send(new DeleteObjectCommand({ Bucket: this.bucket, Key: key }));
  }

  async list(prefix) {
    const keys = [];
    let token;
    do {
      const response = await this.client.send(
        new ListObjectsV2Command({
          Bucket: this.bucket,
          Prefix: prefix,
          ContinuationToken: token,
        }),
      );
      for (const object of response.Contents || []) keys.push(object.Key);
      token = response.IsTruncated ? response.NextContinuationToken : undefined;
    } while (token);
    return keys;
  }
}

class ReadStorage {
  /**
   * Read-only view. Chat uses this so the process has no put or delete method.
   * @param {{ getBytes: Function, getText: Function, getJson: Function, exists: Function, list: Function }} storage
   */
  constructor(storage) {
    this.storage = storage;
  }

  getBytes(key) {
    return this.storage.getBytes(key);
  }

  getText(key) {
    return this.storage.getText(key);
  }

  getJson(key) {
    return this.storage.getJson(key);
  }

  exists(key) {
    return this.storage.exists(key);
  }

  list(prefix) {
    return this.storage.list(prefix);
  }
}

module.exports = { shouldSkip, LocalFsStorage, S3Storage, ReadStorage };