add S3 compatibility

.snaplet/config.json

@@ -0,0 +1,6 @@
+{
+  "targetDatabaseUrl": "postgresql://postgres@localhost:5432/snaplet_development",
+  "sourceDatabaseUrl": "postgresql://postgres@localhost:5432/snaplet_development",
+  "s3Bucket": "snaplet",
+  "s3Region": "weur"
+}

cli/package.json

@@ -16,6 +16,9 @@
     "snaplet_internal": "bin/snaplet"
   },
   "dependencies": {
+    "@aws-sdk/client-s3": "^3.614.0",
+    "@aws-sdk/lib-storage": "3.616.0",
+    "@paralleldrive/cuid2": "^2.2.2",
     "@sentry/integrations": "7.48.0",
     "@sentry/node": "7.48.0",
     "@snaplet/copycat": "5.0.0",
@@ -79,6 +82,7 @@
     "@types/prompts": "2.4.4",
     "@types/semver": "7.3.13",
     "@types/sinon": "10.0.14",
+    "@types/temp": "^0.9.4",
     "@types/uuid": "8.3.4",
     "@types/yargs": "17.0.24",
     "axios": "0.27.2",

cli/src/commands/snapshot/actions/list/lib/storage.ts

@@ -0,0 +1,195 @@
+import { S3Client } from '@aws-sdk/client-s3'
+import {
+  checkIfObjectExists,
+  downloadFileFromBucket,
+  initClient,
+  S3Settings,
+  uploadFileToBucket,
+} from '~/lib/s3.js'
+
+import DatabaseConstructor, { type Database } from 'better-sqlite3'
+
+import fsExtra from 'fs-extra'
+import path from 'path'
+
+const DATABASE_NAME = 'db.sqlite'
+
+/**
+ * we store a list of snapshots in a sqlite db (a ledger), as S3 is missing
+ * support for filtering by tags.
+ */
+class SnapshotListStorage {
+  client: S3Client
+  bucketName: string
+  db: Database | undefined
+  dbPath: string
+
+  constructor(settings: S3Settings, dbPath: string) {
+    this.client = initClient(settings)
+    this.bucketName = settings.bucket
+    this.dbPath = path.join(dbPath, DATABASE_NAME)
+  }
+
+  /**
+   * create a new database instance we
+   * can use to make queries against.
+   */
+  async init() {
+    // TODO_BEFORE_REVIEW: look at saving the file at temp location
+    // is it is being used as a placeholder at the moment.
+    const isFound = await this.downloadDatabaseFile()
+
+    if (isFound) {
+      this.db = new DatabaseConstructor(this.dbPath)
+    } else {
+      this.db = await this.createNewDatabase()
+    }
+    this.db.pragma('journal_mode = WAL')
+  }
+
+  async createNewDatabase() {
+    const newDb = new DatabaseConstructor(this.dbPath)
+
+    newDb.exec(`CREATE TABLE snapshots (
+      id TEXT PRIMARY KEY,
+      name TEXT NOT NULL,
+      tags JSON NOT NULL DEFAULT '[]',
+      created_at DATETIME DEFAULT current_timestamp
+    )`)
+
+    await this.saveDatabaseFile()
+
+    return newDb
+  }
+
+  /**
+   *
+   */
+  async saveDatabaseFile() {
+    const dbFile = await fsExtra.createReadStream(this.dbPath)
+
+    return uploadFileToBucket(dbFile, {
+      bucket: this.bucketName,
+      client: this.client,
+      key: DATABASE_NAME,
+    })
+  }
+
+  /**
+   * download `db.sqlite` file from S3 bucket
+   * and write to the `.snapshots` file.
+   *
+   * @returns {boolean}: returns false if not found
+   */
+  async downloadDatabaseFile() {
+    const isFound = await checkIfObjectExists(
+      this.bucketName,
+      DATABASE_NAME,
+      this.client
+    )
+
+    if (isFound) {
+      const snapshotDb = await downloadFileFromBucket(
+        this.bucketName,
+        DATABASE_NAME,
+        { client: this.client }
+      )
+
+      if (snapshotDb) {
+        await fsExtra.writeFile(this.dbPath, snapshotDb)
+        return true
+      }
+    }
+
+    return false
+  }
+
+  getDatabase() {
+    if (this.db) {
+      return this.db
+    } else {
+      throw new Error('Database instance not found, run `init`')
+    }
+  }
+
+  async destoryDatabaseFile() {
+    return fsExtra.unlink(this.dbPath)
+  }
+}
+
+export type SnapshotListStorage = Awaited<
+  ReturnType<typeof snapshotListStorage>
+>
+
+export const snapshotListStorage = async (
+  settings: S3Settings,
+  /** base `.snaplet` folder */
+  basePath: string
+) => {
+  const storage = new SnapshotListStorage(settings, basePath)
+  await storage.init()
+
+  const db = storage.getDatabase()
+
+  return {
+    getSnapshotsMany: (rules?: { startsWith?: string }) => {
+      return db
+        .prepare(
+          [
+            'SELECT * FROM snapshots',
+            rules?.startsWith ? 'WHERE ? LIKE name || %' : null,
+          ].join(' ')
+        )
+        .all(rules?.startsWith)
+        .map((r: any) => ({
+          id: r.id,
+          name: r.name,
+          tags: r.tags,
+          createdAt: r.createdAt,
+        }))
+    },
+    getLatestSnapshot: () => {
+      return db
+        .prepare('SELECT * FROM snapshots ORDER BY created_at LIMIT 1')
+        .all()
+        .map((r: any) => ({
+          id: r.id,
+          name: r.name,
+          tags: r.tags,
+          createdAt: r.created_at,
+        }))
+    },
+    insertSnapshot: (data: {
+      id: string
+      name: string
+      createdAt: string
+      tags?: string[]
+    }) => {
+      return db.exec(`INSERT INTO snapshots (
+        id, name, created_at, tags
+      ) VALUES (
+        ${[
+          `'${data.id}'`,
+          `'${data.name}'`,
+          `'${data.createdAt}'`,
+          data.tags ? `[${data.tags.toString()}]` : '[]',
+        ].join(',')}
+      )`)
+    },
+    /**
+     * save current database instance to S3
+     * bucket, will automatically delete the
+     */
+    commit: async (
+      opts: {
+        /** delete the file stored  */
+        destroy: boolean
+      } = { destroy: true }
+    ) => {
+      await storage.saveDatabaseFile()
+      if (opts.destroy) {
+        await storage.destoryDatabaseFile()
+      }
+    },
+  }
+}

cli/src/commands/snapshot/actions/share/debugShare.ts

@@ -0,0 +1,3 @@
+import { xdebug } from '@snaplet/sdk/cli'
+
+export const xdebugShare = xdebug.extend('share') // snaplet:share

cli/src/commands/snapshot/actions/share/shareAction.handler.ts

@@ -0,0 +1,87 @@
+import { getHosts } from '~/lib/hosts/hosts.js'
+import type { CommandOptions } from './shareAction.types.js'
+import { findSnapshotSummary } from '~/components/findSnapshotSummary.js'
+import { exitWithError } from '~/lib/exit.js'
+import { encryptAndCompressTables } from './steps/encryptAndCompressTables.js'
+import SnapshotCache from '../restore/lib/SnapshotCache.js'
+
+import { encryptionPayloadStep } from './steps/encryptionPayload.js'
+import { activity } from '~/lib/spinner.js'
+import { uploadSnapshot } from '~/lib/snapshotStorage.js'
+import { getOrCreateSnapshotIdStep } from './steps/getOrCreateSnapshot.js'
+import { needs } from '~/components/needs/index.js'
+import { snapshotListStorage } from '../list/lib/storage.js'
+import path from 'path'
+
+export async function handler(options: CommandOptions) {
+  const { snapshotName: startsWith, latest, tags, noEncrypt } = await options
+
+  const isEncryptionSelected = noEncrypt === false
+
+  const settings = await needs.s3Settings()
+
+  const hosts = await getHosts({ only: ['local', 'abspath'] })
+  const sss = await findSnapshotSummary(
+    {
+      latest,
+      startsWith,
+      tags,
+    },
+    hosts
+  )
+
+  if (!sss.cachePath) {
+    console.log('Error: A snapshot must be cached in order to be shared')
+    await exitWithError('UNHANDLED_ERROR')
+
+    return
+  }
+
+  const cache = new SnapshotCache(sss)
+
+  /**
+   * if encryption is enabled, read the user config and
+   * generate a payload we use to encrypt a snapshot.
+   * */
+  const encryptionPayload = isEncryptionSelected
+    ? await encryptionPayloadStep()
+    : undefined
+
+  await encryptAndCompressTables({ paths: cache.paths, encryptionPayload })
+
+  /**
+   * upload snapshot to object storage
+   */
+  const snapshotId = await getOrCreateSnapshotIdStep(sss)
+  const act = activity('Snapshot', 'Uploading...')
+
+  try {
+    await uploadSnapshot(snapshotId, cache.paths, {
+      onProgress: async (percentage) => {
+        act.info(`Uploading... [${percentage}%]`)
+      },
+      settings,
+    })
+  } catch (err: any) {
+    act.fail(err?.message)
+    await exitWithError('SNAPSHOT_CAPTURE_INCOMPLETE_ERROR')
+  } finally {
+    act.done()
+  }
+
+  const storage = await snapshotListStorage(
+    settings,
+    path.join(cache.paths.base, '..')
+  )
+
+  storage.insertSnapshot({
+    id: snapshotId,
+    name: sss.summary.name,
+    tags: sss.summary.tags,
+    createdAt: new Date().toString(),
+  })
+
+  await storage.commit()
+
+  console.log(`Snapshot "${sss.summary.name}" shared`)
+}

cli/src/commands/snapshot/actions/share/shareAction.ts

@@ -0,0 +1,42 @@
+import type { CommandModule } from 'yargs'
+
+import type { CommandOptions } from './shareAction.types.js'
+
+export const shareAction: CommandModule<unknown, CommandOptions> = {
+  command: 'share [snapshot-name|snapshot-path]',
+  describe: 'Share a snapshot',
+  aliases: ['upload'],
+  // @ts-expect-error
+  builder: (y) => {
+    return y
+      .parserConfiguration({
+        'boolean-negation': false,
+      })
+      .option('no-encrypt', {
+        type: 'boolean',
+        default: false,
+        describe: 'Disable encryption',
+      })
+      .positional('snapshot-name|snapshot-path', {
+        describe: 'the unique name or path of the snapshot',
+        type: 'string',
+      })
+      .option('tags', {
+        describe: 'Filter snapshots by tags',
+        type: 'array',
+        coerce: (values: string[]) => {
+          return values.flatMap((v) => v.split(','))
+        },
+        default: [],
+      })
+      .option('latest', {
+        type: 'boolean',
+        default: false,
+        describe: 'Share the latest snapshot',
+      })
+  },
+  async handler(props) {
+    const { handler } = await import('./shareAction.handler.js')
+    await handler(props)
+  },
+}

cli/src/commands/snapshot/actions/share/shareAction.types.ts

@@ -0,0 +1,6 @@
+export type CommandOptions = {
+  latest: boolean
+  tags: string[]
+  snapshotName?: string
+  noEncrypt: boolean
+}

cli/src/commands/snapshot/actions/share/steps/encryptAndCompressTables.ts

@@ -0,0 +1,31 @@
+import fs from 'fs-extra'
+import fg from 'fast-glob'
+import pMap from 'p-map'
+
+import {
+  EncryptionPayload,
+  encryptAndCompressSnapshotFile,
+  getSnapshotFilePaths,
+} from '@snaplet/sdk/cli'
+
+export const encryptAndCompressTables = async ({
+  paths,
+  encryptionPayload,
+}: {
+  paths: ReturnType<typeof getSnapshotFilePaths>
+  encryptionPayload?: EncryptionPayload
+}) => {
+  const csvFiles = await fg(`*.csv`, { cwd: paths.tables, absolute: true })
+
+  await pMap(
+    csvFiles,
+    async (p) => {
+      await encryptAndCompressSnapshotFile(p, encryptionPayload)
+      await fs.unlink(p)
+    },
+    // There is no need to try to compress/encrypt all files at once it'll just blow up the CPU usage
+    // and create tons of child processes for nothing. Doing it in batches of 2 ensure that one single
+    // long table won't block the other smaller ones.
+    { concurrency: 2 }
+  )
+}