srx-lib-azure 0.1.6__tar.gz → 0.1.8__tar.gz

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,6 +1,6 @@
1
1
  Metadata-Version: 2.4
2
2
  Name: srx-lib-azure
3
- Version: 0.1.6
3
+ Version: 0.1.8
4
4
  Summary: Azure helpers for SRX services: Blob, Email, Table
5
5
  Author-email: SRX <dev@srx.id>
6
6
  Requires-Python: >=3.12
@@ -4,7 +4,7 @@ build-backend = "hatchling.build"
4
4
 
5
5
  [project]
6
6
  name = "srx-lib-azure"
7
- version = "0.1.6"
7
+ version = "0.1.8"
8
8
  description = "Azure helpers for SRX services: Blob, Email, Table"
9
9
  readme = "README.md"
10
10
  requires-python = ">=3.12"
@@ -3,4 +3,3 @@ from .email import EmailService
3
3
  from .table import AzureTableService
4
4
 
5
5
  __all__ = ["AzureBlobService", "EmailService", "AzureTableService"]
6
-
@@ -4,6 +4,11 @@ from datetime import datetime, timedelta, timezone
4
4
  from typing import Optional, BinaryIO, Tuple
5
5
 
6
6
  from azure.storage.blob import BlobServiceClient, BlobSasPermissions, generate_blob_sas
7
+ from azure.core.exceptions import (
8
+ ResourceNotFoundError,
9
+ ClientAuthenticationError,
10
+ HttpResponseError,
11
+ )
7
12
  from fastapi import UploadFile
8
13
 
9
14
  from loguru import logger
@@ -49,9 +54,7 @@ class AzureBlobService:
49
54
  return None, None
50
55
  try:
51
56
  clean = self.connection_string.strip().strip('"').strip("'")
52
- parts = dict(
53
- seg.split("=", 1) for seg in clean.split(";") if "=" in seg
54
- )
57
+ parts = dict(seg.split("=", 1) for seg in clean.split(";") if "=" in seg)
55
58
  account_name = parts.get("AccountName")
56
59
  account_key = parts.get("AccountKey") or self.account_key
57
60
  return account_name, account_key
@@ -93,9 +96,20 @@ class AzureBlobService:
93
96
  if self.base_blob_url:
94
97
  base_url = self.base_blob_url.strip().strip('"').strip("'").rstrip("/")
95
98
  return f"{base_url}/{blob_name}?{sas}"
96
- return f"https://{account_name}.blob.core.windows.net/{self.container_name}/{blob_name}?{sas}"
99
+ return (
100
+ f"https://{account_name}.blob.core.windows.net/{self.container_name}/{blob_name}?{sas}"
101
+ )
97
102
 
98
103
  async def upload_file(self, file: UploadFile, blob_path: str) -> Optional[str]:
104
+ """Upload a file to Azure Blob Storage and return a SAS URL.
105
+
106
+ Args:
107
+ file: File to upload
108
+ blob_path: Destination path in the container
109
+
110
+ Returns:
111
+ SAS URL if successful, None on error
112
+ """
99
113
  if not self.connection_string:
100
114
  logger.error("Azure Storage connection string not configured")
101
115
  return None
@@ -105,13 +119,37 @@ class AzureBlobService:
105
119
  container = client.get_container_client(self.container_name)
106
120
  content = await file.read()
107
121
  blob_client = container.get_blob_client(blob_path)
108
- blob_client.upload_blob(content, overwrite=True, content_type=file.content_type or "application/octet-stream")
122
+ blob_client.upload_blob(
123
+ content,
124
+ overwrite=True,
125
+ content_type=file.content_type or "application/octet-stream",
126
+ )
109
127
  return self._generate_sas_url(blob_path)
128
+ except ClientAuthenticationError as e:
129
+ logger.error(f"Authentication failed uploading {file.filename}: {e}")
130
+ return None
131
+ except HttpResponseError as e:
132
+ logger.error(
133
+ f"Azure service error uploading {file.filename}: {e.status_code} - {e.message}"
134
+ )
135
+ return None
110
136
  except Exception as e:
111
- logger.error(f"Failed to upload file {file.filename}: {e}")
137
+ logger.error(f"Unexpected error uploading {file.filename}: {e}")
112
138
  return None
113
139
 
114
- async def upload_stream(self, stream: BinaryIO, blob_path: str, content_type: str = "application/octet-stream") -> Optional[str]:
140
+ async def upload_stream(
141
+ self, stream: BinaryIO, blob_path: str, content_type: str = "application/octet-stream"
142
+ ) -> Optional[str]:
143
+ """Upload a binary stream to Azure Blob Storage and return a SAS URL.
144
+
145
+ Args:
146
+ stream: Binary stream to upload
147
+ blob_path: Destination path in the container
148
+ content_type: MIME type of the content
149
+
150
+ Returns:
151
+ SAS URL if successful, None on error
152
+ """
115
153
  if not self.connection_string:
