ChromaDB Vector Collection Snapshot and AWS S3 Disaster Recovery Pipeline

id: chroma-vector-database-snapshot-and-s3-backup
namespace: company.ai

concurrency:
  limit: 1
  behavior: QUEUE

inputs:
  - id: chroma_host
    type: STRING
    displayName: ChromaDB Host Endpoint
    required: true
    defaults: "http://chroma.internal.net:8000"
    description: "Base HTTP URL for the persistent ChromaDB vector instance."

  - id: collection_name
    type: STRING
    displayName: Vector Collection Name
    required: true
    defaults: "knowledge-base-embeddings"
    description: "The targeted Chroma collection to export, validate, and back up."

  - id: s3_bucket
    type: STRING
    displayName: S3 Backup Bucket
    required: true
    defaults: "production-vector-backups"
    description: "AWS S3 target bucket to store the encrypted vector snapshot archive."

  - id: s3_region
    type: STRING
    displayName: AWS S3 Region
    required: true
    defaults: "us-east-1"
    description: "AWS Region hosting the backup bucket."

  - id: min_vector_count
    type: INT
    displayName: Minimum Vector Threshold
    required: true
    defaults: 1
    description: "Safety threshold: abort backup if collection has fewer vectors
      than this count to prevent corrupt empty backups."

