Coverage for src/dataknobs_fsm/functions/library/database.py: 26%

1"""Built-in database functions for FSM.

2

3 This module provides database-related functions that can be referenced

4 in FSM configurations, leveraging the dataknobs_data package.

5"""

6

7 from typing import Any, Dict, List

8

9 from dataknobs_fsm.functions.base import ITransformFunction, TransformError

10 from dataknobs_fsm.resources.database import DatabaseResourceAdapter

11

12

13 class DatabaseFetch(ITransformFunction):

14 """Fetch data from a database using a query."""

15

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.

25

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

38

39 async def transform(self, data: Dict[str, Any]) -> Dict[str, Any]:

40 """Transform data by fetching from database.

41

42 Args:

43 data: Input data (can contain query parameters).

44

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 )

54

55 # Merge parameters

56 query_params = {**self.params}

57

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

63

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 )

72

73 # Return result

74 if self.fetch_one:

75 return {"record": result, **data}

76 else:

77 return {"records": result, **data}

78

79 except Exception as e:

80 raise TransformError(f"Database query failed: {e}") from e

81

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]}..."

85

86

87 class DatabaseUpsert(ITransformFunction):

88 """Upsert data into a database table."""

89

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.

99

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

112

113 async def transform(self, data: Dict[str, Any]) -> Dict[str, Any]:

114 """Transform data by upserting to database.

115

116 Args:

117 data: Input data to upsert.

118

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 )

128

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]

137

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 )

147

148 return {

149 "upserted_count": result.get("affected_rows", 0),

150 **data,

151 }

152

153 except Exception as e:

154 raise TransformError(f"Database upsert failed: {e}") from e

155

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}"

159

160

161 class BatchCommit(ITransformFunction):

162 """Commit a batch of records to the database."""

163

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.

171

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

180

181 async def transform(self, data: Dict[str, Any]) -> Dict[str, Any]:

182 """Transform data by committing batch to database.

183

184 Args:

185 data: Input data containing batch to commit.

186

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 )

196

197 # Get batch from data

198 batch = data.get("batch", [])

199 if not batch:

200 return data

201

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)

212

213 return {

214 "committed_count": committed,

215 "batch": [], # Clear batch after commit

216 **data,

217 }

218

219 except Exception as e:

220 raise TransformError(f"Batch commit failed: {e}") from e

221

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})"

225

226

227 class DatabaseQuery(ITransformFunction):

228 """Execute a dynamic database query."""

229

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.

238

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

249

250 async def transform(self, data: Dict[str, Any]) -> Dict[str, Any]:

251 """Transform data by executing dynamic query.

252

253 Args:

254 data: Input data containing query and parameters.

255

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 )

265

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")

270

271 params = data.get(self.params_field, {})

272

273 try:

274 # Execute query

275 result = await resource.execute_query(query, params=params)

276

277 return {

278 **data,

279 self.result_field: result,

280 }

281

282 except Exception as e:

283 raise TransformError(f"Query execution failed: {e}") from e

284

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}'"

288

289

290 class DatabaseTransaction(ITransformFunction):

291 """Manage database transactions."""

292

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.

300

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

309

310 async def transform(self, data: Dict[str, Any]) -> Dict[str, Any]:

311 """Transform data by managing transaction.

312

313 Args:

314 data: Input data.

315

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 )

325

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 }

334

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 }

344

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 }

354

355 else:

356 raise TransformError(f"Unknown action: {self.action}")

357

358 except Exception as e:

359 raise TransformError(f"Transaction {self.action} failed: {e}") from e

360

361 def get_transform_description(self) -> str:

362 """Get a description of the transformation."""

363 return f"Database transaction: {self.action}"

364

365

366 class DatabaseBulkInsert(ITransformFunction):

367 """Perform bulk insert into database."""

368

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.

378

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

391

392 async def transform(self, data: Dict[str, Any]) -> Dict[str, Any]:

393 """Transform data by performing bulk insert.

394

395 Args:

396 data: Input data containing records to insert.

397

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 )

407

408 # Get records to insert

409 records = data.get("records", [])

410 if not records:

411 return {**data, "inserted_count": 0}

412

413 # Determine columns

414 columns = self.columns

415 if not columns and records:

416 columns = list(records[0].keys())

417

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)

430

431 return {

432 **data,

433 "inserted_count": total_inserted,

434 }

435

436 except Exception as e:

437 raise TransformError(f"Bulk insert failed: {e}") from e

438

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})"

442

443

444# Convenience functions for creating database functions

445 def fetch(resource: str, query: str, **kwargs) -> DatabaseFetch:

446 """Create a DatabaseFetch function."""

447 return DatabaseFetch(resource, query, **kwargs)

448

449

450 def upsert(resource: str, table: str, keys: List[str], **kwargs) -> DatabaseUpsert:

451 """Create a DatabaseUpsert function."""

452 return DatabaseUpsert(resource, table, keys, **kwargs)

453

454

455 def commit_batch(resource: str, **kwargs) -> BatchCommit:

456 """Create a BatchCommit function."""

457 return BatchCommit(resource, **kwargs)

458

459

460 def query(resource: str, **kwargs) -> DatabaseQuery:

461 """Create a DatabaseQuery function."""

462 return DatabaseQuery(resource, **kwargs)

463

464

465 def transaction(resource: str, action: str, **kwargs) -> DatabaseTransaction:

466 """Create a DatabaseTransaction function."""

467 return DatabaseTransaction(resource, action, **kwargs)

468

469

470 def bulk_insert(resource: str, table: str, **kwargs) -> DatabaseBulkInsert:

471 """Create a DatabaseBulkInsert function."""

472 return DatabaseBulkInsert(resource, table, **kwargs)