116
154
  logger.error("Azure Storage connection string not configured")
117
155
  return None
@@ -122,12 +160,27 @@ class AzureBlobService:
122
160
  blob_client = container.get_blob_client(blob_path)
123
161
  blob_client.upload_blob(stream, overwrite=True, content_type=content_type)
124
162
  return self._generate_sas_url(blob_path)
163
+ except ClientAuthenticationError as e:
164
+ logger.error(f"Authentication failed uploading stream to {blob_path}: {e}")
165
+ return None
166
+ except HttpResponseError as e:
167
+ logger.error(
168
+ f"Azure service error uploading stream to {blob_path}: {e.status_code} - {e.message}"
169
+ )
170
+ return None
125
171
  except Exception as e:
126
- logger.error(f"Failed to upload stream to {blob_path}: {e}")
172
+ logger.error(f"Unexpected error uploading stream to {blob_path}: {e}")
127
173
  return None
128
174
 
129
175
  async def download_file(self, blob_path: str) -> Optional[bytes]:
130
- """Download a blob's content as bytes."""
176
+ """Download a blob's content as bytes.
177
+
178
+ Returns:
179
+ bytes if successful, None if blob doesn't exist
180
+
181
+ Raises:
182
+ RuntimeError: For connection/auth errors (caller should handle)
183
+ """
131
184
  if not self.connection_string:
132
185
  logger.error("Azure Storage connection string not configured")
133
186
  return None
@@ -139,9 +192,24 @@ class AzureBlobService:
139
192
  content = download_stream.readall()
140
193
  logger.info(f"Successfully downloaded {blob_path}")
141
194
  return content
142
- except Exception as e:
143
- logger.error(f"Failed to download {blob_path}: {e}")
195
+ except ResourceNotFoundError:
196
+ # Blob doesn't exist - this is expected in many scenarios
197
+ logger.info(f"Blob not found: {blob_path}")
144
198
  return None
199
+ except ClientAuthenticationError as e:
200
+ # Auth errors should not be retried - they need credential fixes
201
+ logger.error(f"Authentication failed for {blob_path}: {e}")
202
+ raise RuntimeError(f"Azure authentication failed: {e}") from e
203
+ except HttpResponseError as e:
204
+ # Other Azure service errors (rate limits, service issues, etc.)
205
+ logger.error(
206
+ f"Azure service error downloading {blob_path}: {e.status_code} - {e.message}"
207
+ )
208
+ raise RuntimeError(f"Azure Blob download failed for {blob_path}: {e.message}") from e
209
+ except Exception as e:
210
+ # Catch-all for unexpected errors (network, etc.)
211
+ logger.error(f"Unexpected error downloading {blob_path}: {e}")
212
+ raise RuntimeError(f"Unexpected error downloading {blob_path}: {e}") from e
145
213
 
146
214
  async def download_to_temp_file(self, blob_path: str) -> Optional[str]:
147
215
  """Download a blob to a temporary file and return its path."""
@@ -149,7 +217,9 @@ class AzureBlobService:
149
217
  if content is None:
150
218
  return None
151
219
  try:
152
- with tempfile.NamedTemporaryFile(delete=False, suffix=os.path.splitext(blob_path)[1]) as tf:
220
+ with tempfile.NamedTemporaryFile(
221
+ delete=False, suffix=os.path.splitext(blob_path)[1]
222
+ ) as tf:
153
223
  tf.write(content)
154
224
  path = tf.name
155
225
  logger.info(f"Downloaded {blob_path} to temporary file: {path}")
@@ -172,7 +242,14 @@ class AzureBlobService:
172
242
  return None
173
243
 
174
244
  async def delete_file(self, blob_path: str) -> bool:
175
- """Delete a blob and return True on success."""
245
+ """Delete a blob and return True on success.
246
+
247
+ Args:
248
+ blob_path: Path to the blob to delete
249
+
250
+ Returns:
251
+ True if deleted successfully or blob doesn't exist, False on error
252
+ """
176
253
  if not self.connection_string:
177
254
  logger.error("Azure Storage connection string not configured")
178
255
  return False
@@ -183,12 +260,29 @@ class AzureBlobService:
183
260
  blob_client.delete_blob()
184
261
  logger.info(f"Successfully deleted {blob_path}")
185
262
  return True
