dbagent-cli 0.3.0__tar.gz → 0.3.1__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.
Files changed (73) hide show
  1. {dbagent_cli-0.3.0 → dbagent_cli-0.3.1}/PKG-INFO +2 -2
  2. {dbagent_cli-0.3.0 → dbagent_cli-0.3.1}/dbagent/__init__.py +1 -1
  3. {dbagent_cli-0.3.0 → dbagent_cli-0.3.1}/dbagent/agent/pipeline.py +43 -1
  4. {dbagent_cli-0.3.0 → dbagent_cli-0.3.1}/dbagent/cli.py +175 -43
  5. {dbagent_cli-0.3.0 → dbagent_cli-0.3.1}/dbagent/connectors/base.py +11 -0
  6. {dbagent_cli-0.3.0 → dbagent_cli-0.3.1}/dbagent/connectors/relational.py +11 -0
  7. {dbagent_cli-0.3.0 → dbagent_cli-0.3.1}/dbagent/execution/changeset.py +77 -49
  8. dbagent_cli-0.3.1/dbagent/execution/recovery.py +571 -0
  9. {dbagent_cli-0.3.0 → dbagent_cli-0.3.1}/dbagent_cli.egg-info/PKG-INFO +2 -2
  10. {dbagent_cli-0.3.0 → dbagent_cli-0.3.1}/dbagent_cli.egg-info/SOURCES.txt +2 -0
  11. {dbagent_cli-0.3.0 → dbagent_cli-0.3.1}/pyproject.toml +2 -2
  12. {dbagent_cli-0.3.0 → dbagent_cli-0.3.1}/tests/test_cli.py +13 -0
  13. dbagent_cli-0.3.1/tests/test_recovery.py +265 -0
  14. {dbagent_cli-0.3.0 → dbagent_cli-0.3.1}/tests/test_sandbox.py +1 -1
  15. {dbagent_cli-0.3.0 → dbagent_cli-0.3.1}/README.md +0 -0
  16. {dbagent_cli-0.3.0 → dbagent_cli-0.3.1}/dbagent/agent/confidence.py +0 -0
  17. {dbagent_cli-0.3.0 → dbagent_cli-0.3.1}/dbagent/agent/conversation.py +0 -0
  18. {dbagent_cli-0.3.0 → dbagent_cli-0.3.1}/dbagent/agent/doctor.py +0 -0
  19. {dbagent_cli-0.3.0 → dbagent_cli-0.3.1}/dbagent/agent/error_classifier.py +0 -0
  20. {dbagent_cli-0.3.0 → dbagent_cli-0.3.1}/dbagent/agent/generator.py +0 -0
  21. {dbagent_cli-0.3.0 → dbagent_cli-0.3.1}/dbagent/agent/recommender.py +0 -0
  22. {dbagent_cli-0.3.0 → dbagent_cli-0.3.1}/dbagent/agent/validator.py +0 -0
  23. {dbagent_cli-0.3.0 → dbagent_cli-0.3.1}/dbagent/config.py +0 -0
  24. {dbagent_cli-0.3.0 → dbagent_cli-0.3.1}/dbagent/connectors/factory.py +0 -0
  25. {dbagent_cli-0.3.0 → dbagent_cli-0.3.1}/dbagent/connectors/mongo.py +0 -0
  26. {dbagent_cli-0.3.0 → dbagent_cli-0.3.1}/dbagent/execution/__init__.py +0 -0
  27. {dbagent_cli-0.3.0 → dbagent_cli-0.3.1}/dbagent/execution/cache.py +0 -0
  28. {dbagent_cli-0.3.0 → dbagent_cli-0.3.1}/dbagent/execution/cost.py +0 -0
  29. {dbagent_cli-0.3.0 → dbagent_cli-0.3.1}/dbagent/execution/sandbox.py +0 -0
  30. {dbagent_cli-0.3.0 → dbagent_cli-0.3.1}/dbagent/execution/streaming.py +0 -0
  31. {dbagent_cli-0.3.0 → dbagent_cli-0.3.1}/dbagent/llm/auto_setup.py +0 -0
  32. {dbagent_cli-0.3.0 → dbagent_cli-0.3.1}/dbagent/llm/base.py +0 -0
  33. {dbagent_cli-0.3.0 → dbagent_cli-0.3.1}/dbagent/llm/factory.py +0 -0
  34. {dbagent_cli-0.3.0 → dbagent_cli-0.3.1}/dbagent/llm/gemini_provider.py +0 -0
  35. {dbagent_cli-0.3.0 → dbagent_cli-0.3.1}/dbagent/llm/groq_provider.py +0 -0
  36. {dbagent_cli-0.3.0 → dbagent_cli-0.3.1}/dbagent/llm/mock_provider.py +0 -0
  37. {dbagent_cli-0.3.0 → dbagent_cli-0.3.1}/dbagent/llm/ollama_provider.py +0 -0
  38. {dbagent_cli-0.3.0 → dbagent_cli-0.3.1}/dbagent/llm/openrouter_provider.py +0 -0
  39. {dbagent_cli-0.3.0 → dbagent_cli-0.3.1}/dbagent/schema/cache.py +0 -0
  40. {dbagent_cli-0.3.0 → dbagent_cli-0.3.1}/dbagent/schema/formatter.py +0 -0
  41. {dbagent_cli-0.3.0 → dbagent_cli-0.3.1}/dbagent/schema/graph.py +0 -0
  42. {dbagent_cli-0.3.0 → dbagent_cli-0.3.1}/dbagent/schema/models.py +0 -0
  43. {dbagent_cli-0.3.0 → dbagent_cli-0.3.1}/dbagent/schema/retriever.py +0 -0
  44. {dbagent_cli-0.3.0 → dbagent_cli-0.3.1}/dbagent/schema/selector.py +0 -0
  45. {dbagent_cli-0.3.0 → dbagent_cli-0.3.1}/dbagent/security/__init__.py +0 -0
  46. {dbagent_cli-0.3.0 → dbagent_cli-0.3.1}/dbagent/security/audit.py +0 -0
  47. {dbagent_cli-0.3.0 → dbagent_cli-0.3.1}/dbagent/security/environment.py +0 -0
  48. {dbagent_cli-0.3.0 → dbagent_cli-0.3.1}/dbagent/security/masking.py +0 -0
  49. {dbagent_cli-0.3.0 → dbagent_cli-0.3.1}/dbagent/security/modes.py +0 -0
  50. {dbagent_cli-0.3.0 → dbagent_cli-0.3.1}/dbagent/security/rbac.py +0 -0
  51. {dbagent_cli-0.3.0 → dbagent_cli-0.3.1}/dbagent/semantic/__init__.py +0 -0
  52. {dbagent_cli-0.3.0 → dbagent_cli-0.3.1}/dbagent/semantic/dictionary.py +0 -0
  53. {dbagent_cli-0.3.0 → dbagent_cli-0.3.1}/dbagent/semantic/templates.py +0 -0
  54. {dbagent_cli-0.3.0 → dbagent_cli-0.3.1}/dbagent/sql/__init__.py +0 -0
  55. {dbagent_cli-0.3.0 → dbagent_cli-0.3.1}/dbagent/sql/optimizer.py +0 -0
  56. {dbagent_cli-0.3.0 → dbagent_cli-0.3.1}/dbagent/sql/validator.py +0 -0
  57. {dbagent_cli-0.3.0 → dbagent_cli-0.3.1}/dbagent/ui/console.py +0 -0
  58. {dbagent_cli-0.3.0 → dbagent_cli-0.3.1}/dbagent/ui/viewer.py +0 -0
  59. {dbagent_cli-0.3.0 → dbagent_cli-0.3.1}/dbagent/ui/visualizer.py +0 -0
  60. {dbagent_cli-0.3.0 → dbagent_cli-0.3.1}/dbagent_cli.egg-info/dependency_links.txt +0 -0
  61. {dbagent_cli-0.3.0 → dbagent_cli-0.3.1}/dbagent_cli.egg-info/entry_points.txt +0 -0
  62. {dbagent_cli-0.3.0 → dbagent_cli-0.3.1}/dbagent_cli.egg-info/requires.txt +0 -0
  63. {dbagent_cli-0.3.0 → dbagent_cli-0.3.1}/dbagent_cli.egg-info/top_level.txt +0 -0
  64. {dbagent_cli-0.3.0 → dbagent_cli-0.3.1}/setup.cfg +0 -0
  65. {dbagent_cli-0.3.0 → dbagent_cli-0.3.1}/tests/test_cache.py +0 -0
  66. {dbagent_cli-0.3.0 → dbagent_cli-0.3.1}/tests/test_connectors.py +0 -0
  67. {dbagent_cli-0.3.0 → dbagent_cli-0.3.1}/tests/test_doctor_and_viz.py +0 -0
  68. {dbagent_cli-0.3.0 → dbagent_cli-0.3.1}/tests/test_generator.py +0 -0
  69. {dbagent_cli-0.3.0 → dbagent_cli-0.3.1}/tests/test_graph.py +0 -0
  70. {dbagent_cli-0.3.0 → dbagent_cli-0.3.1}/tests/test_safety.py +0 -0
  71. {dbagent_cli-0.3.0 → dbagent_cli-0.3.1}/tests/test_schema.py +0 -0
  72. {dbagent_cli-0.3.0 → dbagent_cli-0.3.1}/tests/test_security.py +0 -0
  73. {dbagent_cli-0.3.0 → dbagent_cli-0.3.1}/tests/test_semantic.py +0 -0
