021-python-clients 0.1.0__py3-none-any.whl → 0.2.0__py3-none-any.whl

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
@@ -1,7 +1,7 @@
1
1
  Metadata-Version: 2.5
2
2
  Name: 021-python-clients
3
- Version: 0.1.0
4
- Summary: Observed clients for Redis, SQS, QStash, S3, HTTP, email and KMS, built on 021-python-observability
3
+ Version: 0.2.0
4
+ Summary: Observed clients for Redis, SQS, QStash, S3, Vercel Blob, HTTP, email and KMS, built on 021-python-observability
5
5
  Project-URL: Repository, https://github.com/from021-io/python-clients
6
6
  Project-URL: Changelog, https://github.com/from021-io/python-clients/blob/main/CHANGELOG.md
7
7
  Project-URL: Contract, https://github.com/from021-io/python-clients/blob/main/CONTRACT.md
@@ -29,6 +29,9 @@ Provides-Extra: httpx2
29
29
  Requires-Dist: httpx2<3,>=2; extra == 'httpx2'
30
30
  Provides-Extra: redis
31
31
  Requires-Dist: redis<9,>=6; extra == 'redis'
32
+ Provides-Extra: vercel-blob
33
+ Requires-Dist: httpx2<3,>=2; extra == 'vercel-blob'
34
+ Requires-Dist: vercel-blob<0.6,>=0.5; extra == 'vercel-blob'
32
35
  Description-Content-Type: text/markdown
33
36
 
34
37
  # 021-python-clients
@@ -40,20 +43,21 @@ dependency call gets the same span, the same metrics under a stable dependency n
40
43
  the same failure type.
41
44
 
42
45
  ```bash
43
- uv add "021-python-clients[redis,aws,httpx2]" # pick the extras you use
46
+ uv add "021-python-clients[redis,aws,httpx2,vercel-blob]" # pick the extras you use
44
47
  ```
45
48
 
46
49
  | Module | Extra | What it gives |
47
50
  | --- | --- | --- |
48
51
  | `clients_021.http` | none | `HttpClient`: one httpx client per dependency, timeouts, retries on idempotent methods |
49
52
  | `clients_021.httpx2` | `httpx2` | `ObservedTransport` for SDKs built on httpx2 (Anthropic, OpenAI) |
50
- | `clients_021.redis` | `redis` | `RedisClient`: JSON with TTL, one-time reads, windowed counters, ping |
53
+ | `clients_021.cache.redis` | `redis` | `RedisClient`: JSON with TTL, one-time reads, windowed counters, ping |
51
54
  | `clients_021.queue.qstash` | none | `QStashClient`: JSON publishing, deduplication, delivery options and forwarded trace context |
52
55
  | `clients_021.queue.sqs` | `aws` | `SqsClient`: send, receive, delete, visibility, trace carried in the message |