tasks:
  - id: log_start
    type: io.kestra.plugin.core.log.Log
    message: "📦 Initiating snapshot and backup for ChromaDB collection '{{
      inputs.collection_name }}' on '{{ inputs.chroma_host }}'"

  - id: chroma_preflight_check
    type: io.kestra.plugin.scripts.shell.Commands
    description: "Query ChromaDB REST API v2 heartbeat and collection count"
    retry:
      type: constant
      interval: PT10S
      maxAttempts: 3
    commands:
      - |
        curl -fsS "{{ inputs.chroma_host }}/api/v2/heartbeat" > heartbeat.json
        curl -fsS "{{ inputs.chroma_host }}/api/v2/tenants/default_tenant/databases/default_database/collections/{{ inputs.collection_name }}/count" > count.txt
    outputFiles:
      - heartbeat.json
      - count.txt

  - id: validate_collection_health
    type: io.kestra.plugin.scripts.python.Script
    description: "Inspect collection count and enforce non-empty safety barrier"
    taskRunner:
      type: io.kestra.plugin.scripts.runner.docker.Docker
    containerImage: python:3.11-slim
    dependencies:
      - kestra
    inputFiles:
      count.txt: "{{ outputs.chroma_preflight_check.outputFiles['count.txt'] }}"
    script: |
      import os
      from kestra import Kestra

      with open("count.txt", "r") as f:
          raw_count = f.read().strip()

      try:
          vector_count = int(raw_count)
      except ValueError:
          vector_count = 0

      print(f"Verified Chroma collection vector count: {vector_count}")

      Kestra.outputs({
          "vector_count": vector_count,
          "is_valid": vector_count >= {{ inputs.min_vector_count }}
      })

  - id: check_empty_threshold
    type: io.kestra.plugin.core.flow.If
    description: "Fail fast if collection is empty or corrupted"
    condition: "{{ not outputs.validate_collection_health.vars.is_valid }}"
    then:
      - id: fail_empty_collection
        type: io.kestra.plugin.core.execution.Fail
        errorMessage: "Aborting backup: ChromaDB collection '{{ inputs.collection_name
          }}' reported {{ outputs.validate_collection_health.vars.vector_count
          }} vectors, which is below minimum threshold {{
          inputs.min_vector_count }}."

  - id: export_and_compress_collection
    type: io.kestra.plugin.scripts.python.Script
    description: "Stream Chroma vectors and metadata in paginated batches into a
      compressed parquet archive"
    taskRunner:
      type: io.kestra.plugin.scripts.runner.docker.Docker
    containerImage: python:3.11-slim
    dependencies:
      - requests==2.32.3
      - pyarrow==17.0.0
    env:
      CHROMA_URL: "{{ inputs.chroma_host }}"
      COLLECTION: "{{ inputs.collection_name }}"
    script: |
      import os
      import requests
      import pyarrow as pa
      import pyarrow.parquet as pq

      chroma_url = os.environ["CHROMA_URL"].rstrip("/")
      coll_name = os.environ["COLLECTION"]

      coll_url = f"{chroma_url}/api/v2/tenants/default_tenant/databases/default_database/collections/{coll_name}"
      coll_resp = requests.get(coll_url)
      coll_resp.raise_for_status()
      coll_id = coll_resp.json()["id"]

      limit = 1000
      offset = 0
      all_ids = []
      all_embeddings = []
      all_documents = []
      all_metadatas = []

      while True:
          get_url = f"{coll_url}/get"
          payload = {
              "limit": limit,
              "offset": offset,
              "include": ["embeddings", "documents", "metadatas"]
          }
          resp = requests.post(get_url, json=payload)
          resp.raise_for_status()
          batch = resp.json()

          ids = batch.get("ids", [])
          if not ids:
              break

          all_ids.extend(ids)
          all_embeddings.extend(batch.get("embeddings", []))
          all_documents.extend([doc if doc is not None else "" for doc in batch.get("documents", [])])
          all_metadatas.extend([str(m) if m is not None else "{}" for m in batch.get("metadatas", [])])

          offset += len(ids)
          print(f"Extracted {len(all_ids)} vectors so far...")

      table = pa.Table.from_arrays(
          [
              pa.array(all_ids, pa.string()),
              pa.array(all_documents, pa.string()),
              pa.array(all_metadatas, pa.string())
          ],
          names=["id", "document", "metadata"]
      )
      archive_name = f"chroma_{coll_name}_snapshot.parquet"
      pq.write_table(table, archive_name, compression="snappy")
      print(f"Successfully generated compressed snapshot: {archive_name} with {len(all_ids)} records.")
    outputFiles:
      - "*.parquet"

  - id: upload_snapshot_to_s3
    type: io.kestra.plugin.aws.s3.Upload
    description: "Archive the generated Parquet vector snapshot to secure AWS S3 storage"
    from: "{{ outputs.export_and_compress_collection.outputFiles['chroma_' ~
      inputs.collection_name ~ '_snapshot.parquet'] }}"
    bucket: "{{ inputs.s3_bucket }}"
    key: "chroma-backups/{{ inputs.collection_name }}/{{ execution.startDate |
      date('yyyy-MM-dd') }}/chroma_{{ inputs.collection_name }}_{{ execution.id
      }}.parquet"
    region: "{{ inputs.s3_region }}"
    accessKeyId: "{{ secret('AWS_ACCESS_KEY_ID') }}"
    secretKeyId: "{{ secret('AWS_SECRET_ACCESS_KEY') }}"

  - id: record_backup_state
    type: io.kestra.plugin.core.kv.Set
    description: "Persist latest backup metadata into Kestra KV Store for audit and
      disaster recovery discovery"
    key: "chroma_backup_{{ inputs.collection_name }}"
    value: |
      {
        "collection": "{{ inputs.collection_name }}",
        "vector_count": {{ outputs.validate_collection_health.vars.vector_count }},
        "s3_path": "s3://{{ inputs.s3_bucket }}/chroma-backups/{{ inputs.collection_name }}/{{ execution.startDate | date('yyyy-MM-dd') }}/chroma_{{ inputs.collection_name }}_{{ execution.id }}.parquet",
        "timestamp": "{{ execution.startDate }}",
        "execution_id": "{{ execution.id }}"
      }

  - id: notify_slack_success
    type: io.kestra.plugin.slack.notifications.SlackIncomingWebhook
    description: "Notify data platform and MLOps channels of completed vector snapshot"
    url: "{{ secret('SLACK_WEBHOOK_URL') }}"
    payload: |
      {
        "blocks": [
          {
            "type": "header",
            "text": {
              "type": "plain_text",
              "text": "✅ ChromaDB Vector Backup Succeeded",
              "emoji": true
            }
          },
          {
            "type": "section",
            "fields": [
              {
                "type": "mrkdwn",
                "text": "*Collection:* `{{ inputs.collection_name }}`"
              },
              {
                "type": "mrkdwn",
                "text": "*Vector Count:* `{{ outputs.validate_collection_health.vars.vector_count }}`"
              },
              {
                "type": "mrkdwn",
                "text": "*S3 Bucket:* `{{ inputs.s3_bucket }}`"
              },
              {
                "type": "mrkdwn",
                "text": "*Execution:* `{{ execution.id }}`"
              }
            ]
          }
        ]
      }

errors:
  - id: alert_on_backup_failure
    type: io.kestra.plugin.slack.notifications.SlackIncomingWebhook
    description: "Immediate dead-letter alerting if any stage of the vector backup
      flow fails"
    url: "{{ secret('SLACK_WEBHOOK_URL') }}"
    payload: |
      {
        "blocks": [
          {
            "type": "header",
            "text": {
              "type": "plain_text",
              "text": "🚨 ChromaDB Vector Backup FAILED",
              "emoji": true
            }
          },
          {
            "type": "section",
            "text": {
              "type": "mrkdwn",
              "text": "*Execution Failed:* `{{ execution.id }}` on collection `{{ inputs.collection_name }}`\nPlease verify ChromaDB connectivity and AWS S3 permissions."
            }
          }
        ]
      }