@@ -1,7 +1,7 @@
1
1
  Metadata-Version: 2.4
2
2
  Name: dbagent-cli
3
- Version: 0.3.0
4
- Summary: Enterprise AI Database Copilot & Autonomous Data Agent (CLI)
3
+ Version: 0.3.1
4
+ Summary: Enterprise AI Database Copilot with Long-Term Change Recovery (CLI)
5
5
  Author: Santhosh Gupta
6
6
  Requires-Python: >=3.9
7
7
  Description-Content-Type: text/markdown
@@ -2,5 +2,5 @@
2
2
  DB-Agent: Universal Database Introspection and Script Generation AI Agent (CLI).
3
3
  """
4
4
 
5
- __version__ = "0.3.0"
5
+ __version__ = "0.3.1"
6
6
  __author__ = "Santhosh Gupta"
@@ -25,6 +25,7 @@ from dbagent.sql.validator import SQLValidator
25
25
  from dbagent.execution.cost import QueryCostEstimator, QueryComplexity
26
26
  from dbagent.execution.cache import QueryResultCache
27
27
  from dbagent.semantic.dictionary import BusinessDictionary
28
+ from dbagent.execution.recovery import RecoveryManager, RowSnapshotRecoveryStrategy, RecoveryStatus
28
29
  from dbagent.agent.confidence import ConfidenceEvaluator, ConfidenceScore
29
30
  from dbagent.agent.conversation import ConversationContext
30
31
  from dbagent.agent.error_classifier import ErrorClassifier
@@ -45,6 +46,11 @@ class PipelineResult:
45
46
  needs_confirmation: bool = False
46
47
  is_cached: bool = False
47
48
  cache_age_seconds: float = 0.0
49
+ # Long-term Change Recovery
50
+ change_id: Optional[str] = None
51
+ recovery_status: Optional[str] = None
52
+ recovery_message: Optional[str] = None
53
+ expires_at: Optional[str] = None
48
54
  # Conversational chat response
49
55
  is_chat: bool = False
50
56
  message: str = ""
@@ -100,6 +106,7 @@ class QueryPipeline:
100
106
  confirm_callback: Optional[Callable[[str, List[str]], bool]] = None,
101
107
  choice_callback: Optional[Callable[[str, List[str]], Optional[str]]] = None,
102
108
  conversation_context: Optional[ConversationContext] = None,
109
+ recovery_manager: Optional[RecoveryManager] = None,
103
110
  ):
104
111
  self.connector = connector
105
112
  self.generator = ScriptGenerator(llm)
@@ -115,6 +122,7 @@ class QueryPipeline:
115
122
  self.dictionary = BusinessDictionary()
116
123
  self.result_cache = QueryResultCache()
117
124
  self.audit_logger = AuditLogger()
125
+ self.recovery_mgr = recovery_manager or RecoveryManager()
118
126
  self.conversation = conversation_context or ConversationContext()
119
127
 
120
128
  # Database identifier
@@ -272,10 +280,44 @@ class QueryPipeline:
272
280
  elif self.mode == OperationMode.ADMIN:
273
281
  should_execute = True
274
282
 
275
- # --- 11. Execute with Error Retry ---
283
+ # --- 11. Pre-Write Recovery Snapshot Capture & Execution ---
284
+ snapshot_data = None
285
+ target_table_name = exact[0] if exact else (sub_schema.tables[0].name if sub_schema.tables else "unknown")
286
+ table_model = sub_schema.get_table(target_table_name) if sub_schema else None
287
+
288
+ if should_execute and result.query_type in ("write", "ddl") and table_model:
289
+ can_snap, snap, snap_err = RowSnapshotRecoveryStrategy.capture_pre_write_snapshot(
290
+ connector=self.connector,
291
+ table_model=table_model,
292
+ sql=result.sql,
293
+ operation=result.query_type.upper(),
294
+ )
295
+ if can_snap and snap:
296
+ snapshot_data = snap
297
+
276
298
  if should_execute:
277
299
  self._execute_with_retry(result, augmented_prompt, sub_schema)
278
300
 
301
+ # --- Record ChangeSet / Recovery Record for successful write operations ---
302
+ if result.success and result.query_type in ("write", "ddl"):
303
+ affected = len(snapshot_data.get("rows", [])) if snapshot_data else 1
304
+ meta = self.recovery_mgr.create_change_record(
305
+ database=self.db_name,
306
+ sql=result.sql,
307
+ table_name=target_table_name,
308
+ operation=result.query_type.upper(),
309
+ affected_rows=affected,
310
+ snapshot_data=snapshot_data,
311
+ user=self.user_id,
312
+ )
313
+ result.change_id = meta["change_id"]
314
+ result.recovery_status = meta["recovery_status"]
315
+ result.expires_at = meta["expires_at"]
316
+ if meta["recovery_status"] == RecoveryStatus.AVAILABLE.value:
317
+ result.recovery_message = f"Available until {meta['expires_at']} ({self.recovery_mgr.retention_days} days retention)"
318
+ else:
319
+ result.recovery_message = "Not available (unsupported operation or missing primary key)"
320
+
279
321
  # --- 12. Sensitive Data Masking ---
280
322
  if result.success and self.mask_sensitive_data and result.rows:
281
323
  result.rows = SensitiveDataMasker.mask_row_set(result.columns, result.rows)
@@ -36,6 +36,14 @@ from dbagent.security.modes import OperationMode
36
36
  from dbagent.security.audit import AuditLogger
37
37
  from dbagent.security.masking import SensitiveDataMasker
38
38
  from dbagent.execution.changeset import ChangeSetManager
39
+ from dbagent.execution.recovery import (
40
+ RecoveryManager,
41
+ RowSnapshotRecoveryStrategy,
42
+ RecoveryStatus,
43
+ ConflictPolicy,
44
+ RollbackPreview,
45
+ RollbackResult,
46
+ )
39
47
  from dbagent.semantic.dictionary import BusinessDictionary
40
48
  from dbagent.semantic.templates import TemplateManager
41
49
  from dbagent.ui.visualizer import ResultVisualizer
@@ -59,6 +67,7 @@ app = typer.Typer(
59
67
 
60
68
  config_mgr = ConfigManager()
61
69
  cache_mgr = SchemaCacheManager()
70
+ recovery_mgr = RecoveryManager()
62
71
  change_mgr = ChangeSetManager()
63
72
  dict_mgr = BusinessDictionary()
64
73
  template_mgr = TemplateManager()
@@ -169,6 +178,16 @@ def _display_pipeline_result(result: PipelineResult, show_sql: bool = True) -> N
169
178
 
170
179
  if result.retries > 0:
171
180
  console.print(f"[dim](Auto-corrected after {result.retries} retry(ies))[/dim]")
181
+
182
+ # Display recovery info for write changes
183
+ if result.change_id:
184
+ console.print(f"\n[bold green]✓ Change committed successfully[/bold green]")
185
+ console.print(f" * Change ID: [bold white]{result.change_id}[/bold white]")
186
+ if result.recovery_status == "AVAILABLE":
187
+ console.print(f" * Recovery: [bold cyan]{result.recovery_message}[/bold cyan]")
188
+ console.print(f" * Rollback command: [bold]db-agent rollback {result.change_id}[/bold]\n")
189
+ else:
190
+ console.print(f" * Recovery: [yellow]{result.recovery_message or result.recovery_status}[/yellow]\n")
172
191
  elif result.needs_confirmation and not result.was_executed:
173
192
  if result.sql:
174
193
  print_code(result.sql, "sql", f"Generated {result.query_type.upper()} Query")
@@ -931,22 +950,34 @@ def schema_status_command(
931
950
  def changes_command(
932
951
  limit: int = typer.Option(20, "--limit", "-l", help="Number of change sets to show"),
933
952
  ):
934
- """List recorded write change sets."""
953
+ """List recorded write change sets and recovery availability."""
935
954
  print_banner()
936
- sets = change_mgr.list_change_sets(limit=limit)
937
- if not sets:
955
+ records = recovery_mgr.list_changes(limit=limit)
956
+ if not records:
938
957
  print_info("No recorded write change sets found.")
939
958
  return
940
959
 
941
- table = Table(title="[bold cyan]Recorded Write ChangeSets[/bold cyan]", show_header=True, header_style="bold cyan")
942
- table.add_column("ChangeSet ID", style="bold white")
960
+ table = Table(title="[bold cyan]Recorded Write ChangeSets & Recovery Status[/bold cyan]", show_header=True, header_style="bold cyan")
961
+ table.add_column("Change ID", style="bold white")
943
962
  table.add_column("Timestamp", style="dim")
944
963
  table.add_column("Operation", style="yellow")
945
964
  table.add_column("Table", style="cyan")
946
- table.add_column("Affected Rows", justify="right")
965
+ table.add_column("Rows", justify="right")
966
+ table.add_column("Recovery Status", style="bold")
967
+ table.add_column("Expires", style="dim")
947
968
 
948
- for s in sets:
949
- table.add_row(s.change_set_id, s.timestamp, s.operation, s.table_name, str(s.affected_rows))
969
+ for s in records:
970
+ st = s.get("recovery_status", "AVAILABLE")
971
+ st_style = "[green]AVAILABLE[/green]" if st == "AVAILABLE" else ("[yellow]EXPIRED[/yellow]" if st == "EXPIRED" else f"[dim]{st}[/dim]")
972
+ table.add_row(
973
+ s.get("change_id", ""),
974
+ s.get("executed_at", ""),
975
+ s.get("operation", ""),
976
+ s.get("table_name", ""),
977
+ str(s.get("affected_rows", 0)),
978
+ st_style,
979
+ s.get("expires_at", "-"),
980
+ )
950
981
 
951
982
  console.print(table)
952
983
 
@@ -955,64 +986,165 @@ def changes_command(
955
986
 
956
987
  @app.command(name="changes-show")
957
988
  def changes_show_command(
958
- change_set_id: str = typer.Argument(..., help="ChangeSet ID (e.g. CS-20260828-001)"),
989
+ change_id: str = typer.Argument(..., help="Change ID (e.g. CHG-20260828-001)"),
959
990
  ):
960
991
  """Show details and before-snapshot of a recorded ChangeSet."""
961
992
  print_banner()
962
- cs = change_mgr.get_change_set(change_set_id)
963
- if not cs:
964
- print_error(f"ChangeSet '{change_set_id}' not found.")
993
+ meta = recovery_mgr.get_change_record(change_id)
994
+ if not meta:
995
+ print_error(f"Change record '{change_id}' not found.")
965
996
  raise typer.Exit(1)
966
997
 
967
- console.print(f"\n[bold cyan]ChangeSet: {cs.change_set_id}[/bold cyan]")
968
- console.print(f" * Database: {cs.database} | Timestamp: {cs.timestamp}")
969
- console.print(f" * Operation: [{cs.operation}] on Table `{cs.table_name}` ({cs.affected_rows} rows)")
970
- print_code(cs.sql, "sql", "Executed Statement")
998
+ console.print(f"\n[bold cyan]Change Record: {meta['change_id']}[/bold cyan]")
999
+ console.print(f" * Database: {meta.get('database')} | Executed: {meta.get('executed_at')}")
1000
+ console.print(f" * Operation: [{meta.get('operation')}] on Table `{meta.get('table_name')}` ({meta.get('affected_rows')} rows)")
1001
+ console.print(f" * Recovery Status: [bold green]{meta.get('recovery_status')}[/bold green] (Expires: {meta.get('expires_at')})")
1002
+ if meta.get("rollback_of"):
1003
+ console.print(f" * Rollback Of: [bold yellow]{meta['rollback_of']}[/bold yellow]")
1004
+ print_code(meta.get("sql", ""), "sql", "Executed Statement")
971
1005
 
972
- if cs.before_snapshot:
973
- console.print(f"\n[bold yellow]Before-Snapshot ({len(cs.before_snapshot)} rows captured for recovery):[/bold yellow]")
974
- cols = list(cs.before_snapshot[0].keys())
975
- print_results_table(cols, cs.before_snapshot, title="Before Snapshot")
1006
+ snap = recovery_mgr.get_snapshot(change_id)
1007
+ if snap and "rows" in snap and snap["rows"]:
1008
+ rows_disp = [r.get("before", {}) for r in snap["rows"][:20]]
1009
+ if rows_disp and rows_disp[0]:
1010
+ console.print(f"\n[bold yellow]Captured Before-Snapshot ({len(snap['rows'])} rows captured):[/bold yellow]")
1011
+ cols = list(rows_disp[0].keys())
1012
+ print_results_table(cols, rows_disp, title="Snapshot Sample (First 20)")
976
1013
 
977
1014
 
978
1015
  # ---- rollback ----
979
1016
 
980
1017
  @app.command(name="rollback")
981
1018
  def rollback_command(
982
- change_set_id: str = typer.Argument(..., help="ChangeSet ID to roll back"),
1019
+ change_id: Optional[str] = typer.Argument(None, help="Change ID to roll back (e.g. CHG-20260828-001) or use --last"),
983
1020
  db: Optional[str] = typer.Option(None, "--db", "-d", help="Database connection URL or profile name"),
984
- force: bool = typer.Option(False, "--force", "-f", help="Bypass confirmation"),
1021
+ preview: bool = typer.Option(False, "--preview", "-p", help="Preview rollback and detect conflicts without executing"),
1022
+ force: bool = typer.Option(False, "--force", "-f", help="Bypass confirmation and overwrite conflicts"),
1023
+ policy: str = typer.Option("skip", "--conflict-policy", help="Conflict resolution policy: skip, abort, force"),
1024
+ last: bool = typer.Option(False, "--last", "-l", help="Roll back the most recent committed change"),
1025
+ mode: str = typer.Option("CONFIRM", "--mode", help="Operation mode (SAFE, READ_ONLY, CONFIRM, ADMIN)"),
985
1026
  ):
986
- """Generate and execute recovery rollback SQL for a ChangeSet."""
1027
+ """Safely revert and recover committed database modifications with conflict protection."""
987
1028
  print_banner()
988
- cs = change_mgr.get_change_set(change_set_id)
989
- if not cs:
990
- print_error(f"ChangeSet '{change_set_id}' not found.")
1029
+
1030
+ target_id = change_id
1031
+ if last or target_id == "--last":
1032
+ records = recovery_mgr.list_changes(limit=1)
1033
+ if not records:
1034
+ print_error("No recorded change sets available to roll back.")
1035
+ raise typer.Exit(1)
1036
+ target_id = records[0]["change_id"]
1037
+
1038
+ if not target_id:
1039
+ records = recovery_mgr.list_changes(limit=10)
1040
+ if not records:
1041
+ print_info("No recorded change sets found.")
1042
+ raise typer.Exit(0)
1043
+ console.print("[bold cyan]Recent Change Sets:[/bold cyan]")
1044
+ for r in records:
1045
+ console.print(f" * [bold white]{r['change_id']}[/bold white] ({r.get('executed_at')}) -> [{r.get('operation')}] on {r.get('table_name')} ({r.get('affected_rows')} rows)")
1046
+ target_id = Prompt.ask("\nEnter Change ID to roll back")
1047
+
1048
+ meta = recovery_mgr.get_change_record(target_id)
1049
+ if not meta:
1050
+ print_error(f"Change record '{target_id}' not found.")
1051
+ raise typer.Exit(1)
1052
+
1053
+ if meta.get("recovery_status") == RecoveryStatus.EXPIRED.value:
1054
+ print_error(f"Rollback rejected: Recovery snapshot for {target_id} has EXPIRED ({meta.get('expires_at')}).")
1055
+ raise typer.Exit(1)
1056
+
1057
+ if meta.get("recovery_status") == RecoveryStatus.UNAVAILABLE.value:
1058
+ print_error(f"Rollback unavailable: Operation was executed without long-term recovery snapshot.")
1059
+ raise typer.Exit(1)
1060
+
1061
+ snap = recovery_mgr.get_snapshot(target_id)
1062
+ if not snap:
1063
+ print_error(f"Recovery snapshot corrupted or missing checksum verification for '{target_id}'.")
1064
+ raise typer.Exit(1)
1065
+
1066
+ db_url = resolve_db_url(db or meta.get("database"))
1067
+ connector = create_connector(db_url)
1068
+
1069
+ # Concurrency & Conflict Preview
1070
+ with console.status("[bold cyan]Evaluating row eligibility & detecting conflicts...[/bold cyan]"):
1071
+ prev = RowSnapshotRecoveryStrategy.preview_recovery(connector, meta, snap)
1072
+
1073
+ console.print(f"\n{prev.format_summary()}\n")
1074
+
1075
+ if not prev.can_proceed:
1076
+ connector.close()
991
1077
  raise typer.Exit(1)
992
1078
 
993
- stmts = change_mgr.generate_rollback_sql(cs)
994
- if not stmts:
995
- print_warning(f"No recovery snapshot available to roll back ChangeSet '{change_set_id}'.")
1079
+ # If preview only, stop here
1080
+ if preview:
1081
+ console.print("[dim]Preview complete. Use 'db-agent rollback " + target_id + "' to execute recovery.[/dim]\n")
1082
+ connector.close()
996
1083
  return
997
1084
 
998
- full_sql = "\n".join(stmts)
999
- console.print(f"\n[bold cyan]Generated Rollback SQL for {change_set_id}:[/bold cyan]")
1000
- print_code(full_sql, "sql", "Rollback Operations")
1085
+ # Check Safe Mode
1086
+ if mode.upper() in ("SAFE", "READ_ONLY"):
1087
+ print_error(f"Rollback execution blocked in {mode.upper()} mode. Rollback requires write permission. Run with --preview to inspect.")
1088
+ connector.close()
1089
+ raise typer.Exit(1)
1090
+
1091
+ # Conflict Policy enforcement
1092
+ pol_enum = ConflictPolicy(policy.upper()) if policy.upper() in ConflictPolicy.__members__ else ConflictPolicy.SKIP
1093
+ if force:
1094
+ pol_enum = ConflictPolicy.FORCE
1095
+
1096
+ if prev.conflicts_detected > 0 and pol_enum == ConflictPolicy.ABORT:
1097
+ print_error(f"Rollback aborted: {prev.conflicts_detected} concurrency conflict(s) detected with policy=ABORT.")
1098
+ connector.close()
1099
+ raise typer.Exit(1)
1001
1100
 
1002
- if not force and not Confirm.ask("Execute these rollback statements on the database?"):
1101
+ if not force and not Confirm.ask(f"Execute rollback for {target_id} ({prev.rows_eligible} rows eligible)?", default=False):
1003
1102
  console.print("[dim]Rollback cancelled.[/dim]")
1103
+ connector.close()
1004
1104
  return
1005
1105
 
1006
- db_url = resolve_db_url(db or cs.database)
1007
- with console.status("[bold cyan]Executing rollback...[/bold cyan]"):
1008
- try:
1009
- connector = create_connector(db_url)
1010
- for stmt in stmts:
1011
- connector.execute_query(stmt)
1012
- connector.close()
1013
- print_success(f"Rollback completed successfully for {change_set_id}!")
1014
- except Exception as e:
1015
- print_error(f"Rollback execution failed: {str(e)}")
1106
+ with console.status("[bold cyan]Executing rollback recovery...[/bold cyan]"):
1107
+ res = RowSnapshotRecoveryStrategy.execute_recovery(connector, meta, snap, policy=pol_enum)
1108
+ connector.close()
1109
+
1110
+ # Record rollback audit chain
1111
+ recovery_mgr.create_change_record(
1112
+ database=meta.get("database", "database"),
1113
+ sql=f"-- Rollback of {target_id}",
1114
+ table_name=meta.get("table_name", "unknown"),
1115
+ operation="ROLLBACK",
1116
+ affected_rows=res.restored_count,
1117
+ snapshot_data=None,
1118
+ rollback_of=target_id,
1119
+ )
1120
+
1121
+ console.print(f"\n{res.format_summary()}\n")
1122
+
1123
+
1124
+ # ---- recovery ----
1125
+
1126
+ @app.command(name="recovery")
1127
+ def recovery_command(
1128
+ action: str = typer.Argument("list", help="Action: list, cleanup, status"),
1129
+ limit: int = typer.Option(20, "--limit", "-l", help="Number of records to list"),
1130
+ ):
1131
+ """Manage database recovery snapshots, retention, and storage."""
1132
+ print_banner()
1133
+ if action == "list":
1134
+ changes_command(limit=limit)
1135
+ elif action == "cleanup":
1136
+ with console.status("[bold cyan]Cleaning up expired recovery snapshots...[/bold cyan]"):
1137
+ cleaned = recovery_mgr.cleanup_expired()
1138
+ print_success(f"Recovery cleanup complete. Removed {cleaned} expired snapshot file(s).")
1139
+ elif action == "status":
1140
+ st = recovery_mgr.get_storage_status()
1141
+ console.print(f"\n[bold cyan]Recovery Storage & Retention Status:[/bold cyan]")
1142
+ console.print(f" * Total Change Records: [bold white]{st['total_change_records']}[/bold white]")
1143
+ console.print(f" * Storage Used: [bold green]{st['storage_used_mb']} MB[/bold green] ({st['storage_used_bytes']:,} bytes)")
1144
+ console.print(f" * Retention Period: [bold]{st['retention_days']} days[/bold]")
1145
+ console.print(f" * Storage Directory: [dim]{st['directory']}[/dim]\n")
1146
+ else:
1147
+ print_error(f"Unknown recovery action '{action}'. Supported: list, cleanup, status")
1016
1148
 
1017
1149
 
1018
1150
  # ---- define ----
@@ -69,7 +69,18 @@ class BaseConnector(ABC):
69
69
  """