53
- | `clients_021.s3` | `aws` | `S3Client`: put, get, head, delete, presigned URLs |
56
+ | `clients_021.storage.s3` | `aws` | `S3Client`: put, get, head, delete, presigned URLs |
57
+ | `clients_021.storage.vercel_blob` | `vercel-blob` | `VercelBlobClient`: put, get, head and conditional delete for public/private stores |
54
58
  | `clients_021.email` | none | `EmailMessage`, the `EmailSender` protocol, `ResendSender` (Resend's HTTP API) |
55
59
  | `clients_021.email.ses` | `aws` | `SesSender`: the same protocol over Amazon SES v2 |
56
- | `clients_021.kms` | `aws` | `KmsClient`: encrypt, decrypt and data keys bound to an encryption context; AES-256-GCM envelopes |
60
+ | `clients_021.crypto.kms` | `aws` | `KmsClient`: encrypt, decrypt and data keys bound to an encryption context; AES-256-GCM envelopes |
57
61
 
58
62
  Every client takes the service's `Observability` (from `setup_observability`) and a
59
63
  `dependency` name, which becomes the `dependency` label of its metrics.
@@ -61,11 +65,16 @@ Every client takes the service's `Observability` (from `setup_observability`) an
61
65
  ## Organization
62
66
 
63
67
  Providers with the same purpose live under one capability package: `email/` contains
64
- Resend and SES; `queue/` contains SQS and QStash. Import queue clients explicitly:
68
+ Resend and SES; `queue/` contains SQS and QStash; `storage/` contains S3 and Vercel Blob; `cache/`
69
+ contains Redis; `crypto/` contains KMS. HTTP transports remain at the root. Import providers
70
+ explicitly:
65
71
 
66
72
  ```python
67
73
  from clients_021.queue.sqs import SqsClient
68
74
  from clients_021.queue.qstash import QStashClient
75
+ from clients_021.storage.s3 import S3Client
76
+ from clients_021.cache.redis import RedisClient
77
+ from clients_021.crypto.kms import KmsClient
69
78
  ```
70
79
 
71
80
  SQS supports polling, receipts and visibility changes; QStash delivers to an HTTP endpoint.
@@ -74,16 +83,20 @@ because both implement the same send operation. A shared package does not imply
74
83
  provider implements an interchangeable consumption protocol. Importing `queue` or QStash
75
84
  does not load the optional AWS SDK.
76
85
 
86
+ The 0.1.0 imports `clients_021.redis`, `clients_021.s3` and `clients_021.kms` remain
87
+ compatible aliases for the same classes and result types. Importing a capability package
88
+ does not load its optional SDKs.
89
+
77
90
  ## What every call gets
78
91
 
79
92
  - A CLIENT span named `<dependency> <operation>` (`HTTP <dependency>` for HTTP), with the
80
93
  system attribute (`db.system.name=redis`, `messaging.system=aws_sqs|resend|aws_ses`,
81
- `rpc.service=S3|KMS`).
94
+ `rpc.service=S3|KMS`, `rpc.system=vercel-blob`).
82
95
  - A counter and a duration histogram: `db.client.operation.*` (Redis, shared with SQL),
83
- `messaging.client.operation.*` (SQS, QStash, email), `storage.client.operation.*` (S3),
96
+ `messaging.client.operation.*` (SQS, QStash, email), `storage.client.operation.*` (S3, Vercel Blob),
84
97
  `crypto.client.operation.*` (KMS), `http.client.request.*` (HTTP, and the HTTP request
85
98
  under QStash and Resend). Labels are bounded; see [CONTRACT.md](CONTRACT.md).
86
- - Redis, SQS, QStash, S3, email and KMS emit a `dependency.call_failed` warning on failure,
99
+ - Redis, SQS, QStash, S3, Vercel Blob, email and KMS emit a `dependency.call_failed` warning on failure,
87
100
  with dependency, operation and error type only. HTTP transports emit spans and metrics.
88
101
  - `DependencyUnavailable(dependency, operation, reason)` when the service timed out,
89
102
  refused the connection, throttled or failed with a 5xx. Map it to a 503. Errors that mean
@@ -100,8 +113,8 @@ them from a tenant, user or request. All clients apply observability's identifie
100
113
 
101
114
  ```python
102
115
  from clients_021 import DependencyUnavailable, HttpClient
103
- from clients_021.redis import RedisClient
104
- from clients_021.s3 import S3Client
116
+ from clients_021.cache.redis import RedisClient
117
+ from clients_021.storage.s3 import S3Client
105
118
  from clients_021.queue.sqs import SqsClient
106
119
 
107
120
  redis = RedisClient.from_url(obs, settings.redis_url)
@@ -123,6 +136,59 @@ http = HttpClient(obs, timeout_seconds=10)
123
136
  response = await http.request("GET", jwks_url, dependency="identity-provider")
124
137
  ```
125
138
 
