Coverage for src/dataknobs_fsm/functions/library/database.py: 0%
164 statements
« prev ^ index » next coverage.py v7.10.6, created at 2025-09-20 16:46 -0600
1"""Built-in database functions for FSM.
3This module provides database-related functions that can be referenced
4in FSM configurations, leveraging the dataknobs_data package.
5"""
7from typing import Any, Dict, List
9from dataknobs_fsm.functions.base import ITransformFunction, TransformError
10from dataknobs_fsm.resources.database import DatabaseResourceAdapter
13class DatabaseFetch(ITransformFunction):
14 """Fetch data from a database using a query."""
16 def __init__(
17 self,
18 resource_name: str,
19 query: str,
20 params: Dict[str, Any] | None = None,
21 fetch_one: bool = False,
22 as_dict: bool = True,
23 ):
24 """Initialize the database fetch function.
26 Args:
27 resource_name: Name of the database resource to use.
28 query: SQL query to execute.
29 params: Query parameters for parameterized queries.
30 fetch_one: If True, fetch only one record.
31 as_dict: If True, return records as dictionaries.
32 """
33 self.resource_name = resource_name
34 self.query = query
35 self.params = params or {}
36 self.fetch_one = fetch_one
37 self.as_dict = as_dict
39 async def transform(self, data: Dict[str, Any]) -> Dict[str, Any]:
40 """Transform data by fetching from database.
42 Args:
43 data: Input data (can contain query parameters).
45 Returns:
46 Data with database query results.
47 """
48 # Get resource from context (injected during execution)
49 resource = data.get("_resources", {}).get(self.resource_name)
50 if not resource or not isinstance(resource, DatabaseResourceAdapter):
51 raise TransformError(
52 f"Database resource '{self.resource_name}' not found"
53 )
55 # Merge parameters
56 query_params = {**self.params}
58 # Allow dynamic parameters from input data
59 for key, value in data.items():
60 if key.startswith("param_"):
61 param_name = key[6:] # Remove "param_" prefix
62 query_params[param_name] = value
64 try:
65 # Execute query
66 result = await resource.execute_query(
67 self.query,
68 params=query_params,
69 fetch_one=self.fetch_one,
70 as_dict=self.as_dict,
71 )
73 # Return result
74 if self.fetch_one:
75 return {"record": result, **data}
76 else:
77 return {"records": result, **data}
79 except Exception as e:
80 raise TransformError(f"Database query failed: {e}") from e
82 def get_transform_description(self) -> str:
83 """Get a description of the transformation."""
84 return f"Fetch data from {self.resource_name} using query: {self.query[:50]}..."
87class DatabaseUpsert(ITransformFunction):
88 """Upsert data into a database table."""
90 def __init__(
91 self,
92 resource_name: str,
93 table: str,
94 key_columns: List[str],
95 value_columns: List[str] | None = None,
96 on_conflict: str = "update", # "update", "ignore", "error"
97 ):
98 """Initialize the database upsert function.
100 Args:
101 resource_name: Name of the database resource to use.
102 table: Table name to upsert into.
103 key_columns: Columns that form the unique key.
104 value_columns: Columns to update (if None, update all non-key columns).
105 on_conflict: Action on conflict ("update", "ignore", "error").
106 """
107 self.resource_name = resource_name
108 self.table = table
109 self.key_columns = key_columns
110 self.value_columns = value_columns
111 self.on_conflict = on_conflict
113 async def transform(self, data: Dict[str, Any]) -> Dict[str, Any]:
114 """Transform data by upserting to database.
116 Args:
117 data: Input data to upsert.
119 Returns:
120 Data with upsert result.
121 """
122 # Get resource from context
123 resource = data.get("_resources", {}).get(self.resource_name)
124 if not resource or not isinstance(resource, DatabaseResourceAdapter):
125 raise TransformError(
126 f"Database resource '{self.resource_name}' not found"
127 )
129 # Extract record(s) to upsert
130 if "records" in data:
131 records = data["records"]
132 elif "record" in data:
133 records = [data["record"]]
134 else:
135 # Use the entire data as a single record
136 records = [data]
138 try:
139 # Perform upsert
140 result = await resource.upsert(
141 table=self.table,
142 records=records,
143 key_columns=self.key_columns,
144 value_columns=self.value_columns,
145 on_conflict=self.on_conflict,
146 )
148 return {
149 "upserted_count": result.get("affected_rows", 0),
150 **data,
151 }
153 except Exception as e:
154 raise TransformError(f"Database upsert failed: {e}") from e
156 def get_transform_description(self) -> str:
157 """Get a description of the transformation."""
158 return f"Upsert data into {self.table} table in {self.resource_name}"
161class BatchCommit(ITransformFunction):
162 """Commit a batch of records to the database."""
164 def __init__(
165 self,
166 resource_name: str,
167 batch_size: int = 1000,
168 use_transaction: bool = True,
169 ):
170 """Initialize the batch commit function.
172 Args:
173 resource_name: Name of the database resource to use.
174 batch_size: Number of records per batch.
175 use_transaction: Whether to use a transaction for each batch.
176 """
177 self.resource_name = resource_name
178 self.batch_size = batch_size
179 self.use_transaction = use_transaction
181 async def transform(self, data: Dict[str, Any]) -> Dict[str, Any]:
182 """Transform data by committing batch to database.
184 Args:
185 data: Input data containing batch to commit.
187 Returns:
188 Data with commit result.
189 """
190 # Get resource from context
191 resource = data.get("_resources", {}).get(self.resource_name)
192 if not resource or not isinstance(resource, DatabaseResourceAdapter):
193 raise TransformError(
194 f"Database resource '{self.resource_name}' not found"
195 )
197 # Get batch from data
198 batch = data.get("batch", [])
199 if not batch:
200 return data
202 try:
203 if self.use_transaction:
204 # Commit with transaction
205 async with resource.transaction() as tx:
206 await tx.commit_batch(batch)
207 committed = len(batch)
208 else:
209 # Direct commit
210 result = await resource.commit_batch(batch)
211 committed = result.get("affected_rows", 0)
213 return {
214 "committed_count": committed,
215 "batch": [], # Clear batch after commit
216 **data,
217 }
219 except Exception as e:
220 raise TransformError(f"Batch commit failed: {e}") from e
222 def get_transform_description(self) -> str:
223 """Get a description of the transformation."""
224 return f"Commit batch to {self.resource_name} (batch_size={self.batch_size})"
227class DatabaseQuery(ITransformFunction):
228 """Execute a dynamic database query."""
230 def __init__(
231 self,
232 resource_name: str,
233 query_field: str = "query",
234 params_field: str = "params",
235 result_field: str = "result",
236 ):
237 """Initialize the database query function.
239 Args:
240 resource_name: Name of the database resource to use.
241 query_field: Field containing the query to execute.
242 params_field: Field containing query parameters.
243 result_field: Field to store results in.
244 """
245 self.resource_name = resource_name
246 self.query_field = query_field
247 self.params_field = params_field
248 self.result_field = result_field
250 async def transform(self, data: Dict[str, Any]) -> Dict[str, Any]:
251 """Transform data by executing dynamic query.
253 Args:
254 data: Input data containing query and parameters.
256 Returns:
257 Data with query results.
258 """
259 # Get resource from context
260 resource = data.get("_resources", {}).get(self.resource_name)
261 if not resource or not isinstance(resource, DatabaseResourceAdapter):
262 raise TransformError(
263 f"Database resource '{self.resource_name}' not found"
264 )
266 # Get query and parameters
267 query = data.get(self.query_field)
268 if not query:
269 raise TransformError(f"Query field '{self.query_field}' not found")
271 params = data.get(self.params_field, {})
273 try:
274 # Execute query
275 result = await resource.execute_query(query, params=params)
277 return {
278 **data,
279 self.result_field: result,
280 }
282 except Exception as e:
283 raise TransformError(f"Query execution failed: {e}") from e
285 def get_transform_description(self) -> str:
286 """Get a description of the transformation."""
287 return f"Execute dynamic query from field '{self.query_field}'"
290class DatabaseTransaction(ITransformFunction):
291 """Manage database transactions."""
293 def __init__(
294 self,
295 resource_name: str,
296 action: str = "begin", # "begin", "commit", "rollback"
297 savepoint: str | None = None,
298 ):
299 """Initialize the database transaction function.
301 Args:
302 resource_name: Name of the database resource to use.
303 action: Transaction action to perform.
304 savepoint: Optional savepoint name.
305 """
306 self.resource_name = resource_name
307 self.action = action
308 self.savepoint = savepoint
310 async def transform(self, data: Dict[str, Any]) -> Dict[str, Any]:
311 """Transform data by managing transaction.
313 Args:
314 data: Input data.
316 Returns:
317 Data with transaction status.
318 """
319 # Get resource from context
320 resource = data.get("_resources", {}).get(self.resource_name)
321 if not resource or not isinstance(resource, DatabaseResourceAdapter):
322 raise TransformError(
323 f"Database resource '{self.resource_name}' not found"
324 )
326 try:
327 if self.action == "begin":
328 tx = await resource.begin_transaction()
329 return {
330 **data,
331 "_transaction": tx,
332 "transaction_active": True,
333 }
335 elif self.action == "commit":
336 tx = data.get("_transaction")
337 if tx:
338 await tx.commit()
339 return {
340 **data,
341 "_transaction": None,
342 "transaction_active": False,
343 }
345 elif self.action == "rollback":
346 tx = data.get("_transaction")
347 if tx:
348 await tx.rollback()
349 return {
350 **data,
351 "_transaction": None,
352 "transaction_active": False,
353 }
355 else:
356 raise TransformError(f"Unknown action: {self.action}")
358 except Exception as e:
359 raise TransformError(f"Transaction {self.action} failed: {e}") from e
361 def get_transform_description(self) -> str:
362 """Get a description of the transformation."""
363 return f"Database transaction: {self.action}"
366class DatabaseBulkInsert(ITransformFunction):
367 """Perform bulk insert into database."""
369 def __init__(
370 self,
371 resource_name: str,
372 table: str,
373 columns: List[str] | None = None,
374 chunk_size: int = 1000,
375 on_duplicate: str = "error", # "error", "ignore", "update"
376 ):
377 """Initialize the bulk insert function.
379 Args:
380 resource_name: Name of the database resource to use.
381 table: Table to insert into.
382 columns: Columns to insert (if None, use all columns from first record).
383 chunk_size: Number of records per chunk.
384 on_duplicate: Action on duplicate key.
385 """
386 self.resource_name = resource_name
387 self.table = table
388 self.columns = columns
389 self.chunk_size = chunk_size
390 self.on_duplicate = on_duplicate
392 async def transform(self, data: Dict[str, Any]) -> Dict[str, Any]:
393 """Transform data by performing bulk insert.
395 Args:
396 data: Input data containing records to insert.
398 Returns:
399 Data with insert results.
400 """
401 # Get resource from context
402 resource = data.get("_resources", {}).get(self.resource_name)
403 if not resource or not isinstance(resource, DatabaseResourceAdapter):
404 raise TransformError(
405 f"Database resource '{self.resource_name}' not found"
406 )
408 # Get records to insert
409 records = data.get("records", [])
410 if not records:
411 return {**data, "inserted_count": 0}
413 # Determine columns
414 columns = self.columns
415 if not columns and records:
416 columns = list(records[0].keys())
418 try:
419 # Perform bulk insert in chunks
420 total_inserted = 0
421 for i in range(0, len(records), self.chunk_size):
422 chunk = records[i:i + self.chunk_size]
423 result = await resource.bulk_insert(
424 table=self.table,
425 records=chunk,
426 columns=columns,
427 on_duplicate=self.on_duplicate,
428 )
429 total_inserted += result.get("affected_rows", 0)
431 return {
432 **data,
433 "inserted_count": total_inserted,
434 }
436 except Exception as e:
437 raise TransformError(f"Bulk insert failed: {e}") from e
439 def get_transform_description(self) -> str:
440 """Get a description of the transformation."""
441 return f"Bulk insert into {self.table} table (chunk_size={self.chunk_size})"
444# Convenience functions for creating database functions
445def fetch(resource: str, query: str, **kwargs) -> DatabaseFetch:
446 """Create a DatabaseFetch function."""
447 return DatabaseFetch(resource, query, **kwargs)
450def upsert(resource: str, table: str, keys: List[str], **kwargs) -> DatabaseUpsert:
451 """Create a DatabaseUpsert function."""
452 return DatabaseUpsert(resource, table, keys, **kwargs)
455def commit_batch(resource: str, **kwargs) -> BatchCommit:
456 """Create a BatchCommit function."""
457 return BatchCommit(resource, **kwargs)
460def query(resource: str, **kwargs) -> DatabaseQuery:
461 """Create a DatabaseQuery function."""
462 return DatabaseQuery(resource, **kwargs)
465def transaction(resource: str, action: str, **kwargs) -> DatabaseTransaction:
466 """Create a DatabaseTransaction function."""
467 return DatabaseTransaction(resource, action, **kwargs)
470def bulk_insert(resource: str, table: str, **kwargs) -> DatabaseBulkInsert:
471 """Create a DatabaseBulkInsert function."""
472 return DatabaseBulkInsert(resource, table, **kwargs)