70
70
  pass
71
71
 
72
+ def get_recovery_capabilities(self) -> Dict[str, bool]:
73
+ """Returns map of supported recovery strategies."""
74
+ return {
75
+ "transaction_rollback": True,
76
+ "row_snapshot": True,
77
+ "database_snapshot": False,
78
+ "point_in_time_recovery": False,
79
+ "long_term_recovery": True,
80
+ }
81
+
72
82
  @abstractmethod
73
83
  def close(self) -> None:
74
84
  """Close any open connections."""
75
85
  pass
86
+
@@ -473,6 +473,17 @@ class RelationalConnector(BaseConnector):
473
473
  except Exception as e:
474
474
  return [], [], str(e)
475
475
 
476
+ def get_recovery_capabilities(self) -> Dict[str, bool]:
477
+ is_sqlite = self.engine.dialect.name.lower() == "sqlite"
478
+ return {
479
+ "transaction_rollback": True,
480
+ "row_snapshot": True,
481
+ "database_snapshot": is_sqlite,
482
+ "point_in_time_recovery": False,
483
+ "long_term_recovery": True,
484
+ }
485
+
476
486
  def close(self) -> None:
477
487
  if self._engine:
478
488
  self._engine.dispose()
489
+
@@ -1,15 +1,25 @@
1
1
  """
2
2
  Change Set & Recovery Tracker.
3
- Captures snapshots of rows modified/deleted by write operations and stores them in ~/.dbagent/changesets/
4
- Enables auditable inspection (db-agent changes show <id>) and recovery rollback (db-agent rollback <id>).
3
+ Stores persistent Change Records with rollback snapshots, metadata, and checksums in ~/.dbagent/recovery/
4
+ Provides backward-compatible ChangeSetManager and links to RecoveryManager.
5
5
  """