139
+ ### Vercel Blob
140
+
141
+ ```python
142
+ from clients_021.storage.vercel_blob import VercelBlobClient
143
+
144
+ files = VercelBlobClient(
145
+ obs,
146
+ token=settings.blob_token,
147
+ store_id=settings.blob_store_id,
148
+ access="private", # Required; must match the store. Public is also supported.
149
+ dependency="documents",
150
+ timeout_seconds=10,
151
+ )
152
+ try:
153
+ uploaded = await files.put("reports/report.pdf", pdf_bytes, content_type="application/pdf")
154
+ content = await files.get(uploaded.url) # bytes, or None for a missing object
155
+ info = await files.head(uploaded.url) # metadata, or None for a missing object
156
+ if info is not None:
157
+ await files.delete(uploaded.url, if_match=info.etag)
158
+ finally:
159
+ await files.aclose()
160
+ ```
161
+
162
+ Uses the [official Python Blob SDK](https://vercel.com/docs/vercel-blob/python-sdk-reference),
163
+ `vercel-blob>=0.5,<0.6`, installed with the `vercel-blob` extra. This is the current functional
164
+ SDK API; it does not depend on the earlier `AsyncBlobClient` API or the full `vercel` bundle.
165
+ Credentials and store access are explicit; no environment or default SDK session is used
166
+ for requests. A read-write token's store must match `store_id` (the `store_` prefix is allowed).
167
+
168
+ `put` returns `BlobUpload` (the SDK's `PutResult`); `head` returns `BlobInfo` (its `HeadResult`).
169
+ These are provider-specific types, not a common S3 protocol. Uploads require a pathname;
170
+ reads, metadata and deletes accept a pathname or a Vercel delivery URL. The SDK validates
171
+ those URLs and does not follow redirects. `get` buffers the entire body in memory.
172
+
173
+ Each operation has one attempt and a finite total deadline, including body consumption.
174
+ The client reuses its connection pool across calls and closes it through `aclose()`,
175
+ including an injected `httpx2.AsyncBaseTransport`. Call `aclose()` only after in-flight calls
176
+ finish. Calls after close fail locally. SDK sessions are scoped to each operation so
177
+ concurrent clients keep their credentials separate. No trace headers are sent by default.
178
+
179
+ `allow_overwrite=False` and `add_random_suffix=False` by default. Both can be explicitly
180
+ enabled, along with `cache_control_max_age`. A timeout may follow a successful upload:
181
+ retrying with a random suffix can create another object; overwriting a stable pathname can
182
+ race with other writers. Deleting a missing object succeeds. `if_match` guards a delete
183
+ against a changed ETag. CDN caches can retain deleted content briefly.
184
+
185
+ 429, 5xx, transport failures and deadlines become `DependencyUnavailable`. Authentication,
186
+ missing stores, precondition failures, invalid inputs and malformed responses retain SDK
187
+ exceptions; only a missing object becomes `None` for `get`/`head`. Cancellation propagates.
188
+ The storage operation includes the complete request and response body; this client does
189
+ not emit an additional HTTP metric. Signed URLs, multipart and streaming APIs stay outside
190
+ this initial client.
191
+
126
192
  ### QStash
127
193
 
128
194
  ```python
@@ -207,7 +273,7 @@ The encryption context is required on every call, so a ciphertext only decrypts
207
273
  purpose it was made for:
208
274
 
209
275
  ```python
210
- from clients_021.kms import KmsClient
276
+ from clients_021.crypto.kms import KmsClient
211
277
 
212
278
  keys = KmsClient.create(obs, settings.session_key_arn, region="us-east-1")
213
279
  context = {"purpose": "refresh-token", "tenant": tenant_id}
@@ -292,6 +358,8 @@ Retry policies are explicit:
292
358
  `DependencyUnavailable`; other rejected statuses and invalid acknowledgements raise
293
359
  `QStashResponseError` with a status code and no provider body. Delivery timeouts and
294
360
  delays use integer seconds and remain subject to the provider plan's limits.
361
+ - Vercel Blob sends each operation once, with a finite total deadline (10 seconds by default).
362
+ An injected transport must also avoid automatic retries of uncertain writes.
295
363
  - httpx2 SDKs own their retries; `ObservedTransport` adds no extra retry loop.
296
364
 
297
365
  `DependencyUnavailable` signals a transient dependency failure, **not permission to replay
@@ -0,0 +1,28 @@
1
+ clients_021/__init__.py,sha256=sF_9ZsxPumQbZ-9iCXoLP8JcU_tITDT6Xk1mXF9PVuQ,816
2
+ clients_021/_aws.py,sha256=sj2kCfod-vz9J93UyeVzNVLGK3M25yipKdW15KyIkw4,1722
3
+ clients_021/_observe.py,sha256=Xjg5Qck-e8Gt5NQHOND6nzr7DvTn46PmBdccubiD3Dk,3792
4
+ clients_021/errors.py,sha256=GS_fN8mjr7MCOkhqPhH6LweIvQ20Vd29U-aZY8zBeJQ,972
5
+ clients_021/http.py,sha256=-AsGLfngd-eEJnTuZEPdVBI9myYiuggmnrvjQjaMpT8,5045
6
+ clients_021/httpx2.py,sha256=w248IcEtDZtsEDkkvj4Bpfl7MU7IJ08EHqHWhryC_3k,3880
7
+ clients_021/kms.py,sha256=YcZXURpqIZTkjcP0fpWMD6X_adMAYwN7zAJNhUImR4s,183
8
+ clients_021/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
9
+ clients_021/redis.py,sha256=xC3tji_ldJAwaBntkZTuPhNHLIzl3F3nDzkikMrbJCc,169
10
+ clients_021/s3.py,sha256=6Q5_1emngSB00IwvBdrmV3iAg-Q_Ndco1HGzYObL5FI,231
11
+ clients_021/cache/__init__.py,sha256=3OyQQ-uZeghpgJvvJ2k56_t50q3CCwbAp3dBrCMMcSU,76
12
+ clients_021/cache/redis.py,sha256=DRXCMc8MqY8KsZ50j4X4mUsd0lvvmP-UDuUD2wJXyW8,4932
13
+ clients_021/crypto/__init__.py,sha256=qESFBPQBKcHhdFXqeJzSyEP41yveG7rFrxnpITmqlNU,77
14
+ clients_021/crypto/kms.py,sha256=My5Xyg0bOiT9YHYWojQSszpREEh3bSYMMB26wy6fNEI,7261
15
+ clients_021/email/__init__.py,sha256=ivmdu9hhNDwPqBn3G7uNG0Zpc3DEMnlnWBVVoVQ-fQA,790
16
+ clients_021/email/_message.py,sha256=ulTX2cXPxSVaq_5FVF0e6glePVJ7wiO5P3M4WUWN3N0,3258
17
+ clients_021/email/resend.py,sha256=ugENkEzw9eUlYih3BZcCVQswTF1MDikPgZyuqdapCjc,5347
18
+ clients_021/email/ses.py,sha256=0ukHvuVE0nTU5mYWRy2EaHY5v9dUwi4GSKES93pjqBg,4668
19
+ clients_021/queue/__init__.py,sha256=CesweOrXmX5cyU-mRvqu-4wLkNJI15zi_Lows6SPupY,338
20
+ clients_021/queue/qstash.py,sha256=j3vz_S9-ed0YQZCaWwQg5wPwf8Bv5P1voNsx8WD8Zhc,8315
21
+ clients_021/queue/sqs.py,sha256=jnccgbxyglzs2mR_DsH04IisEwygDlw6s3RxR0UtMTc,6693
22
+ clients_021/storage/__init__.py,sha256=KrJViRHBK7PkIYnagg7vyXp95SzMFp9cUCKEnMeFr84,78
23
+ clients_021/storage/s3.py,sha256=NDtG9QzHSXvrcazTYTq1EH5VjuRd2xMMsMvaumqci4Q,5946
24
+ clients_021/storage/vercel_blob.py,sha256=FlnoKiahqtk3PpLsGpyu_Gy87otg15i4BlJZsIB3LfM,6511
25
+ 021_python_clients-0.2.0.dist-info/METADATA,sha256=gx7Jt2wjPNVFhyUHCze-tQ5YVu-Kqrt2GwAyia6HdXs,20595
26
+ 021_python_clients-0.2.0.dist-info/WHEEL,sha256=W3fkpkm7-wf9vBI5Z-7s0eWkeM-spu78I8Neb98DeEg,87
27
+ 021_python_clients-0.2.0.dist-info/licenses/LICENSE,sha256=z8d0m5b2O9McPEK1xHG_dWgUBT6EfBDz6wA0F7xSPTA,11358
28
+ 021_python_clients-0.2.0.dist-info/RECORD,,
clients_021/__init__.py CHANGED
@@ -1,9 +1,10 @@
1
1
  """Observed clients for the external services Python services call.
2
2
 
3
3
  `clients_021.http`, `clients_021.queue.qstash` and `clients_021.email` (Resend) need no extras;
4
- `clients_021.redis` needs the `redis` extra; `clients_021.queue.sqs`, `clients_021.s3`,
5
- `clients_021.email.ses` and `clients_021.kms` the `aws` extra; `clients_021.httpx2` the
6
- `httpx2` extra.
4
+ `clients_021.cache.redis` needs the `redis` extra; `clients_021.queue.sqs`,
5
+ `clients_021.storage.s3`, `clients_021.email.ses` and `clients_021.crypto.kms` the `aws` extra;
6
+ `clients_021.httpx2` the `httpx2` extra.
7
+ `clients_021.storage.vercel_blob` needs the `vercel-blob` extra.
7
8
  """
8
9
 
9
10
  from importlib.metadata import version
clients_021/_observe.py CHANGED
@@ -2,7 +2,7 @@
2
2
 
3
3
  The `db` family reuses the instruments 021-python-observability already creates for SQL
4
4
  (`db.client.operation.*`), so Redis and Postgres share dashboards. `messaging` (SQS, email),
5
- `storage` (S3) and `crypto` (KMS) get their own pair on the same meter. Attributes are
5
+ `storage` (S3, Vercel Blob) and `crypto` (KMS) get their own pair on the same meter. Attributes are
6
6
  bounded: dependency, operation, status and error.type; never keys, URLs, bodies or ids.
7
7
  """
8
8
 
@@ -0,0 +1 @@
1
+ """Cache clients. Import providers explicitly; SDKs are optional extras."""
@@ -0,0 +1,124 @@
1
+ """Redis, observed: JSON values with a TTL, one-time reads, windowed counters and a ping.
2
+
3
+ from clients_021.cache.redis import RedisClient
4
+
5
+ redis = RedisClient.from_url(obs, settings.redis_url, dependency="redis")
6
+ await redis.set_json("session:abc", {"user": "ana"}, ttl_seconds=1800)
7
+ value = await redis.get_json("session:abc")
8
+
9
+ Needs the `redis` extra. Calls record `db.client.operation.*` with `dependency` and
10
+ `operation` (`get_json`, `set_json`, `take_json`, `delete`, `incr_window`, `ping`).
11
+ Timeouts and connection failures raise DependencyUnavailable; a wrong command or type
12
+ (`ResponseError`) is a bug and keeps its type.
13
+ """
14
+
15
+ from __future__ import annotations
16
+
17
+ import json
18
+ from collections.abc import Awaitable
19
+ from typing import Any, cast
20
+
21
+ from observability_021 import Observability
22
+ from observability_021.config import identifier
23
+ from redis.asyncio import Redis
24
+ from redis.exceptions import AuthenticationError, AuthorizationError
25
+ from redis.exceptions import ConnectionError as RedisConnectionError
26
+ from redis.exceptions import TimeoutError as RedisTimeoutError
27
+
28
+ from clients_021._observe import observed
29
+ from clients_021.errors import DependencyUnavailable
30
+
31
+ SYSTEM = {"db.system.name": "redis"}
32
+
33
+
34
+ def _classify(error: BaseException) -> str | None:
35
+ # Both inherit ConnectionError, but retrying cannot fix invalid credentials.
36
+ if isinstance(error, AuthenticationError | AuthorizationError):
37
+ return None
38
+ if isinstance(error, RedisTimeoutError):
39
+ return "timeout"
40
+ if isinstance(error, RedisConnectionError):
41
+ return "connection_error"
42
+ return None
43
+
44
+
45
+ class RedisClient:
46
+ def __init__(self, obs: Observability, redis: Redis, *, dependency: str = "redis") -> None:
47
+ self._obs = obs
48
+ self._redis = redis
49
+ self.dependency = identifier(dependency)
50
+
51
+ @classmethod
52
+ def from_url(
53
+ cls,
54
+ obs: Observability,
55
+ url: str,
56
+ *,
57
+ dependency: str = "redis",
58
+ timeout_seconds: float = 1.0,
59
+ ) -> RedisClient:
60
+ # redis-py's from_url is only partially typed.
61
+ redis: Redis = Redis.from_url( # pyright: ignore[reportUnknownMemberType]
62
+ url, socket_timeout=timeout_seconds, socket_connect_timeout=timeout_seconds
63
+ )
64
+ return cls(obs, redis, dependency=dependency)
65
+
66
+ def _observed(self, operation: str): # noqa: ANN202 - an async context manager
67
+ return observed(
68
+ self._obs,
69
+ family="db",
70
+ system=SYSTEM,
71
+ dependency=self.dependency,
72
+ operation=operation,
73
+ classify=_classify,
74
+ )
75
+
76
+ async def get_json(self, key: str) -> Any:
77
+ """The decoded value, or None when the key is missing or expired."""
78
+ async with self._observed("get_json"):
79
+ raw = await self._redis.get(key)
80
+ return None if raw is None else json.loads(raw)
81
+
82
+ async def set_json(self, key: str, value: Any, *, ttl_seconds: int) -> None:
83
+ """Every value expires: there is no way to store one forever by mistake."""
84
+ if ttl_seconds < 1:
85
+ raise ValueError("ttl_seconds must be at least 1")
86
+ data = json.dumps(value, separators=(",", ":"))
87
+ async with self._observed("set_json"):
88
+ await self._redis.set(key, data, ex=ttl_seconds)
89
+
90
+ async def take_json(self, key: str) -> Any:
91
+ """Reads and deletes in one step (GETDEL): a second caller gets None."""
92
+ async with self._observed("take_json"):
93
+ raw = await cast(Awaitable[Any], self._redis.getdel(key))
94
+ return None if raw is None else json.loads(raw)
95
+
96
+ async def delete(self, key: str) -> None:
97
+ async with self._observed("delete"):
98
+ await self._redis.delete(key)
99
+
100
+ async def incr_window(self, key: str, *, by: int = 1, ttl_seconds: int) -> int:
101
+ """Adds `by` and returns the total; the key expires `ttl_seconds` after its first hit.
102
+
103
+ INCRBY and EXPIRE NX run in one transaction, so a crash between them cannot leave
104
+ a counter that never expires.
105
+ """
106
+ if ttl_seconds < 1:
107
+ raise ValueError("ttl_seconds must be at least 1")
108
+ async with self._observed("incr_window"), self._redis.pipeline(transaction=True) as pipe:
109
+ pipe.incrby(key, by)
110
+ pipe.expire(key, ttl_seconds, nx=True)
111
+ total, _ = await pipe.execute()
112
+ return int(total)
113
+
114
+ async def ping(self) -> bool:
115
+ """True when Redis answers; False (recorded as an error) when it does not."""
116
+ try:
117
+ async with self._observed("ping"):
118
+ # redis-py types ping() for both clients; the asyncio one returns an awaitable.
119
+ return bool(await self._redis.ping()) # pyright: ignore[reportUnknownMemberType]
120
+ except DependencyUnavailable:
121
+ return False
122
+
123
+ async def aclose(self) -> None:
124
+ await self._redis.aclose()
@@ -0,0 +1 @@
1
+ """Crypto clients. Import providers explicitly; SDKs are optional extras."""
@@ -0,0 +1,178 @@
1
+ """AWS KMS, observed: encrypt, decrypt and data keys, always bound to an encryption context.
2
+
3
+ from clients_021.crypto.kms import KmsClient
4
+
5
+ keys = KmsClient.create(obs, settings.session_key_id, region="us-east-1")
6
+ context = {"purpose": "session", "tenant": tenant_id}
7
+ blob = await keys.encrypt(secret, context=context)
8
+ secret = await keys.decrypt(blob, context=context)
9
+ envelope = await keys.seal(large_payload, context=context) # any size, AES-256-GCM
10
+ large_payload = await keys.unseal(envelope, context=context)
11
+
12
+ Needs the `aws` extra. The encryption context is required: decrypting with another context
13
+ fails, so a ciphertext only opens for the purpose it was made for. KMS logs it in CloudTrail
14
+ in plaintext: it identifies the purpose, it is never a secret. Calls record
15
+ `crypto.client.operation.*` with `operation` (`encrypt`, `decrypt`, `generate_data_key`) and
16
+ a span with `rpc.service=KMS`; never key ids, contexts, plaintext or ciphertext.
17
+ Timeouts, connection failures, throttling and 5xx raise DependencyUnavailable. Access
18
+ denied, a disabled key, a wrong context or a tampered ciphertext (`InvalidCiphertextException`)
19
+ and a ciphertext from another key (`IncorrectKeyException`) keep their botocore type.
20
+ """
21
+
22
+ from __future__ import annotations
23
+
24
+ import asyncio
25
+ import json
26
+ import os
27
+ import struct
28
+ from collections.abc import Mapping
29
+ from dataclasses import dataclass
30
+ from typing import TYPE_CHECKING, Any
31
+
32
+ import boto3
33
+ from cryptography.hazmat.primitives.ciphers.aead import AESGCM
34
+ from observability_021 import Observability
35
+ from observability_021.config import identifier
36
+
37
+ from clients_021 import _aws
38
+ from clients_021._observe import observed
39
+
40
+ if TYPE_CHECKING:
41
+ from mypy_boto3_kms import KMSClient
42
+ else:
43
+ KMSClient = Any
44
+
45
+ SYSTEM = {"rpc.system": "aws-api", "rpc.service": "KMS"}
46
+ # Envelope layout: version (1 byte), wrapped key length (2 bytes, big endian), wrapped key,
47
+ # nonce (12 bytes), AES-256-GCM ciphertext and tag. The header and context are the AAD.
48
+ _ENVELOPE = struct.Struct(">BH")
49
+ _VERSION = 1
50
+ _NONCE = 12
51
+ _TAG = 16
52
+
53
+
54
+ @dataclass(frozen=True, repr=False)
55
+ class DataKey:
56
+ """A data key: use `plaintext` locally and drop it; store only `ciphertext`."""
57
+
58
+ plaintext: bytes
59
+ ciphertext: bytes
60
+
61
+ def __repr__(self) -> str:
62
+ return "DataKey(<redacted>)"
63
+
64
+
65
+ def _context(context: Mapping[str, str]) -> dict[str, str]:
66
+ if not context or not all(context):
67
+ raise ValueError("context must be a non-empty mapping with non-empty keys")
68
+ return dict(context)
69
+
70
+
71
+ def _aad(header: bytes, context: dict[str, str]) -> bytes:
72
+ canonical = json.dumps(context, sort_keys=True, separators=(",", ":"), ensure_ascii=False)
73
+ return header + canonical.encode()
74
+
75
+
76
+ class KmsClient:
77
+ def __init__(
78
+ self, obs: Observability, client: KMSClient, key_id: str, *, dependency: str = "kms"
79
+ ) -> None:
80
+ if not key_id:
81
+ raise ValueError("key_id is required")
82
+ self._obs = obs
83
+ self._client = client
84
+ self._key_id = key_id
85
+ self.dependency = identifier(dependency)
86
+
87
+ @classmethod
88
+ def create(
89
+ cls,
90
+ obs: Observability,
91
+ key_id: str,
92
+ *,
93
+ region: str | None = None,
94
+ endpoint_url: str | None = None,
95
+ dependency: str = "kms",
96
+ timeout_seconds: float = 5.0,
97
+ ) -> KmsClient:
98
+ client: KMSClient = boto3.client( # pyright: ignore[reportUnknownMemberType]
99
+ "kms",
100
+ region_name=region,
101
+ endpoint_url=endpoint_url,
102
+ config=_aws.config(
103
+ connect_seconds=min(timeout_seconds, 3.0), read_seconds=timeout_seconds
104
+ ),
105
+ )
106
+ return cls(obs, client, key_id, dependency=dependency)
107
+
108
+ def _observed(self, operation: str): # noqa: ANN202 - an async context manager
109
+ return observed(
110
+ self._obs,
111
+ family="crypto",
112
+ system=SYSTEM,
113
+ dependency=self.dependency,
114
+ operation=operation,
115
+ classify=_aws.classify,
116
+ )
117
+
118
+ async def encrypt(self, plaintext: bytes, *, context: Mapping[str, str]) -> bytes:
119
+ """Up to 4096 bytes (a KMS limit); use `seal` for anything larger."""
120
+ encryption_context = _context(context)
121
+ async with self._observed("encrypt"):
122
+ response = await asyncio.to_thread(
123
+ self._client.encrypt,
124
+ KeyId=self._key_id,
125
+ Plaintext=plaintext,
126
+ EncryptionContext=encryption_context,
127
+ )
128
+ return response["CiphertextBlob"]
129
+
130
+ async def decrypt(self, ciphertext: bytes, *, context: Mapping[str, str]) -> bytes:
131
+ """Names the key, as KMS recommends: a ciphertext from another key is rejected."""
132
+ encryption_context = _context(context)
133
+ async with self._observed("decrypt"):
134
+ response = await asyncio.to_thread(
135
+ self._client.decrypt,
136
+ KeyId=self._key_id,
137
+ CiphertextBlob=ciphertext,
138
+ EncryptionContext=encryption_context,
139
+ )
140
+ return response["Plaintext"]
141
+
142
+ async def generate_data_key(self, *, context: Mapping[str, str]) -> DataKey:
143
+ """A fresh AES-256 key, in plaintext and wrapped under this client's KMS key."""
144
+ encryption_context = _context(context)
145
+ async with self._observed("generate_data_key"):
146
+ response = await asyncio.to_thread(
147
+ self._client.generate_data_key,
148
+ KeyId=self._key_id,
149
+ KeySpec="AES_256",
150
+ EncryptionContext=encryption_context,
151
+ )
152
+ return DataKey(response["Plaintext"], response["CiphertextBlob"])
153
+
154
+ async def seal(self, plaintext: bytes, *, context: Mapping[str, str]) -> bytes:
155
+ """Envelope encryption: a new data key per call (one KMS request), AES-256-GCM locally."""
156
+ encryption_context = _context(context)
157
+ key = await self.generate_data_key(context=encryption_context)
158
+ header = _ENVELOPE.pack(_VERSION, len(key.ciphertext)) + key.ciphertext
159
+ nonce = os.urandom(_NONCE)
160
+ sealed = AESGCM(key.plaintext).encrypt(nonce, plaintext, _aad(header, encryption_context))
161
+ return header + nonce + sealed
162
+
163
+ async def unseal(self, envelope: bytes, *, context: Mapping[str, str]) -> bytes:
164
+ """Raises ValueError for a malformed envelope, cryptography's InvalidTag if tampered."""
165
+ encryption_context = _context(context)
166
+ if len(envelope) < _ENVELOPE.size:
167
+ raise ValueError("not a sealed envelope")
168
+ version, size = _ENVELOPE.unpack_from(envelope)
169
+ end = _ENVELOPE.size + size
170
+ if version != _VERSION or size == 0 or len(envelope) < end + _NONCE + _TAG:
171
+ raise ValueError("not a sealed envelope")
172
+ key = await self.decrypt(envelope[_ENVELOPE.size : end], context=encryption_context)
173
+ nonce, sealed = envelope[end : end + _NONCE], envelope[end + _NONCE :]
174
+ return AESGCM(key).decrypt(nonce, sealed, _aad(envelope[:end], encryption_context))
175
+
176
+ async def aclose(self) -> None:
177
+ """Release the underlying SDK connection pools after in-flight calls finish."""
178
+ await asyncio.to_thread(self._client.close)