@effect-app/infra 4.0.0-beta.99 → 4.0.1

package/src/Store/Cosmos.ts

@@ -1,16 +1,25 @@
 /* eslint-disable @typescript-eslint/no-explicit-any */
 
-import { Array, Duration, Effect, Layer, type NonEmptyReadonlyArray, Option, pipe, Redacted, Struct } from "effect-app"
-import { toNonEmptyArray } from "effect-app/Array"
+import * as Array from "effect-app/Array"
+import type { NonEmptyReadonlyArray } from "effect-app/Array"
+import * as Effect from "effect-app/Effect"
+import * as Layer from "effect-app/Layer"
+import type { FieldValues } from "effect-app/Model/filter/types"
+import { type ComputedProjectionIrExpression, type RawQuery } from "effect-app/Model/query"
+import * as Option from "effect-app/Option"
+import { type FilterArgs, type PersistenceModelType, type StorageConfig, type Store, type StoreConfig, storeId, StoreMaker } from "effect-app/Store"
 import { dropUndefinedT, mutable } from "effect-app/utils"
-import { CosmosClient, CosmosClientLayer } from "../adapters/cosmos-client.js"
-import { OptimisticConcurrencyException } from "../errors.js"
-import { InfraLogger } from "../logger.js"
-import type { FieldValues } from "../Model/filter/types.js"
-import { type RawQuery } from "../Model/query.js"
-import { buildWhereCosmosQuery3, logQuery } from "./Cosmos/query.js"
-import { storeId } from "./Memory.js"
-import { type FilterArgs, type PersistenceModelType, type StorageConfig, type Store, type StoreConfig, StoreMaker } from "./service.js"
+import * as Duration from "effect/Duration"
+import { pipe } from "effect/Function"
+import * as Redacted from "effect/Redacted"
+import * as Struct from "effect/Struct"
+import { CosmosClient, CosmosClientLayer, createContainerIfNotExists } from "../cosmos-client.ts"
+import { DatabaseError, OptimisticConcurrencyException } from "../errors.ts"
+import { InfraLogger } from "../logger.ts"
+import { annotateCosmosResponse, annotateDb } from "../otel.ts"
+import { buildWhereCosmosQuery3, logQuery } from "./Cosmos/query.ts"
+import { decodeStoredMany, decodeStoredOption, makeJsonDocumentCodec, makeStoredDecode } from "./jsonDocument.ts"
+import { makeJsonLower } from "./utils.ts"
 
 const makeMapId =
   <IdKey extends keyof Encoded, Encoded extends FieldValues>(idKey: IdKey) => ({ [idKey]: id, ...e }: Encoded) => ({
@@ -22,547 +31,623 @@ const makeReverseMapId =
   ({ id, ...t }: PersistenceModelType<Omit<Encoded, IdKey> & { id: string }>) =>
     ({ ...t, [idKey]: id }) as any as PersistenceModelType<Encoded>
 
-class CosmosDbOperationError {
-  constructor(readonly message: string, readonly raw?: unknown) {}
-} // TODO: Retry operation when running into RU limit.
+// Cosmos statuses worth retrying: request timeout, throttle, retry-with,
+// internal, and service unavailable.
+const RETRYABLE_COSMOS_STATUS = new Set([408, 429, 449, 500, 503])
+const statusIsTransient = (code: number) => RETRYABLE_COSMOS_STATUS.has(code)
 
-function makeCosmosStore({ prefix }: StorageConfig) {
-  return Effect.gen(function*() {
-    const { db } = yield* CosmosClient
-    return {
-      make: <IdKey extends keyof Encoded, Encoded extends FieldValues, R = never, E = never>(
-        name: string,
-        idKey: IdKey,
-        seed?: Effect.Effect<Iterable<Encoded>, E, R>,
-        config?: StoreConfig<Encoded>
-      ) =>
-        Effect.gen(function*() {
-          const mapId = makeMapId<IdKey, Encoded>(idKey)
-          const mapReverseId = makeReverseMapId<IdKey, Encoded>(idKey)
-          type PM = PersistenceModelType<Encoded>
-          type PMCosmos = PersistenceModelType<Omit<Encoded, IdKey> & { id: string }>
-          const containerId = `${prefix}${name}`
-          yield* Effect.promise(() =>
-            db.containers.createIfNotExists(dropUndefinedT({
-              id: containerId,
-              uniqueKeyPolicy: config?.uniqueKeys
-                ? { uniqueKeys: config.uniqueKeys }
-                : undefined
-            }))
-          )
+// A thrown SDK error is transient if its status/code is retryable or its message
+// looks like a timeout / throttle / dropped connection.
+const thrownIsTransient = (e: unknown): boolean => {
+  const code = (e as any)?.code ?? (e as any)?.statusCode
+  if (typeof code === "number" && RETRYABLE_COSMOS_STATUS.has(code)) return true
+  const msg = e instanceof Error ? e.message : String(e)
+  return /request timed out|timed out|\btimeout\b|ETIMEDOUT|ECONNRESET|ECONNREFUSED|EAI_AGAIN|socket hang up|ServiceUnavailable|RequestTimeout|TooManyRequests|throttl/i
+    .test(msg)
+}
+
+// Wrap a Cosmos SDK promise so a rejection becomes a typed, serializable
+// `DatabaseError` (transient-classified) instead of a bare-`Error` defect — the
+// latter both bypasses retry and breaks JSON encoding of the failure channel.
+const tryCosmos = <A>(f: () => Promise<A>): Effect.Effect<A, DatabaseError> =>
+  Effect.tryPromise({
+    try: f,
+    catch: (e) =>
+      new DatabaseError({
+        message: `Cosmos request failed: ${e instanceof Error ? e.message : String(e)}`,
+        transient: thrownIsTransient(e),
+        cause: e
+      })
+  })
+
+const respBytes = (
+  resp: { diagnostics?: { clientSideRequestStatistics?: { totalResponsePayloadLengthInBytes?: number } } }
+) => resp.diagnostics?.clientSideRequestStatistics?.totalResponsePayloadLengthInBytes ?? 0
+
+const annotateFeed = (resp: {
+  resources: readonly unknown[]
+  requestCharge?: number
+  diagnostics?: { clientSideRequestStatistics?: { totalResponsePayloadLengthInBytes?: number } }
+}) =>
+  annotateCosmosResponse({
+    requestCharge: resp.requestCharge,
+    returnedRows: resp.resources.length,
+    responseBytes: respBytes(resp)
+  })
+
+const annotateItem = (resp: {
+  requestCharge?: number
+  statusCode?: number
+  diagnostics?: { clientSideRequestStatistics?: { totalResponsePayloadLengthInBytes?: number } }
+}) =>
+  annotateCosmosResponse({
+    requestCharge: resp.requestCharge,
+    statusCode: resp.statusCode,
+    responseBytes: respBytes(resp)
+  })
+
+const makeCosmosStore = Effect.fnUntraced(function*({ autoscaleMaxThroughput, prefix }: StorageConfig) {
+  const { db } = yield* CosmosClient
+  return {
+    make: Effect.fnUntraced(function*<IdKey extends keyof Encoded, Encoded extends FieldValues, R = never, E = never>(
+      name: string,
+      idKey: IdKey,
+      seed?: Effect.Effect<Iterable<Encoded>, E, R>,
+      config?: StoreConfig<Encoded>
+    ) {
+      const mapId = makeMapId<IdKey, Encoded>(idKey)
+      const mapReverseId = makeReverseMapId<IdKey, Encoded>(idKey)
+      const codec = makeJsonDocumentCodec<Encoded>(config?.schema)
+      const decodeStored = makeStoredDecode<Encoded>(config?.schema, config?.jitM)
+      const json = makeJsonLower(config)
+      const defaultValues = json.toJson(config?.defaultValues ?? {}) as Partial<Encoded>
+      // stored JSON -> defaultValues -> jitM -> JSON→Encoded decode
+      const fromStored = (raw: Encoded) => decodeStored({ ...defaultValues, ...mapReverseId(raw as any) })
+      type PM = PersistenceModelType<Encoded>
+      type PMCosmos = PersistenceModelType<Omit<Encoded, IdKey> & { id: string }>
+      const containerId = `${prefix}${name}`
+      yield* Effect.promise(() =>
+        createContainerIfNotExists(
+          db,
+          dropUndefinedT({
+            id: containerId,
+            uniqueKeyPolicy: config?.uniqueKeys
+              ? { uniqueKeys: config.uniqueKeys }
+              : undefined,
+            partitionKey: {
+              paths: ["/_partitionKey"],
+              version: 2 // support large partitionkeys so that the hash is not based on just the first 100 bytes!
+            }
+          }),
+          { autoscaleMaxThroughput }
+        )
+      )
 
-          const basePartitionKey = config?.partitionValue() ?? "primary"
-          const nsPrefix = (ns: string) => ns === "primary" ? "" : `${ns}::`
-          const nsPartitionValue = (ns: string, e?: Encoded) => {
-            const base = config?.partitionValue(e) ?? "primary"
-            return `${nsPrefix(ns)}${base}`
+      const basePartitionKey = config?.partitionValue() ?? "primary"
+      const nsPrefix = (ns: string) => ns === "primary" ? "" : `${ns}::`
+      const nsPartitionValue = (ns: string, e?: Encoded) => {
+        const base = config?.partitionValue(e) ?? "primary"
+        return `${nsPrefix(ns)}${base}`
+      }
+      const nsBasePartitionKey = (ns: string) => `${nsPrefix(ns)}${basePartitionKey}`
+      const resolveNamespace = !config?.allowNamespace
+        ? Effect.succeed("primary")
+        : storeId.pipe(Effect.map((namespace) => {
+          if (namespace !== "primary" && !config.allowNamespace!(namespace)) {
+            throw new Error(`Namespace ${namespace} not allowed!`)
           }
-          const nsBasePartitionKey = (ns: string) => `${nsPrefix(ns)}${basePartitionKey}`
-          const resolveNamespace = !config?.allowNamespace
-            ? Effect.succeed("primary")
-            : storeId.asEffect().pipe(Effect.map((namespace) => {
-              if (namespace !== "primary" && !config.allowNamespace!(namespace)) {
-                throw new Error(`Namespace ${namespace} not allowed!`)
-              }
-              return namespace
-            }))
+          return namespace
+        }))
 
-          const defaultValues = config?.defaultValues ?? {}
-          const container = db.container(containerId)
-          const bulk = container.items.bulk.bind(container.items)
-          const execBatch = container.items.batch.bind(container.items)
-          // TODO: move the marker to a separate container and get rid of the checks on every query
-          // then need to clean up the actual data.. perhaps first do with a config toggle to prescribe to it.
-          const importedMarkerId = containerId
+      const container = db.container(containerId)
+      const bulk = container.items.bulk.bind(container.items)
+      const execBatch = container.items.batch.bind(container.items)
+      // TODO: move the marker to a separate container and get rid of the checks on every query
+      // then need to clean up the actual data.. perhaps first do with a config toggle to prescribe to it.
+      const importedMarkerId = containerId
 
-          const ctx = yield* Effect.context<R>()
-          const seedCache = new Map<string, Effect.Effect<void>>()
-          const makeSeedEffect = (ns: string) => {
-            const markerId = ns === "primary" ? importedMarkerId : `${importedMarkerId}::${ns}`
-            return Effect
-              .promise(() =>
-                container
-                  .item(markerId, markerId)
-                  .read<{ id: string }>()
-                  .then(({ resource }) => Option.fromNullishOr(resource))
-              )
-              .pipe(
-                Effect.flatMap((marker) => {
-                  if (Option.isSome(marker)) return Effect.void
-                  return InfraLogger.logInfo(`Creating mock data for ${name} (namespace: ${ns})`).pipe(
-                    Effect.andThen(seed!),
-                    Effect.flatMap((m) =>
-                      Effect.flatMapOption(
-                        Effect.succeed(toNonEmptyArray([...m])),
-                        (a) => bulkSetInternal(a, ns).pipe(Effect.orDie)
-                      )
-                    ),
-                    Effect.andThen(
-                      Effect.promise(() =>
-                        container.items.create({
-                          _partitionKey: markerId,
-                          id: markerId,
-                          ttl: -1
-                        })
-                      )
-                    ),
-                    Effect.provide(ctx),
-                    Effect.orDie
+      const ctx = yield* Effect.context<R>()
+      const seedCache = new Map<string, Effect.Effect<void>>()
+      const makeSeedEffect = (ns: string) => {
+        const markerId = ns === "primary" ? importedMarkerId : `${importedMarkerId}::${ns}`
+        return Effect
+          .promise(() =>
+            container
+              .item(markerId, markerId)
+              .read<{ id: string }>()
+              .then(({ resource }) => Option.fromNullishOr(resource))
+          )
+          .pipe(
+            Effect.flatMap((marker) => {
+              if (Option.isSome(marker)) return Effect.void
+              return InfraLogger.logInfo(`Creating mock data for ${name} (namespace: ${ns})`).pipe(
+                Effect.andThen(seed!),
+                Effect.flatMap((m) =>
+                  Effect.flatMapOption(
+                    Effect.succeed(Option.liftPredicate(Array.isReadonlyArrayNonEmpty)([...m])),
+                    (a) => bulkSetInternal(a, ns).pipe(Effect.orDie)
+                  )
+                ),
+                Effect.andThen(
+                  Effect.promise(() =>
+                    container.items.create({
+                      _partitionKey: markerId,
+                      id: markerId,
+                      ttl: -1
+                    })
                   )
-                }),
-                Effect.withLogSpan(`Cosmos.seedCheck ${name} in ${ns} [effect-app/infra/Store]`),
-                Effect.withSpan("Cosmos.seed [effect-app/infra/Store]", { attributes: { name, namespace: ns } })
+                ),
+                Effect.provide(ctx),
+                Effect.orDie
               )
-          }
-          const seedNamespace = Effect.fn("seedNamespace")(function*(ns: string) {
-            if (!seed) return
-            let cached = seedCache.get(ns)
-            if (!cached) {
-              cached = yield* Effect.cached(makeSeedEffect(ns))
-              seedCache.set(ns, cached)
-            }
-            yield* cached
-          })
-          const resolveAndSeed = resolveNamespace.pipe(
-            Effect.tap((ns) => seedNamespace(ns))
+            }),
+            Effect.withLogSpan(`Cosmos.seedCheck ${name} in ${ns} [effect-app/infra/Store]`),
+            annotateDb({
+              operation: "seed",
+              system: "cosmosdb",
+              collection: containerId,
+              namespace: ns,
+              entity: name
+            })
           )
-
-          const bulkSetInternal = (items: NonEmptyReadonlyArray<PM>, ns: string) =>
-            Effect
-              .gen(function*() {
-                // TODO: disable batching if need atomicity
-                // we delay and batch to keep low amount of RUs
-                const b = [...items]
-                  .map(
-                    (x) =>
-                      [
-                        x,
-                        Option.match(Option.fromNullishOr(x._etag), {
-                          onNone: () =>
-                            dropUndefinedT({
-                              operationType: "Create" as const,
-                              resourceBody: {
-                                ...Struct.omit(x, ["_etag", idKey]),
-                                id: x[idKey],
-                                _partitionKey: nsPartitionValue(ns, x)
-                              }
-                              // don't use this or we get an error that the request and some item partition key dont match - makese no sense
-                              // partitionKey: config?.partitionValue(x)
-                            }),
-                          onSome: (eTag) =>
-                            dropUndefinedT({
-                              operationType: "Replace" as const,
-                              id: x[idKey],
-                              resourceBody: {
-                                ...Struct.omit(x, ["_etag", idKey]),
-                                id: x[idKey],
-                                _partitionKey: nsPartitionValue(ns, x)
-                              },
-                              ifMatch: eTag
-                              // don't use this or we get an error that the request and some item partition key dont match - makese no sense
-                              // partitionKey: config?.partitionValue(x)
-                            })
+      }
+      const seedNamespace = Effect.fn("seedNamespace")(function*(ns: string) {
+        if (!seed) return
+        let cached = seedCache.get(ns)
+        if (!cached) {
+          cached = yield* Effect.cached(Effect.uninterruptible(makeSeedEffect(ns)))
+          seedCache.set(ns, cached)
+        }
+        yield* cached
+      })
+      const bulkSetInternal = (items: NonEmptyReadonlyArray<PM>, ns: string) =>
+        Effect
+          .gen(function*() {
+            // TODO: disable batching if need atomicity
+            // we delay and batch to keep low amount of RUs
+            const b = [...items]
+              .map(
+                (x) =>
+                  [
+                    x,
+                    Option.match(Option.fromNullishOr(x._etag), {
+                      onNone: () =>
+                        dropUndefinedT({
+                          operationType: "Create" as const,
+                          resourceBody: {
+                            ...Struct.omit(codec.encode(x), ["_etag", idKey]),
+                            id: x[idKey],
+                            _partitionKey: nsPartitionValue(ns, x)
+                          }
+                          // don't use this or we get an error that the request and some item partition key dont match - makese no sense
+                          // partitionKey: config?.partitionValue(x)
+                        }),
+                      onSome: (eTag) =>
+                        dropUndefinedT({
+                          operationType: "Replace" as const,
+                          id: x[idKey],
+                          resourceBody: {
+                            ...Struct.omit(codec.encode(x), ["_etag", idKey]),
+                            id: x[idKey],
+                            _partitionKey: nsPartitionValue(ns, x)
+                          },
+                          ifMatch: eTag
+                          // don't use this or we get an error that the request and some item partition key dont match - makese no sense
+                          // partitionKey: config?.partitionValue(x)
                         })
-                      ] as const
-                  )
-                const batches = Array.chunksOf(b, config?.maxBulkSize ?? 10)
+                    })
+                  ] as const
+              )
+            const batches = Array.chunksOf(b, config?.maxBulkSize ?? 10)
 
-                const batchResult = yield* Effect.forEach(
-                  batches
-                    .map((x, i) => [i, x] as const),
-                  ([i, batch]) =>
+            const batchResult = yield* Effect.forEach(
+              batches
+                .map((x, i) => [i, x] as const),
+              ([i, batch]) =>
+                tryCosmos(() => bulk(batch.map(([, op]) => op)))
+                  .pipe(
                     Effect
-                      .promise(() => bulk(batch.map(([, op]) => op)))
-                      .pipe(
-                        Effect
-                          .delay(Duration.millis(i === 0 ? 0 : 150)),
-                        Effect
-                          .flatMap((responses) =>
-                            Effect.gen(function*() {
-                              const r = responses.find((x) =>
-                                x.statusCode === 412 || x.statusCode === 404 || x.statusCode === 409
-                              )
-                              if (r) {
-                                return yield* Effect.fail(
-                                  new OptimisticConcurrencyException(
-                                    {
-                                      type: name,
-                                      id: JSON.stringify(r.resourceBody?.["id"]),
-                                      code: r.statusCode,
-                                      raw: responses
-                                    }
-                                  )
-                                )
-                              }
-                              const r2 = responses.find(
-                                (x) => x.statusCode !== 424 && (x.statusCode > 299 || x.statusCode < 200)
-                              )
-                              if (r2) {
-                                return yield* Effect.die(
-                                  new CosmosDbOperationError(
-                                    "not able to update records: " + r2.statusCode,
-                                    responses
-                                  )
-                                )
-                              }
-                              const r3 = responses.find(
-                                (x) => x.statusCode > 299 || x.statusCode < 200
+                      .delay(Duration.millis(i === 0 ? 0 : 150)),
+                    Effect
+                      .flatMap((responses) =>
+                        Effect.gen(function*() {
+                          const r = responses.find((x) =>
+                            x.statusCode === 412 || x.statusCode === 404 || x.statusCode === 409
+                          )
+                          if (r) {
+                            return yield* Effect.fail(
+                              new OptimisticConcurrencyException(
+                                {
+                                  type: name,
+                                  id: JSON.stringify(r.resourceBody?.["id"]),
+                                  code: r.statusCode,
+                                  raw: responses
+                                }
                               )
-                              if (r3) {
-                                return yield* Effect.die(
-                                  new CosmosDbOperationError(
-                                    "not able to update records: " + r3.statusCode,
-                                    responses
-                                  )
-                                )
-                              }
-                              return batch.map(([e], i) => ({
-                                ...e,
-                                // eslint-disable-next-line @typescript-eslint/no-non-null-assertion
-                                _etag: responses[i]!.eTag
-                              }))
+                            )
+                          }
+                          const r2 = responses.find(
+                            (x) => x.statusCode !== 424 && (x.statusCode > 299 || x.statusCode < 200)
+                          )
+                          if (r2) {
+                            return yield* new DatabaseError({
+                              message: "not able to update records: " + r2.statusCode,
+                              transient: statusIsTransient(r2.statusCode),
+                              cause: responses
                             })
+                          }
+                          const r3 = responses.find(
+                            (x) => x.statusCode > 299 || x.statusCode < 200
                           )
+                          if (r3) {
+                            return yield* new DatabaseError({
+                              message: "not able to update records: " + r3.statusCode,
+                              transient: statusIsTransient(r3.statusCode),
+                              cause: responses
+                            })
+                          }
+                          return batch.map(([e], i) => ({
+                            ...e,
+                            // eslint-disable-next-line @typescript-eslint/no-non-null-assertion
+                            _etag: responses[i]!.eTag
+                          }))
+                        })
                       )
-                )
-
-                return batchResult.flat() as unknown as NonEmptyReadonlyArray<Encoded>
-              })
-              .pipe(
-                Effect.withSpan("Cosmos.bulkSet [effect-app/infra/Store]", {
-                  attributes: { "repository.container_id": containerId, "repository.model_name": name, namespace: ns }
-                }, { captureStackTrace: false })
-              )
-
-          const bulkSet = (items: NonEmptyReadonlyArray<PM>) =>
-            resolveAndSeed.pipe(Effect.flatMap((ns) => bulkSetInternal(items, ns)))
+                  )
+            )
 
-          const batchSet = (items: NonEmptyReadonlyArray<PM>) => {
-            return resolveAndSeed
-              .pipe(Effect.flatMap((ns) =>
-                Effect
-                  .suspend(() => {
-                    const batch = [...items].map(
-                      (x) =>
-                        [
-                          x,
-                          Option.match(Option.fromNullishOr(x._etag), {
-                            onNone: () => ({
-                              operationType: "Create" as const,
-                              resourceBody: {
-                                ...Struct.omit(x, ["_etag", idKey]),
-                                id: x[idKey],
-                                _partitionKey: nsPartitionValue(ns, x)
-                              }
-                              // don't use this or we get an error that the request and some item partition key dont match - makese no sense
-                              // partitionKey: config?.partitionValue(x)
-                            }),
-                            onSome: (eTag) => ({
-                              operationType: "Replace" as const,
-                              id: x[idKey],
-                              resourceBody: {
-                                ...Struct.omit(x, ["_etag", idKey]),
-                                id: x[idKey],
-                                _partitionKey: nsPartitionValue(ns, x)
-                              },
-                              // don't use this or we get an error that the request and some item partition key dont match - makese no sense
-                              // partitionKey: config?.partitionValue(x)
-                              ifMatch: eTag
-                            })
-                          })
-                        ] as const
-                    )
+            return batchResult.flat() as unknown as NonEmptyReadonlyArray<Encoded>
+          })
+          .pipe(
+            annotateDb({
+              operation: "bulkSet",
+              system: "cosmosdb",
+              collection: containerId,
+              namespace: ns,
+              entity: name
+            })
+          )
 
-                    const ex = batch.map(([, c]) => c)
+      const bulkSet = (items: NonEmptyReadonlyArray<PM>) =>
+        resolveNamespace.pipe(Effect.flatMap((ns) => bulkSetInternal(items, ns)))
 
-                    return Effect
-                      .promise(() => execBatch(ex, ex[0]?.resourceBody._partitionKey))
-                      .pipe(Effect.flatMap(Effect.fnUntraced(function*(x) {
-                        const result = x.result ?? []
-                        const firstFailed = result.find(
-                          (x: any) => x.statusCode > 299 || x.statusCode < 200
-                        )
-                        if (firstFailed) {
-                          const code = firstFailed.statusCode ?? 0
-                          if (code === 412 || code === 404 || code === 409) {
-                            return yield* new OptimisticConcurrencyException({ type: name, id: "batch", code })
+      const batchSet = (items: NonEmptyReadonlyArray<PM>) => {
+        return resolveNamespace
+          .pipe(Effect.flatMap((ns) =>
+            Effect
+              .suspend(() => {
+                const batch = [...items].map(
+                  (x) =>
+                    [
+                      x,
+                      Option.match(Option.fromNullishOr(x._etag), {
+                        onNone: () => ({
+                          operationType: "Create" as const,
+                          resourceBody: {
+                            ...Struct.omit(codec.encode(x), ["_etag", idKey]),
+                            id: x[idKey],
+                            _partitionKey: nsPartitionValue(ns, x)
                           }
+                          // don't use this or we get an error that the request and some item partition key dont match - makese no sense
+                          // partitionKey: config?.partitionValue(x)
+                        }),
+                        onSome: (eTag) => ({
+                          operationType: "Replace" as const,
+                          id: x[idKey],
+                          resourceBody: {
+                            ...Struct.omit(codec.encode(x), ["_etag", idKey]),
+                            id: x[idKey],
+                            _partitionKey: nsPartitionValue(ns, x)
+                          },
+                          // don't use this or we get an error that the request and some item partition key dont match - makese no sense
+                          // partitionKey: config?.partitionValue(x)
+                          ifMatch: eTag
+                        })
+                      })
+                    ] as const
+                )
 
-                          return yield* Effect.die(
-                            new CosmosDbOperationError("not able to update record: " + code)
-                          )
-                        }
+                const ex = batch.map(([, c]) => c)
 
-                        return batch.map(([e], i) => ({
-                          ...e,
-                          _etag: result[i]?.eTag
-                        })) as unknown as NonEmptyReadonlyArray<Encoded>
-                      })))
-                  })
-                  .pipe(Effect
-                    .withSpan("Cosmos.batchSet [effect-app/infra/Store]", {
-                      attributes: {
-                        "repository.container_id": containerId,
-                        "repository.model_name": name,
-                        namespace: ns
+                return tryCosmos(() => execBatch(ex, ex[0]?.resourceBody._partitionKey))
+                  .pipe(Effect.flatMap(Effect.fnUntraced(function*(x) {
+                    const result = x.result ?? []
+                    const firstFailed = result.find(
+                      (x: any) => x.statusCode > 299 || x.statusCode < 200
+                    )
+                    if (firstFailed) {
+                      const code = firstFailed.statusCode ?? 0
+                      if (code === 412 || code === 404 || code === 409) {
+                        return yield* new OptimisticConcurrencyException({ type: name, id: "batch", code })
                       }
-                    }, { captureStackTrace: false }))
-              ))
-          }
 
-          const s: Store<IdKey, Encoded> = {
-            queryRaw: <Out>(query: RawQuery<Encoded, Out>) =>
-              Effect
-                .all({ q: Effect.sync(() => query.cosmos({ name })), ns: resolveAndSeed })
-                .pipe(
-                  Effect.tap(({ q }) => logQuery(q)),
-                  Effect.flatMap(({ ns, q }) =>
-                    Effect
-                      .promise(() =>
-                        container
-                          .items
-                          .query<Out>(q, { partitionKey: nsBasePartitionKey(ns) })
-                          .fetchAll()
-                          .then(({ resources }) =>
-                            resources.map(
-                              (_) => ({ ...defaultValues, ...mapReverseId(_ as any) }) as Out
-                            )
-                          )
-                      )
-                      .pipe(
-                        Effect.withSpan("Cosmos.queryRaw [effect-app/infra/Store]", {
-                          attributes: {
-                            "repository.container_id": containerId,
-                            "repository.model_name": name,
-                            namespace: ns
-                          }
-                        }, { captureStackTrace: false })
-                      )
-                  )
-                ),
-            batchRemove: (ids, partitionKey?: string) =>
-              resolveAndSeed.pipe(Effect.flatMap((ns) =>
+                      return yield* new DatabaseError({
+                        message: "not able to update record: " + code,
+                        transient: statusIsTransient(code),
+                        cause: x
+                      })
+                    }
+
+                    return batch.map(([e], i) => ({
+                      ...e,
+                      _etag: result[i]?.eTag
+                    })) as unknown as NonEmptyReadonlyArray<Encoded>
+                  })))
+              })
+              .pipe(annotateDb({
+                operation: "batchSet",
+                system: "cosmosdb",
+                collection: containerId,
+                namespace: ns,
+                entity: name
+              }))
+          ))
+      }
+
+      const s: Store<IdKey, Encoded> = {
+        seedNamespace: (ns) => seedNamespace(ns),
+
+        queryRaw: <Out>(query: RawQuery<Encoded, Out>) =>
+          Effect
+            .all({ q: Effect.sync(() => query.cosmos({ name })), ns: resolveNamespace })
+            .pipe(
+              Effect.tap(({ q }) => logQuery(q)),
+              Effect.flatMap(({ ns, q }) =>
                 Effect
-                  .promise(() =>
-                    execBatch(
-                      mutable(ids.map((id) =>
-                        dropUndefinedT({
-                          operationType: "Delete" as const,
-                          id
-                          // don't use this or we get an error that the request and some item partition key dont match - makese no sense
-                          // partitionKey: config?.partitionValue({ [idKey]: id } as Encoded)
-                        })
-                      )),
-                      partitionKey ?? nsBasePartitionKey(ns)
+                  .gen(function*() {
+                    const response = yield* tryCosmos(() =>
+                      container.items.query<Out>(q, { partitionKey: nsBasePartitionKey(ns) }).fetchAll()
                     )
-                  )
+                    yield* annotateFeed(response)
+                    return response.resources.map(
+                      (_) => ({ ...defaultValues, ...mapReverseId(_ as any) }) as Out
+                    )
+                  })
                   .pipe(
-                    Effect.withSpan("Cosmos.batchRemove [effect-app/infra/Store]", {
-                      attributes: {
-                        "repository.container_id": containerId,
-                        "repository.model_name": name,
-                        namespace: ns
-                      }
-                    }, { captureStackTrace: false })
+                    annotateDb({
+                      operation: "queryRaw",
+                      system: "cosmosdb",
+                      collection: containerId,
+                      namespace: ns,
+                      entity: name,
+                      query: q.query
+                    })
                   )
-              )),
-            all: Effect
-              .all({
-                q: Effect.sync(() => ({
-                  query: `SELECT * FROM ${name}`,
-                  parameters: []
-                })),
-                ns: resolveAndSeed
-              })
+              )
+            ),
+        batchRemove: (ids, partitionKey?: string) =>
+          resolveNamespace.pipe(Effect.flatMap((ns) =>
+            tryCosmos(() =>
+              execBatch(
+                mutable(ids.map((id) =>
+                  dropUndefinedT({
+                    operationType: "Delete" as const,
+                    id
+                    // don't use this or we get an error that the request and some item partition key dont match - makese no sense
+                    // partitionKey: config?.partitionValue({ [idKey]: id } as Encoded)
+                  })
+                )),
+                partitionKey ?? nsBasePartitionKey(ns)
+              )
+            )
               .pipe(
-                Effect.tap(({ q }) => logQuery(q)),
-                Effect.flatMap(({ ns, q }) =>
-                  Effect
-                    .promise(() =>
-                      container
-                        .items
-                        .query<PMCosmos>(q, { partitionKey: nsBasePartitionKey(ns) })
-                        .fetchAll()
-                        .then(({ resources }) =>
-                          resources.map(
-                            (_) => ({ ...defaultValues, ...mapReverseId(_) })
-                          )
-                        )
-                    )
-                    .pipe(
-                      Effect.withSpan("Cosmos.all [effect-app/infra/Store]", {
-                        attributes: {
-                          "repository.container_id": containerId,
-                          "repository.model_name": name,
-                          namespace: ns
-                        }
-                      }, { captureStackTrace: false })
-                    )
-                )
-              ),
-            /**
-             * May return duplicate results for "join_find", when matching more than once.
-             */
-            filter: <U extends keyof Encoded = never>(
-              f: FilterArgs<Encoded, U>
-            ) => {
-              const skip = f?.skip
-              const limit = f?.limit
-              const filter = f.filter
-              type M = U extends undefined ? Encoded : Pick<Encoded, U>
-              return Effect
-                .all({
-                  q: Effect.sync(() =>
-                    buildWhereCosmosQuery3(
-                      idKey,
-                      filter ? [{ t: "where-scope", result: filter, relation: "some" }] : [],
-                      name,
-                      defaultValues,
-                      f.select as
-                        | NonEmptyReadonlyArray<string | { key: string; subKeys: readonly string[] }>
-                        | undefined,
-                      f.order as NonEmptyReadonlyArray<{ key: string; direction: "ASC" | "DESC" }> | undefined,
-                      skip,
-                      limit
-                    )
-                  ),
-                  ns: resolveAndSeed
+                annotateDb({
+                  operation: "batchRemove",
+                  system: "cosmosdb",
+                  collection: containerId,
+                  namespace: ns,
+                  entity: name
+                })
+              )
+          )),
+        all: Effect
+          .all({
+            q: Effect.sync(() => ({
+              query: `SELECT * FROM ${name}`,
+              parameters: []
+            })),
+            ns: resolveNamespace
+          })
+          .pipe(
+            Effect.tap(({ q }) => logQuery(q)),
+            Effect.flatMap(({ ns, q }) =>
+              Effect
+                .gen(function*() {
+                  const response = yield* tryCosmos(() =>
+                    container.items.query<PMCosmos>(q, { partitionKey: nsBasePartitionKey(ns) }).fetchAll()
+                  )
+                  yield* annotateFeed(response)
+                  return yield* Effect.fromResult(
+                    decodeStoredMany(response.resources, (_) => fromStored(_ as unknown as Encoded))
+                  )
                 })
                 .pipe(
-                  Effect.tap(({ q }) => logQuery(q)),
+                  annotateDb({
+                    operation: "all",
+                    system: "cosmosdb",
+                    collection: containerId,
+                    namespace: ns,
+                    entity: name,
+                    query: q.query
+                  })
+                )
+            )
+          ),
+        /**
+         * May return duplicate results for "join_find", when matching more than once.
+         */
+        filter: <U extends keyof Encoded = never>(
+          f: FilterArgs<Encoded, U>
+        ) => {
+          const skip = f?.skip
+          const limit = f?.limit
+          const filter = f.filter
+          type M = U extends undefined ? Encoded : Pick<Encoded, U>
+          return Effect
+            .all({
+              q: Effect.sync(() =>
+                buildWhereCosmosQuery3(
+                  idKey,
+                  filter ? [{ t: "where-scope", result: filter, relation: "some" }] : [],
+                  name,
+                  defaultValues,
+                  f.select as
+                    | NonEmptyReadonlyArray<
+                      string | {
+                        key: string
+                        subKeys: readonly string[]
+                      } | {
+                        key: string
+                        computed: ComputedProjectionIrExpression
+                      }
+                    >
+                    | undefined,
+                  f.order as NonEmptyReadonlyArray<{ key: string; direction: "ASC" | "DESC" }> | undefined,
+                  skip,
+                  limit,
+                  json
+                )
+              ),
+              ns: resolveNamespace
+            })
+            .pipe(
+              Effect.tap(({ q }) => logQuery(q)),
+              Effect
+                .flatMap(({ ns, q }) =>
                   Effect
-                    .flatMap(({ ns, q }) =>
-                      Effect
-                        .promise(() =>
-                          f.select
-                            ? container
-                              .items
-                              .query<M>(q, { partitionKey: nsBasePartitionKey(ns) })
-                              .fetchAll()
-                              .then(({ resources }) =>
-                                resources.map((_) => ({
-                                  ...pipe(
-                                    defaultValues,
-                                    Struct.pick(f.select!.filter((_) => typeof _ === "string") as never[])
-                                  ),
-                                  ...mapReverseId(_ as any)
-                                }))
-                              )
-                            : container
-                              .items
-                              .query<{ f: M }>(q, { partitionKey: nsBasePartitionKey(ns) })
-                              .fetchAll()
-                              .then(({ resources }) =>
-                                resources.map(({ f }) => ({ ...defaultValues, ...mapReverseId(f as any) }) as any)
-                              )
-                        )
-                        .pipe(
-                          Effect.withSpan("Cosmos.filter [effect-app/infra/Store]", {
-                            attributes: {
-                              "repository.container_id": containerId,
-                              "repository.model_name": name,
-                              namespace: ns
-                            }
-                          }, { captureStackTrace: false })
+                    .gen(function*() {
+                      if (f.select) {
+                        const response = yield* tryCosmos(() =>
+                          container.items.query<M>(q, { partitionKey: nsBasePartitionKey(ns) }).fetchAll()
                         )
+                        yield* annotateFeed(response)
+                        return response.resources.map((_) => ({
+                          ...pipe(
+                            defaultValues,
+                            Struct.pick(f.select!.filter((_) => typeof _ === "string") as never[])
+                          ),
+                          ...mapReverseId(_ as any)
+                        }))
+                      }
+                      const response = yield* tryCosmos(() =>
+                        container.items.query<{ f: M }>(q, { partitionKey: nsBasePartitionKey(ns) }).fetchAll()
+                      )
+                      yield* annotateFeed(response)
+                      const decoded = yield* Effect.fromResult(
+                        decodeStoredMany(response.resources, ({ f }) => fromStored(f as Encoded))
+                      )
+                      return decoded as any
+                    })
+                    .pipe(
+                      annotateDb({
+                        operation: "filter",
+                        system: "cosmosdb",
+                        collection: containerId,
+                        namespace: ns,
+                        entity: name,
+                        query: q.query
+                      })
                     )
                 )
-            },
-            find: (id) =>
-              resolveAndSeed.pipe(Effect.flatMap((ns) =>
+            )
+        },
+        find: (id) =>
+          resolveNamespace.pipe(Effect.flatMap((ns) =>
+            Effect
+              .gen(function*() {
+                const response = yield* tryCosmos(() =>
+                  container
+                    .item(id, nsPartitionValue(ns, { [idKey]: id } as Encoded))
+                    .read<Encoded>()
+                )
+                yield* annotateItem(response)
+                return yield* Effect.fromResult(
+                  decodeStoredOption(Option.fromNullishOr(response.resource), fromStored)
+                )
+              })
+              .pipe(annotateDb({
+                operation: "find",
+                system: "cosmosdb",
+                collection: containerId,
+                namespace: ns,
+                entity: name,
+                extra: {
+                  "azure.cosmosdb.operation.partition_key": nsPartitionValue(ns, { [idKey]: id } as Encoded),
+                  "app.entity.id": id
+                }
+              }))
+          )),
+        set: (e) =>
+          resolveNamespace.pipe(Effect.flatMap((ns) =>
+            // Annotate the result so the create/replace promise union collapses to
+            // one type (otherwise `A` infers as `unknown` and the chain degrades).
+            tryCosmos((): Promise<{ statusCode: number; etag: string }> =>
+              Option
+                .match(
+                  Option
+                    .fromNullishOr(e._etag),
+                  {
+                    onNone: () =>
+                      container.items.create({
+                        ...mapId(codec.encode(e)),
+                        _partitionKey: nsPartitionValue(ns, e)
+                      }),
+                    onSome: (eTag) =>
+                      container.item(e[idKey], nsPartitionValue(ns, e)).replace(
+                        { ...mapId(codec.encode(e)), _partitionKey: nsPartitionValue(ns, e) },
+                        {
+                          accessCondition: {
+                            type: "IfMatch",
+                            condition: eTag
+                          }
+                        }
+                      )
+                  }
+                )
+            )
+              .pipe(
                 Effect
-                  .promise(() =>
-                    container
-                      .item(id, nsPartitionValue(ns, { [idKey]: id } as Encoded))
-                      .read<Encoded>()
-                      .then(({ resource }) =>
-                        Option.fromNullishOr(resource).pipe(
-                          Option.map((_) => ({ ...defaultValues, ...mapReverseId(_) }))
-                        )
+                  .flatMap((x): Effect.Effect<PM, OptimisticConcurrencyException | DatabaseError> => {
+                    if (x.statusCode === 412 || x.statusCode === 404 || x.statusCode === 409) {
+                      return Effect.fail(
+                        new OptimisticConcurrencyException({ type: name, id: e[idKey], code: x.statusCode })
                       )
-                  )
-                  .pipe(Effect
-                    .withSpan("Cosmos.find [effect-app/infra/Store]", {
-                      attributes: {
-                        "repository.container_id": containerId,
-                        "repository.model_name": name,
-                        partitionValue: nsPartitionValue(ns, { [idKey]: id } as Encoded),
-                        namespace: ns,
-                        id
-                      }
-                    }, { captureStackTrace: false }))
-              )),
-            set: (e) =>
-              resolveAndSeed.pipe(Effect.flatMap((ns) =>
-                Option
-                  .match(
-                    Option
-                      .fromNullishOr(e._etag),
-                    {
-                      onNone: () =>
-                        Effect.promise(() =>
-                          container.items.create({
-                            ...mapId(e),
-                            _partitionKey: nsPartitionValue(ns, e)
-                          })
-                        ),
-                      onSome: (eTag) =>
-                        Effect.promise(() =>
-                          container.item(e[idKey], nsPartitionValue(ns, e)).replace(
-                            { ...mapId(e), _partitionKey: nsPartitionValue(ns, e) },
-                            {
-                              accessCondition: {
-                                type: "IfMatch",
-                                condition: eTag
-                              }
-                            }
-                          )
-                        )
                     }
-                  )
-                  .pipe(
-                    Effect
-                      .flatMap((x) => {
-                        if (x.statusCode === 412 || x.statusCode === 404 || x.statusCode === 409) {
-                          return Effect.fail(
-                            new OptimisticConcurrencyException({ type: name, id: e[idKey], code: x.statusCode })
-                          )
-                        }
-                        if (x.statusCode > 299 || x.statusCode < 200) {
-                          return Effect.die(
-                            new CosmosDbOperationError(
-                              "not able to update record: " + x.statusCode
-                            )
-                          )
-                        }
-                        return Effect.sync(() => ({
-                          ...e,
-                          _etag: x.etag
-                        }))
-                      }),
-                    Effect
-                      .withSpan("Cosmos.set [effect-app/infra/Store]", {
-                        attributes: {
-                          "repository.container_id": containerId,
-                          "repository.model_name": name,
-                          namespace: ns,
-                          id: e[idKey]
-                        }
-                      }, { captureStackTrace: false })
-                  )
-              )),
-            batchSet,
-            bulkSet
-          }
+                    if (x.statusCode > 299 || x.statusCode < 200) {
+                      return Effect.fail(
+                        new DatabaseError({
+                          message: "not able to update record: " + x.statusCode,
+                          transient: statusIsTransient(x.statusCode),
+                          cause: x
+                        })
+                      )
+                    }
+                    return Effect.sync(() => ({
+                      ...e,
+                      _etag: x.etag
+                    }))
+                  }),
+                annotateDb({
+                  operation: "set",
+                  system: "cosmosdb",
+                  collection: containerId,
+                  namespace: ns,
+                  entity: name,
+                  extra: { "app.entity.id": e[idKey] }
+                })
+              )
+          )),
+        batchSet,
+        bulkSet
+      }
 
-          // Eagerly seed primary namespace on initialization
-          yield* seedNamespace("primary")
+      // Eagerly seed primary namespace on initialization
+      yield* seedNamespace("primary")
 
-          return s
-        })
-    }
-  })
-}
+      return s
+    })
+  }
+})
 
 export function CosmosStoreLayer(cfg: StorageConfig) {
   return StoreMaker