6
6
 
7
7
  import json
8
8
  import time
9
9
  import uuid
10
+ import hashlib
10
11
  from pathlib import Path
11
12
  from typing import Dict, Any, List, Optional, Tuple
12
13
 
14
+ from dbagent.execution.recovery import (
15
+ RecoveryManager,
16
+ RecoveryStatus,
17
+ ConflictPolicy,
18
+ RollbackPreview,
19
+ RollbackResult,
20
+ RowSnapshotRecoveryStrategy,
21
+ )
22
+
13
23
 
14
24
  class ChangeSet:
15
25
  """Represents a recorded set of database modifications with recovery data."""
@@ -19,11 +29,13 @@ class ChangeSet:
19
29
  database: str,
20
30
  sql: str,
21
31
  table_name: str,
22
- operation: str, # UPDATE, DELETE, INSERT
32
+ operation: str,
23
33
  timestamp: str,
24
34
  affected_rows: int,
25
35
  before_snapshot: List[Dict[str, Any]],
26
- status: str = "COMMITTED",
36
+ status: str = "AVAILABLE",
37
+ rollback_of: Optional[str] = None,
38
+ expires_at: Optional[str] = None,
27
39
  ):
28
40
  self.change_set_id = change_set_id
29
41
  self.database = database
@@ -34,6 +46,8 @@ class ChangeSet:
34
46
  self.affected_rows = affected_rows