263
+ except ResourceNotFoundError:
264
+ # Blob already doesn't exist - this is still success
265
+ logger.info(f"Blob {blob_path} already deleted or doesn't exist")
266
+ return True
267
+ except ClientAuthenticationError as e:
268
+ logger.error(f"Authentication failed when deleting {blob_path}: {e}")
269
+ return False
270
+ except HttpResponseError as e:
271
+ logger.error(f"Azure service error deleting {blob_path}: {e.status_code} - {e.message}")
272
+ return False
186
273
  except Exception as e:
187
- logger.error(f"Failed to delete {blob_path}: {e}")
274
+ logger.error(f"Unexpected error deleting {blob_path}: {e}")
188
275
  return False
189
276
 
190
277
  async def file_exists(self, blob_path: str) -> bool:
191
- """Check if a blob exists in the container."""
278
+ """Check if a blob exists in the container.
279
+
280
+ Args:
281
+ blob_path: Path to the blob to check
282
+
283
+ Returns:
284
+ True if blob exists, False otherwise (including on errors)
285
+ """
192
286
  if not self.connection_string:
193
287
  logger.error("Azure Storage connection string not configured")
194
288
  return False
@@ -197,6 +291,12 @@ class AzureBlobService:
197
291
  container = client.get_container_client(self.container_name)
198
292
  blob_client = container.get_blob_client(blob_path)
199
293
  return blob_client.exists()
294
+ except ClientAuthenticationError as e:
295
+ logger.error(f"Authentication failed checking {blob_path}: {e}")
296
+ return False
297
+ except HttpResponseError as e:
298
+ logger.error(f"Azure service error checking {blob_path}: {e.status_code} - {e.message}")
299
+ return False
200
300
  except Exception as e:
201
- logger.error(f"Failed to check existence of {blob_path}: {e}")
301
+ logger.error(f"Unexpected error checking existence of {blob_path}: {e}")
202
302
  return False
@@ -4,8 +4,16 @@ from typing import Dict, Any
4
4
 
5
5
  try:
6
6
  from azure.communication.email.aio import EmailClient
7
+ from azure.core.exceptions import (
8
+ ClientAuthenticationError,
9
+ HttpResponseError,
10
+ ServiceRequestError,
11
+ )
7
12
  except Exception: # pragma: no cover - optional dependency at import time
8
13
  EmailClient = None # type: ignore
14
+ ClientAuthenticationError = None # type: ignore
15
+ HttpResponseError = None # type: ignore
16
+ ServiceRequestError = None # type: ignore
9
17
 
10
18
  logger = logging.getLogger(__name__)
11
19
 
@@ -39,10 +47,27 @@ class EmailService:
39
47
  self.email_client = None
40
48
  logger.warning("EmailService initialization failed: %s", e)
41
49
 
42
- async def send_notification(self, recipient: str, subject: str, body: str, html: bool = False) -> Dict[str, Any]:
50
+ async def send_notification(
51
+ self, recipient: str, subject: str, body: str, html: bool = False
52
+ ) -> Dict[str, Any]:
53
+ """Send an email notification via Azure Communication Services.
54
+
55
+ Args:
56
+ recipient: Email address of the recipient
57
+ subject: Email subject line
58
+ body: Email body content
59
+ html: If True, send as HTML; otherwise plain text
60
+
61
+ Returns:
62
+ Dict with status, message, and optional message_id
63
+ - status: "success" | "error" | "skipped"
64
+ - message: Human-readable message
65
+ - message_id: Azure message ID (only on success)
66
+ """
43
67
  if not self.email_client or not self.sender_address:
44
68
  logger.warning("Email skipped: service not configured")
45
69
  return {"status": "skipped", "message": "Email service not configured"}
