learning_ai_common_plat/packages/field-encrypt/src/dek-store-cosmos.ts

/**
 * @bytelyst/field-encrypt — Cosmos DB DEK store
 *
 * Production DEK store backed by Azure Cosmos DB.
 * Container: `_encryption_keys` (partition key: /dekId)
 *
 * Each WrappedDek is stored as a document with dekId as both id and partition key.
 */

import type { Container } from '@azure/cosmos';
import type { DekStore, WrappedDek } from './types.js';

interface CosmosDekDoc {
  id: string;
  dekId: string;
  wrappedKey: string;
  mekVersion: string;
  createdAt: string;
}

export class CosmosDekStore implements DekStore {
  constructor(private readonly container: Container) {}

  async get(dekId: string): Promise<WrappedDek | null> {
    try {
      const { resource } = await this.container.item(dekId, dekId).read<CosmosDekDoc>();
      if (!resource) return null;
      return {
        dekId: resource.dekId,
        wrappedKey: resource.wrappedKey,
        mekVersion: resource.mekVersion,
        createdAt: resource.createdAt,
      };
    } catch (err: unknown) {
      if (isNotFound(err)) return null;
      throw err;
    }
  }

  async put(dek: WrappedDek): Promise<void> {
    const doc: CosmosDekDoc = {
      id: dek.dekId,
      dekId: dek.dekId,
      wrappedKey: dek.wrappedKey,
      mekVersion: dek.mekVersion,
      createdAt: dek.createdAt,
    };
    await this.container.items.upsert(doc);
  }

  async listIds(): Promise<string[]> {
    const { resources } = await this.container.items
      .query<{ dekId: string }>('SELECT c.dekId FROM c')
      .fetchAll();
    return resources.map(r => r.dekId);
  }

  async delete(dekId: string): Promise<void> {
    try {
      await this.container.item(dekId, dekId).delete();
    } catch (err: unknown) {
      if (isNotFound(err)) return; // Already deleted — idempotent
      throw err;
    }
  }
}

function isNotFound(err: unknown): boolean {
  return (
    typeof err === 'object' &&
    err !== null &&
    'code' in err &&
    (err as { code: number }).code === 404
  );
}