35
47
  self.before_snapshot = before_snapshot
36
48
  self.status = status
49
+ self.rollback_of = rollback_of
50
+ self.expires_at = expires_at
37
51
 
38
52
  def to_dict(self) -> Dict[str, Any]:
39
53
  return {
@@ -46,31 +60,38 @@ class ChangeSet:
46
60
  "affected_rows": self.affected_rows,
47
61
  "before_snapshot": self.before_snapshot,
48
62
  "status": self.status,
63
+ "rollback_of": self.rollback_of,
64
+ "expires_at": self.expires_at,
49
65
  }
50
66
 
51
67
  @classmethod
52
68
  def from_dict(cls, data: Dict[str, Any]) -> "ChangeSet":
69
+ cid = data.get("change_id") or data.get("change_set_id")
70
+ ts = data.get("executed_at") or data.get("timestamp") or ""
71
+ st = data.get("recovery_status") or data.get("status") or "AVAILABLE"
53
72
  return cls(
54
- change_set_id=data["change_set_id"],
55
- database=data["database"],
56
- sql=data["sql"],
57
- table_name=data["table_name"],
58
- operation=data["operation"],
59
- timestamp=data["timestamp"],
73
+ change_set_id=cid,
74
+ database=data.get("database", "database"),
75
+ sql=data.get("sql", ""),
76
+ table_name=data.get("table_name", "unknown"),
77
+ operation=data.get("operation", "UPDATE"),
78
+ timestamp=ts,
60
79
  affected_rows=data.get("affected_rows", 0),
61
80
  before_snapshot=data.get("before_snapshot", []),
62
- status=data.get("status", "COMMITTED"),
81
+ status=st,
82
+ rollback_of=data.get("rollback_of"),
83
+ expires_at=data.get("expires_at"),
63
84
  )