70
+
46
71
  message = {
47
72
  "content": {"subject": subject},
48
73
  "recipients": {"to": [{"address": recipient}]},
@@ -52,15 +77,41 @@ class EmailService:
52
77
  message["content"]["html"] = body
53
78
  else:
54
79
  message["content"]["plainText"] = body
80
+
55
81
  try:
56
82
  poller = await self.email_client.begin_send(message)
57
83
  result = await poller.result()
58
84
  message_id = result.get("id")
59
85
  if message_id:
60
86
  logger.info("Email sent to %s with Message ID: %s", recipient, message_id)
61
- return {"status": "success", "message": "Email sent successfully", "message_id": message_id}
87
+ return {
88
+ "status": "success",
89
+ "message": "Email sent successfully",
90
+ "message_id": message_id,
91
+ }
62
92
  logger.error("Failed to send email. Result: %s", result)
63
93
  return {"status": "error", "message": f"Failed to send email: {result}"}
94
+ except ClientAuthenticationError as e:
95
+ # Auth errors should not be retried - they need credential fixes
96
+ logger.error("Authentication failed sending email to %s: %s", recipient, e)
97
+ return {"status": "error", "message": f"Authentication failed: {e}"}
98
+ except HttpResponseError as e:
99
+ # Azure service errors (rate limits, invalid recipient, etc.)
100
+ logger.error(
101
+ "Azure service error sending email to %s: %s - %s",
102
+ recipient,
103
+ e.status_code,
104
+ e.message,
105
+ )
106
+ return {
107
+ "status": "error",
108
+ "message": f"Azure service error ({e.status_code}): {e.message}",
109
+ }
110
+ except ServiceRequestError as e:
111
+ # Network/connection errors - may be retryable
112
+ logger.error("Network error sending email to %s: %s", recipient, e)
113
+ return {"status": "error", "message": f"Network error: {e}"}
64
114
  except Exception as e:
65
- logger.error("Email send exception: %s", e)
66
- return {"status": "error", "message": str(e)}
115
+ # Catch-all for unexpected errors
116
+ logger.error("Unexpected error sending email to %s: %s", recipient, e)
117
+ return {"status": "error", "message": f"Unexpected error: {e}"}
@@ -0,0 +1,498 @@
1
+ from __future__ import annotations
2
+
3
+ import os
4
+ from dataclasses import dataclass
5
+ from datetime import datetime, timezone
6
+ from typing import Any, Dict, Iterable, List, Optional
7
+
8
+ from loguru import logger
9
+
10
+ try:
11
+ from azure.data.tables import TableServiceClient
12
+ from azure.core.exceptions import (
13
+ ResourceNotFoundError,
14
+ ResourceExistsError,
15
+ ClientAuthenticationError,
16
+ HttpResponseError,
17
+ )
18
+ except Exception: # pragma: no cover
19
+ TableServiceClient = None # type: ignore
20
+ ResourceNotFoundError = None # type: ignore
21
+ ResourceExistsError = None # type: ignore
22
+ ClientAuthenticationError = None # type: ignore
23
+ HttpResponseError = None # type: ignore
24
+
25
+
26
+ def _now_iso() -> str:
27
+ return datetime.now(timezone.utc).isoformat()
28
+
29
+
30
+ @dataclass
31
+ class AzureTableService:
32
+ connection_string: Optional[str] = None
33
+
34
+ def __init__(self, connection_string: Optional[str] = None) -> None:
35
+ # Constructor injection preferred; fallback to env only if not provided
36
+ self.connection_string = connection_string or os.getenv("AZURE_STORAGE_CONNECTION_STRING")
37
+
38
+ def _get_client(self) -> "TableServiceClient":
39
+ if not self.connection_string:
40
+ raise RuntimeError("AZURE_STORAGE_CONNECTION_STRING not configured")
41
+ if TableServiceClient is None:
42
+ raise RuntimeError("azure-data-tables not installed; install to use table operations")
43
+ clean = self.connection_string.strip().strip('"').strip("'")
44
+ return TableServiceClient.from_connection_string(conn_str=clean)
45
+
46
+ def ensure_table(self, table_name: str) -> None:
47
+ client = self._get_client()
48
+ try:
49
+ client.create_table_if_not_exists(table_name=table_name)
50
+ except Exception as e:
51
+ logger.warning("ensure_table(%s) warning: %s", table_name, e)
52
+
53
+ def list_tables(self) -> List[str]:
54
+ """List all tables in the storage account.
55
+
56
+ Returns:
57
+ List of table names
58
+ """
59
+ client = self._get_client()
60
+ try:
61
+ tables = [table.name for table in client.list_tables()]
62
+ logger.info("Listed %d tables", len(tables))
63
+ return tables
64
+ except Exception as exc:
65
+ logger.error("Failed to list tables: %s", exc)
66
+ return []
67
+
68
+ def delete_table(self, table_name: str) -> bool:
69
+ """Delete a table.
70
+
71
+ Args:
72
+ table_name: Name of the table to delete
73
+
74
+ Returns:
75
+ True if deleted successfully or table doesn't exist, False on other errors
76
+ """
77
+ client = self._get_client()
78
+ try:
79
+ client.delete_table(table_name=table_name)
80
+ logger.info("Deleted table: %s", table_name)
81
+ return True
82
+ except ResourceNotFoundError:
83
+ # Table already deleted or doesn't exist - this is still success
84
+ logger.info(f"Table {table_name} already deleted or doesn't exist")
85
+ return True
86
+ except ClientAuthenticationError as exc:
87
+ logger.error("Authentication failed deleting table %s: %s", table_name, exc)
88
+ return False
89
+ except HttpResponseError as exc:
90
+ logger.error("Azure service error deleting table %s: %s", table_name, exc.message)
91
+ return False
92
+ except Exception as exc:
93
+ logger.error("Unexpected error deleting table %s: %s", table_name, exc)
94
+ return False
95
+
96
+ def table_exists(self, table_name: str) -> bool:
97
+ """Check if a table exists.
98
+
99
+ Args:
100
+ table_name: Name of the table
101
+
102
+ Returns:
103
+ True if table exists, False otherwise
104
+ """
105
+ return table_name in self.list_tables()
106
+
107
+ def put_entity(self, table_name: str, entity: Dict[str, Any]) -> Dict[str, Any]:
108
+ """Insert a new entity into a table.
109
+
110
+ Args:
111
+ table_name: Name of the table
112
+ entity: Entity dictionary with PartitionKey and RowKey
113
+
114
+ Returns:
115
+ Dict with etag and timestamp
116
+
117
+ Raises:
118
+ RuntimeError: If entity already exists or other errors occur
119
+ """
120
+ client = self._get_client()
121
+ table = client.get_table_client(table_name)
122
+ try:
123
+ res = table.create_entity(entity=entity)
124
+ logger.info(
125
+ "Inserted entity into %s: PK=%s RK=%s",
126
+ table_name,
127
+ entity.get("PartitionKey"),
128
+ entity.get("RowKey"),
129
+ )
130
+ return {"etag": getattr(res, "etag", None), "ts": _now_iso()}
131
+ except ResourceExistsError as e:
132
+ pk, rk = entity.get("PartitionKey"), entity.get("RowKey")
133
+ logger.error(f"Entity already exists in {table_name}: PK={pk} RK={rk}")
134
+ raise RuntimeError(f"Entity already exists: {e}") from e
135
+ except ClientAuthenticationError as e:
136
+ logger.error(f"Authentication failed inserting entity into {table_name}: {e}")
137
+ raise RuntimeError(f"Authentication failed: {e}") from e
138
+ except HttpResponseError as e:
139
+ logger.error(f"Azure service error inserting entity into {table_name}: {e.message}")
140
+ raise RuntimeError(f"Failed to insert entity: {e.message}") from e
141
+
142
+ def upsert_entity(self, table_name: str, entity: Dict[str, Any]) -> Dict[str, Any]:
143
+ """Insert or update an entity in a table.
144
+
145
+ Args:
146
+ table_name: Name of the table
147
+ entity: Entity dictionary with PartitionKey and RowKey
148
+
149
+ Returns:
150
+ Dict with etag and timestamp
151
+
152
+ Raises:
153
+ RuntimeError: If authentication or service errors occur
154
+ """
155
+ client = self._get_client()
156
+ table = client.get_table_client(table_name)
157
+ try:
158
+ res = table.upsert_entity(entity=entity, mode="merge")
159
+ logger.info(
160
+ "Upserted entity into %s: PK=%s RK=%s",
161
+ table_name,
162
+ entity.get("PartitionKey"),
163
+ entity.get("RowKey"),
164
+ )
165
+ return {"etag": getattr(res, "etag", None), "ts": _now_iso()}
166
+ except ClientAuthenticationError as e:
167
+ logger.error(f"Authentication failed upserting entity into {table_name}: {e}")
168
+ raise RuntimeError(f"Authentication failed: {e}") from e
169
+ except HttpResponseError as e:
170
+ logger.error(f"Azure service error upserting entity into {table_name}: {e.message}")
171
+ raise RuntimeError(f"Failed to upsert entity: {e.message}") from e
172
+
173
+ def delete_entity(self, table_name: str, partition_key: str, row_key: str) -> bool:
174
+ """Delete an entity from a table.
175
+
176
+ Args:
177
+ table_name: Name of the table
178
+ partition_key: Partition key of the entity
179
+ row_key: Row key of the entity
180
+
181
+ Returns:
182
+ True if deleted successfully or entity doesn't exist, False on other errors
183
+
184
+ Note:
185
+ Azure Tables delete_entity doesn't fail if entity doesn't exist
186
+ """
187
+ client = self._get_client()
188
+ table = client.get_table_client(table_name)
189
+ try:
190
+ table.delete_entity(partition_key=partition_key, row_key=row_key)
191
+ logger.info("Deleted entity in %s (%s/%s)", table_name, partition_key, row_key)
192
+ return True
193
+ except ResourceNotFoundError:
194
+ # Entity already deleted or doesn't exist - this is still success
195
+ logger.info(
196
+ f"Entity in {table_name} ({partition_key}/{row_key}) already deleted or doesn't exist"
197
+ )
198
+ return True
199
+ except ClientAuthenticationError as exc:
200
+ logger.error("Authentication failed deleting entity in %s: %s", table_name, exc)
201
+ return False
202
+ except HttpResponseError as exc:
203
+ logger.error("Azure service error deleting entity in %s: %s", table_name, exc.message)
204
+ return False
205
+ except Exception as exc:
206
+ logger.error("Unexpected error deleting entity in %s: %s", table_name, exc)
207
+ return False
208
+
209
+ def get_entity(
210
+ self, table_name: str, partition_key: str, row_key: str
211
+ ) -> Optional[Dict[str, Any]]:
212
+ """Retrieve a single entity by partition and row key.
213
+
214
+ Args:
215
+ table_name: Name of the table
216
+ partition_key: Partition key of the entity
217
+ row_key: Row key of the entity
218
+
219
+ Returns:
220
+ Entity dict if found, None otherwise (including on auth/service errors)
221
+ """
222
+ client = self._get_client()
223
+ table = client.get_table_client(table_name)
224
+ try:
225
+ entity = table.get_entity(partition_key=partition_key, row_key=row_key)
226
+ logger.info("Retrieved entity from %s: PK=%s RK=%s", table_name, partition_key, row_key)
227
+ return dict(entity)
228
+ except ResourceNotFoundError:
229
+ logger.info("Entity not found in %s (%s/%s)", table_name, partition_key, row_key)
230
+ return None
231
+ except ClientAuthenticationError as exc:
232
+ logger.error("Authentication failed retrieving entity from %s: %s", table_name, exc)
233
+ return None
234
+ except HttpResponseError as exc:
235
+ logger.error(
236
+ "Azure service error retrieving entity from %s: %s", table_name, exc.message
237
+ )
238
+ return None
239
+ except Exception as exc:
240
+ logger.error("Unexpected error retrieving entity from %s: %s", table_name, exc)
241
+ return None
242
+
243
+ def entity_exists(self, table_name: str, partition_key: str, row_key: str) -> bool:
244
+ """Check if an entity exists without retrieving it.
245
+
246
+ Args:
247
+ table_name: Name of the table
248
+ partition_key: Partition key of the entity
249
+ row_key: Row key of the entity
250
+
251
+ Returns:
252
+ True if entity exists, False otherwise
253
+ """
254
+ return self.get_entity(table_name, partition_key, row_key) is not None
255
+
256
+ def batch_insert_entities(
257
+ self, table_name: str, entities: List[Dict[str, Any]]
258
+ ) -> Dict[str, Any]:
259
+ """Insert multiple entities in a batch operation.
260
+
261
+ Note: All entities must have the same PartitionKey for batch operations.
262
+
263
+ Args:
264
+ table_name: Name of the table
265
+ entities: List of entity dictionaries to insert
266
+
267
+ Returns:
268
+ Dict with count of successful operations and any errors
269
+ """
270
+ if not entities:
271
+ return {"success": 0, "errors": []}
272
+
273
+ client = self._get_client()
274
+ table = client.get_table_client(table_name)
275
+
276
+ # Group by partition key (batch requirement)
277
+ from collections import defaultdict
278
+
279
+ by_partition: defaultdict[str, List[Dict[str, Any]]] = defaultdict(list)
280
+ for entity in entities:
281
+ pk = entity.get("PartitionKey")
282
+ if not pk:
283
+ logger.error("Entity missing PartitionKey, skipping")
284
+ continue
285
+ by_partition[pk].append(entity)
286
+
287
+ success_count = 0
288
+ errors = []
289
+
290
+ for partition_key, partition_entities in by_partition.items():
291
+ # Process in chunks of 100 (Azure limit)
292
+ for i in range(0, len(partition_entities), 100):
293
+ chunk = partition_entities[i : i + 100]
294
+ operations = [("create", entity) for entity in chunk]
295
+
296
+ try:
297
+ table.submit_transaction(operations)
298
+ success_count += len(chunk)
299
+ logger.info(
300
+ "Batch inserted %d entities into %s (PK=%s)",
301
+ len(chunk),
302
+ table_name,
303
+ partition_key,
304
+ )
305
+ except Exception as exc:
306
+ error_msg = f"Batch insert failed for PK={partition_key}: {exc}"
307
+ logger.error(error_msg)
308
+ errors.append(error_msg)
309
+
310
+ return {"success": success_count, "errors": errors, "ts": _now_iso()}
311
+
312
+ def batch_upsert_entities(
313
+ self, table_name: str, entities: List[Dict[str, Any]]
314
+ ) -> Dict[str, Any]:
315
+ """Upsert multiple entities in a batch operation.
316
+
317
+ Note: All entities must have the same PartitionKey for batch operations.
318
+
319
+ Args:
320
+ table_name: Name of the table
321
+ entities: List of entity dictionaries to upsert
322
+
323
+ Returns:
324
+ Dict with count of successful operations and any errors
325
+ """
326
+ if not entities:
327
+ return {"success": 0, "errors": []}
328
+
329
+ client = self._get_client()
330
+ table = client.get_table_client(table_name)
331
+
332
+ # Group by partition key
333
+ from collections import defaultdict
334
+
335
+ by_partition: defaultdict[str, List[Dict[str, Any]]] = defaultdict(list)
336
+ for entity in entities:
337
+ pk = entity.get("PartitionKey")
338
+ if not pk:
339
+ logger.error("Entity missing PartitionKey, skipping")
340
+ continue
341
+ by_partition[pk].append(entity)
342
+
343
+ success_count = 0
344
+ errors = []
345
+
346
+ for partition_key, partition_entities in by_partition.items():
347
+ # Process in chunks of 100
348
+ for i in range(0, len(partition_entities), 100):
349
+ chunk = partition_entities[i : i + 100]
350
+ operations = [("upsert", entity, {"mode": "merge"}) for entity in chunk]
351
+
352
+ try:
353
+ table.submit_transaction(operations)
354
+ success_count += len(chunk)
355
+ logger.info(
356
+ "Batch upserted %d entities into %s (PK=%s)",
357
+ len(chunk),
358
+ table_name,
359
+ partition_key,
360
+ )
361
+ except Exception as exc:
362
+ error_msg = f"Batch upsert failed for PK={partition_key}: {exc}"
363
+ logger.error(error_msg)
364
+ errors.append(error_msg)
365
+
366
+ return {"success": success_count, "errors": errors, "ts": _now_iso()}
367
+
368
+ def batch_delete_entities(self, table_name: str, keys: List[tuple[str, str]]) -> Dict[str, Any]:
369
+ """Delete multiple entities in a batch operation.
370
+
371
+ Note: All entities must have the same PartitionKey for batch operations.
372
+
373
+ Args:
374
+ table_name: Name of the table
375
+ keys: List of (partition_key, row_key) tuples
376
+
377
+ Returns:
378
+ Dict with count of successful operations and any errors
379
+ """
380
+ if not keys:
381
+ return {"success": 0, "errors": []}
382
+
383
+ client = self._get_client()
384
+ table = client.get_table_client(table_name)
385
+
386
+ # Group by partition key
387
+ from collections import defaultdict
388
+
389
+ by_partition: defaultdict[str, List[tuple[str, str]]] = defaultdict(list)
390
+ for pk, rk in keys:
391
+ by_partition[pk].append((pk, rk))
392
+
393
+ success_count = 0
394
+ errors = []
395
+
396
+ for partition_key, partition_keys in by_partition.items():
397
+ # Process in chunks of 100
398
+ for i in range(0, len(partition_keys), 100):
399
+ chunk = partition_keys[i : i + 100]
400
+ operations = [("delete", {"PartitionKey": pk, "RowKey": rk}) for pk, rk in chunk]
401
+
402
+ try:
403
+ table.submit_transaction(operations)
404
+ success_count += len(chunk)
405
+ logger.info(
406
+ "Batch deleted %d entities from %s (PK=%s)",
407
+ len(chunk),
408
+ table_name,
409
+ partition_key,
410
+ )
411
+ except Exception as exc:
412
+ error_msg = f"Batch delete failed for PK={partition_key}: {exc}"
413
+ logger.error(error_msg)
414
+ errors.append(error_msg)
415
+
416
+ return {"success": success_count, "errors": errors, "ts": _now_iso()}
417
+
418
+ def query(self, table_name: str, filter_query: str) -> Iterable[Dict[str, Any]]:
419
+ """Query entities with a filter.
420
+
421
+ Args:
422
+ table_name: Name of the table
423
+ filter_query: OData filter query string
424
+
425
+ Yields:
426
+ Entity dictionaries matching the filter
427
+ """
428
+ client = self._get_client()
429
+ table = client.get_table_client(table_name)
430
+ for entity in table.query_entities(filter=filter_query):
431
+ yield dict(entity)
432
+
433
+ def query_with_options(
434
+ self,
435
+ table_name: str,
436
+ filter_query: Optional[str] = None,
437
+ select: Optional[List[str]] = None,
438
+ top: Optional[int] = None,
439
+ ) -> Iterable[Dict[str, Any]]:
440
+ """Query entities with advanced options.
441
+
442
+ Args:
443
+ table_name: Name of the table
444
+ filter_query: Optional OData filter query string
445
+ select: Optional list of property names to return (projection)
446
+ top: Optional maximum number of entities to return
447
+
448
+ Yields:
449
+ Entity dictionaries matching the criteria
450
+ """
451
+ client = self._get_client()
452
+ table = client.get_table_client(table_name)
453
+
454
+ kwargs: Dict[str, Any] = {}
455
+ if filter_query:
456
+ kwargs["filter"] = filter_query
457
+ if select:
458
+ kwargs["select"] = select
459
+ if top:
460
+ kwargs["results_per_page"] = top
461
+
462
+ for entity in table.query_entities(**kwargs):
463
+ yield dict(entity)
464
+
465
+ def query_all(
466
+ self,
467
+ table_name: str,
468
+ filter_query: Optional[str] = None,
469
+ select: Optional[List[str]] = None,
470
+ ) -> List[Dict[str, Any]]:
471
+ """Query all entities and return as a list.
472
+
473
+ Warning: This loads all results into memory. Use query() for large result sets.
474
+
475
+ Args:
476
+ table_name: Name of the table
477
+ filter_query: Optional OData filter query string
478
+ select: Optional list of property names to return
479
+
480
+ Returns:
481
+ List of entity dictionaries
482
+ """
483
+ return list(self.query_with_options(table_name, filter_query, select))
484
+
485
+ def count_entities(self, table_name: str, filter_query: Optional[str] = None) -> int:
486
+ """Count entities matching a filter.
487
+
488
+ Args:
489
+ table_name: Name of the table
490
+ filter_query: Optional OData filter query string
491
+
492
+ Returns:
493
+ Count of matching entities
494
+ """
495
+ count = 0
496
+ for _ in self.query_with_options(table_name, filter_query):
497
+ count += 1
498
+ return count
@@ -1,82 +0,0 @@
1
- from __future__ import annotations
2
-
3
- import os
4
- from dataclasses import dataclass
5
- from datetime import datetime, timezone
6
- from typing import Any, Dict, Iterable, Optional
7
-
8
- from loguru import logger
9
-
10
- try:
11
- from azure.data.tables import TableServiceClient
12
- except Exception: # pragma: no cover
13
- TableServiceClient = None # type: ignore
14
-
15
-
16
- def _now_iso() -> str:
17
- return datetime.now(timezone.utc).isoformat()
18
-
19
-
20
- @dataclass
21
- class AzureTableService:
22
- connection_string: Optional[str] = None
23
-
24
- def __init__(self, connection_string: Optional[str] = None) -> None:
25
- # Constructor injection preferred; fallback to env only if not provided
26
- self.connection_string = connection_string or os.getenv("AZURE_STORAGE_CONNECTION_STRING")
27
-
28
- def _get_client(self) -> "TableServiceClient":
29
- if not self.connection_string:
30
- raise RuntimeError("AZURE_STORAGE_CONNECTION_STRING not configured")
31
- if TableServiceClient is None:
32
- raise RuntimeError("azure-data-tables not installed; install to use table operations")
33
- clean = self.connection_string.strip().strip('"').strip("'")
34
- return TableServiceClient.from_connection_string(conn_str=clean)
35
-
36
- def ensure_table(self, table_name: str) -> None:
37
- client = self._get_client()
38
- try:
39
- client.create_table_if_not_exists(table_name=table_name)
40
- except Exception as e:
41
- logger.warning("ensure_table(%s) warning: %s", table_name, e)
42
-
43
- def put_entity(self, table_name: str, entity: Dict[str, Any]) -> Dict[str, Any]:
44
- client = self._get_client()
45
- table = client.get_table_client(table_name)
46
- res = table.create_entity(entity=entity)
47
- logger.info(
48
- "Inserted entity into %s: PK=%s RK=%s",
49
- table_name,
50
- entity.get("PartitionKey"),
51
- entity.get("RowKey"),
52
- )
53
- return {"etag": getattr(res, "etag", None), "ts": _now_iso()}
54
-
55
- def upsert_entity(self, table_name: str, entity: Dict[str, Any]) -> Dict[str, Any]:
56
- client = self._get_client()
57
- table = client.get_table_client(table_name)
58
- res = table.upsert_entity(entity=entity, mode="merge")
59
- logger.info(
60
- "Upserted entity into %s: PK=%s RK=%s",
61
- table_name,
62
- entity.get("PartitionKey"),
63
- entity.get("RowKey"),
64
- )
65
- return {"etag": getattr(res, "etag", None), "ts": _now_iso()}
66
-
67
- def delete_entity(self, table_name: str, partition_key: str, row_key: str) -> bool:
68
- client = self._get_client()
69
- table = client.get_table_client(table_name)
70
- try:
71
- table.delete_entity(partition_key=partition_key, row_key=row_key)
72
- logger.info("Deleted entity in %s (%s/%s)", table_name, partition_key, row_key)
73
- return True
74
- except Exception as exc:
75
- logger.error("Failed to delete entity in %s: %s", table_name, exc)
76
- return False
77
-
78
- def query(self, table_name: str, filter_query: str) -> Iterable[Dict[str, Any]]:
79
- client = self._get_client()
80
- table = client.get_table_client(table_name)
81
- for entity in table.query_entities(filter=filter_query):
82
- yield dict(entity)
File without changes
File without changes