64
85
 
65
86
 
66
87
  class ChangeSetManager:
67
88
  """
68
- Manages persistent ChangeSet history on disk in ~/.dbagent/changesets/.
89
+ Manages persistent ChangeSet history on disk in ~/.dbagent/recovery/.
69
90
  """
70
91
 
71
- def __init__(self, storage_dir: Optional[Path] = None):
72
- self.storage_dir = storage_dir or Path.home() / ".dbagent" / "changesets"
73
- self.storage_dir.mkdir(parents=True, exist_ok=True)
92
+ def __init__(self, storage_dir: Optional[Path] = None, retention_days: int = 30):
93
+ self.recovery_mgr = RecoveryManager(storage_dir=storage_dir, retention_days=retention_days)
94
+ self.storage_dir = self.recovery_mgr.storage_dir
74
95
 
75
96
  def create_change_set(
76
97
  self,
@@ -80,58 +101,65 @@ class ChangeSetManager:
80
101
  operation: str,
81
102
  affected_rows: int,
82
103
  before_snapshot: List[Dict[str, Any]],
104
+ schema_checksum: Optional[str] = None,
105
+ rollback_of: Optional[str] = None,
83
106
  ) -> ChangeSet:
84
107
  """Generates a new ChangeSet and persists it to disk."""
85
- date_str = time.strftime("%Y%m%d")
86
- short_id = uuid.uuid4().hex[:6].upper()
87
- cs_id = f"CS-{date_str}-{short_id}"
108
+ snapshot_dict = None
109
+ if before_snapshot:
110
+ snapshot_dict = {
111
+ "table": table_name,
112
+ "operation": operation,
113
+ "pk_columns": ["id"],
114
+ "rows": [{"pk": {"id": r.get("id")}, "before": r} for r in before_snapshot],
115
+ }
116
+
117
+ meta = self.recovery_mgr.create_change_record(
118
+ database=database,
119
+ sql=sql,
120
+ table_name=table_name,
121
+ operation=operation,
122
+ affected_rows=affected_rows,
123
+ snapshot_data=snapshot_dict,
124
+ schema_checksum=schema_checksum,
125
+ rollback_of=rollback_of,
126
+ )
88
127
 
89
- cs = ChangeSet(
90
- change_set_id=cs_id,
128
+ return ChangeSet(
129
+ change_set_id=meta["change_id"],
91
130
  database=database,
92
131
  sql=sql,
93
132
  table_name=table_name,
94
- operation=operation.upper(),
95
- timestamp=time.strftime("%Y-%m-%d %H:%M:%S"),
133
+ operation=operation,
134
+ timestamp=meta["executed_at"],
96
135
  affected_rows=affected_rows,
97
136
  before_snapshot=before_snapshot,
98
- status="COMMITTED",
137
+ status=meta["recovery_status"],
138
+ rollback_of=rollback_of,
139
+ expires_at=meta["expires_at"],
99
140
  )
100
141
 
101
- file_path = self.storage_dir / f"{cs_id}.json"
102
- with open(file_path, "w", encoding="utf-8") as f:
103
- json.dump(cs.to_dict(), f, indent=2)
104
-
105
- return cs
106
-
107
142
  def get_change_set(self, change_set_id: str) -> Optional[ChangeSet]:
108
- """Retrieves a ChangeSet by its ID."""
109
- file_path = self.storage_dir / f"{change_set_id}.json"
110
- if not file_path.exists():
111
- return None
112
- try:
113
- with open(file_path, "r", encoding="utf-8") as f:
114
- data = json.load(f)
115
- return ChangeSet.from_dict(data)
116
- except Exception:
143
+ meta = self.recovery_mgr.get_change_record(change_set_id)
144
+ if not meta:
117
145
  return None
146
+ snap = self.recovery_mgr.get_snapshot(change_set_id)
147
+ before_rows = []
148
+ if snap and "rows" in snap:
149
+ for r in snap["rows"]:
150
+ before_rows.append(r.get("before", {}))
151
+ meta["before_snapshot"] = before_rows
152
+ return ChangeSet.from_dict(meta)
118
153
 
119
154
  def list_change_sets(self, limit: int = 20) -> List[ChangeSet]:
120
- """Lists recent change sets sorted newest first."""
155
+ records = self.recovery_mgr.list_changes(limit=limit)
121
156
  results = []
122
- files = sorted(self.storage_dir.glob("CS-*.json"), reverse=True)
123
- for f in files[:limit]:
124
- try:
125
- with open(f, "r", encoding="utf-8") as fp:
126
- results.append(ChangeSet.from_dict(json.load(fp)))
127
- except Exception:
128
- continue
157
+ for r in records:
158
+ results.append(ChangeSet.from_dict(r))
129
159
  return results
130
160
 
131
161
  def generate_rollback_sql(self, change_set: ChangeSet, primary_key: str = "id") -> List[str]:
132
- """
133
- Generates inverse SQL statements to restore the before_snapshot.
134
- """
162
+ """Generates inverse SQL statements to restore the before_snapshot."""
135
163
  if not change_set.before_snapshot:
136
164
  return []
137
165