dbagent-cli 0.3.0__tar.gz → 0.3.2__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.2}/PKG-INFO +2 -2
  2. {dbagent_cli-0.3.0 → dbagent_cli-0.3.2}/dbagent/__init__.py +1 -1
  3. {dbagent_cli-0.3.0 → dbagent_cli-0.3.2}/dbagent/agent/pipeline.py +87 -12
  4. {dbagent_cli-0.3.0 → dbagent_cli-0.3.2}/dbagent/cli.py +197 -53
  5. {dbagent_cli-0.3.0 → dbagent_cli-0.3.2}/dbagent/connectors/base.py +11 -0
  6. {dbagent_cli-0.3.0 → dbagent_cli-0.3.2}/dbagent/connectors/relational.py +11 -0
  7. {dbagent_cli-0.3.0 → dbagent_cli-0.3.2}/dbagent/execution/changeset.py +77 -49
  8. dbagent_cli-0.3.2/dbagent/execution/recovery.py +571 -0
  9. {dbagent_cli-0.3.0 → dbagent_cli-0.3.2}/dbagent_cli.egg-info/PKG-INFO +2 -2
  10. {dbagent_cli-0.3.0 → dbagent_cli-0.3.2}/dbagent_cli.egg-info/SOURCES.txt +2 -0
  11. {dbagent_cli-0.3.0 → dbagent_cli-0.3.2}/pyproject.toml +2 -2
  12. {dbagent_cli-0.3.0 → dbagent_cli-0.3.2}/tests/test_cli.py +13 -0
  13. dbagent_cli-0.3.2/tests/test_recovery.py +265 -0
  14. {dbagent_cli-0.3.0 → dbagent_cli-0.3.2}/tests/test_sandbox.py +1 -1
  15. {dbagent_cli-0.3.0 → dbagent_cli-0.3.2}/README.md +0 -0
  16. {dbagent_cli-0.3.0 → dbagent_cli-0.3.2}/dbagent/agent/confidence.py +0 -0
  17. {dbagent_cli-0.3.0 → dbagent_cli-0.3.2}/dbagent/agent/conversation.py +0 -0
  18. {dbagent_cli-0.3.0 → dbagent_cli-0.3.2}/dbagent/agent/doctor.py +0 -0
  19. {dbagent_cli-0.3.0 → dbagent_cli-0.3.2}/dbagent/agent/error_classifier.py +0 -0
  20. {dbagent_cli-0.3.0 → dbagent_cli-0.3.2}/dbagent/agent/generator.py +0 -0
  21. {dbagent_cli-0.3.0 → dbagent_cli-0.3.2}/dbagent/agent/recommender.py +0 -0
  22. {dbagent_cli-0.3.0 → dbagent_cli-0.3.2}/dbagent/agent/validator.py +0 -0
  23. {dbagent_cli-0.3.0 → dbagent_cli-0.3.2}/dbagent/config.py +0 -0
  24. {dbagent_cli-0.3.0 → dbagent_cli-0.3.2}/dbagent/connectors/factory.py +0 -0
  25. {dbagent_cli-0.3.0 → dbagent_cli-0.3.2}/dbagent/connectors/mongo.py +0 -0
  26. {dbagent_cli-0.3.0 → dbagent_cli-0.3.2}/dbagent/execution/__init__.py +0 -0
  27. {dbagent_cli-0.3.0 → dbagent_cli-0.3.2}/dbagent/execution/cache.py +0 -0
  28. {dbagent_cli-0.3.0 → dbagent_cli-0.3.2}/dbagent/execution/cost.py +0 -0
  29. {dbagent_cli-0.3.0 → dbagent_cli-0.3.2}/dbagent/execution/sandbox.py +0 -0
  30. {dbagent_cli-0.3.0 → dbagent_cli-0.3.2}/dbagent/execution/streaming.py +0 -0
  31. {dbagent_cli-0.3.0 → dbagent_cli-0.3.2}/dbagent/llm/auto_setup.py +0 -0
  32. {dbagent_cli-0.3.0 → dbagent_cli-0.3.2}/dbagent/llm/base.py +0 -0
  33. {dbagent_cli-0.3.0 → dbagent_cli-0.3.2}/dbagent/llm/factory.py +0 -0
  34. {dbagent_cli-0.3.0 → dbagent_cli-0.3.2}/dbagent/llm/gemini_provider.py +0 -0
  35. {dbagent_cli-0.3.0 → dbagent_cli-0.3.2}/dbagent/llm/groq_provider.py +0 -0
  36. {dbagent_cli-0.3.0 → dbagent_cli-0.3.2}/dbagent/llm/mock_provider.py +0 -0
  37. {dbagent_cli-0.3.0 → dbagent_cli-0.3.2}/dbagent/llm/ollama_provider.py +0 -0
  38. {dbagent_cli-0.3.0 → dbagent_cli-0.3.2}/dbagent/llm/openrouter_provider.py +0 -0
  39. {dbagent_cli-0.3.0 → dbagent_cli-0.3.2}/dbagent/schema/cache.py +0 -0
  40. {dbagent_cli-0.3.0 → dbagent_cli-0.3.2}/dbagent/schema/formatter.py +0 -0
  41. {dbagent_cli-0.3.0 → dbagent_cli-0.3.2}/dbagent/schema/graph.py +0 -0
  42. {dbagent_cli-0.3.0 → dbagent_cli-0.3.2}/dbagent/schema/models.py +0 -0
  43. {dbagent_cli-0.3.0 → dbagent_cli-0.3.2}/dbagent/schema/retriever.py +0 -0
  44. {dbagent_cli-0.3.0 → dbagent_cli-0.3.2}/dbagent/schema/selector.py +0 -0
  45. {dbagent_cli-0.3.0 → dbagent_cli-0.3.2}/dbagent/security/__init__.py +0 -0
  46. {dbagent_cli-0.3.0 → dbagent_cli-0.3.2}/dbagent/security/audit.py +0 -0
  47. {dbagent_cli-0.3.0 → dbagent_cli-0.3.2}/dbagent/security/environment.py +0 -0
  48. {dbagent_cli-0.3.0 → dbagent_cli-0.3.2}/dbagent/security/masking.py +0 -0
  49. {dbagent_cli-0.3.0 → dbagent_cli-0.3.2}/dbagent/security/modes.py +0 -0
  50. {dbagent_cli-0.3.0 → dbagent_cli-0.3.2}/dbagent/security/rbac.py +0 -0
  51. {dbagent_cli-0.3.0 → dbagent_cli-0.3.2}/dbagent/semantic/__init__.py +0 -0
  52. {dbagent_cli-0.3.0 → dbagent_cli-0.3.2}/dbagent/semantic/dictionary.py +0 -0
  53. {dbagent_cli-0.3.0 → dbagent_cli-0.3.2}/dbagent/semantic/templates.py +0 -0
  54. {dbagent_cli-0.3.0 → dbagent_cli-0.3.2}/dbagent/sql/__init__.py +0 -0
  55. {dbagent_cli-0.3.0 → dbagent_cli-0.3.2}/dbagent/sql/optimizer.py +0 -0
  56. {dbagent_cli-0.3.0 → dbagent_cli-0.3.2}/dbagent/sql/validator.py +0 -0
  57. {dbagent_cli-0.3.0 → dbagent_cli-0.3.2}/dbagent/ui/console.py +0 -0
  58. {dbagent_cli-0.3.0 → dbagent_cli-0.3.2}/dbagent/ui/viewer.py +0 -0
  59. {dbagent_cli-0.3.0 → dbagent_cli-0.3.2}/dbagent/ui/visualizer.py +0 -0
  60. {dbagent_cli-0.3.0 → dbagent_cli-0.3.2}/dbagent_cli.egg-info/dependency_links.txt +0 -0
  61. {dbagent_cli-0.3.0 → dbagent_cli-0.3.2}/dbagent_cli.egg-info/entry_points.txt +0 -0
  62. {dbagent_cli-0.3.0 → dbagent_cli-0.3.2}/dbagent_cli.egg-info/requires.txt +0 -0
  63. {dbagent_cli-0.3.0 → dbagent_cli-0.3.2}/dbagent_cli.egg-info/top_level.txt +0 -0
  64. {dbagent_cli-0.3.0 → dbagent_cli-0.3.2}/setup.cfg +0 -0
  65. {dbagent_cli-0.3.0 → dbagent_cli-0.3.2}/tests/test_cache.py +0 -0
  66. {dbagent_cli-0.3.0 → dbagent_cli-0.3.2}/tests/test_connectors.py +0 -0
  67. {dbagent_cli-0.3.0 → dbagent_cli-0.3.2}/tests/test_doctor_and_viz.py +0 -0
  68. {dbagent_cli-0.3.0 → dbagent_cli-0.3.2}/tests/test_generator.py +0 -0
  69. {dbagent_cli-0.3.0 → dbagent_cli-0.3.2}/tests/test_graph.py +0 -0
  70. {dbagent_cli-0.3.0 → dbagent_cli-0.3.2}/tests/test_safety.py +0 -0
  71. {dbagent_cli-0.3.0 → dbagent_cli-0.3.2}/tests/test_schema.py +0 -0
  72. {dbagent_cli-0.3.0 → dbagent_cli-0.3.2}/tests/test_security.py +0 -0
  73. {dbagent_cli-0.3.0 → dbagent_cli-0.3.2}/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.2
4
+ Summary: Enterprise AI Database Copilot with High-Speed Schema Caching & Long-Term 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.2"
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,12 +122,47 @@ 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
121
129
  self.db_name = "database"
122
130
  if hasattr(connector, "engine") and connector.engine:
123
131
  self.db_name = connector.engine.url.database or "database"
132
+ self.db_id = str(connector.engine.url)
133
+ else:
134
+ self.db_id = self.db_name
135
+ self.cache_mgr = SchemaCacheManager()
136
+
137
+ def _get_fast_schema_context(
138
+ self,
139
+ user_prompt: str,
140
+ exact: List[str],
141
+ fuzzy: List[str],
142
+ ) -> Tuple[DatabaseSchema, List[JoinStep]]:
143
+ """Fast schema sub-graph retrieval using cache (<5ms) or targeted introspection (<50ms)."""
144
+ # 1. Try local schema cache (< 5ms)
145
+ cached = self.cache_mgr.load_cached_schema(self.db_id)
146
+ if cached:
147
+ retriever = ContextRetriever(cached)
148
+ return retriever.retrieve_context(
149
+ user_prompt=user_prompt,
150
+ exact_tables=exact,
151
+ fuzzy_tables=fuzzy,
152
+ max_tables=8,
153
+ )
154
+
155
+ # 2. Targeted on-demand table introspection (~50ms instead of 5000ms for full scan)
156
+ target_names = list(set(exact + fuzzy))
157
+ if target_names:
158
+ targeted = self.connector.inspect_schema(table_names=target_names, include_samples=False, max_samples=0)
159
+ retriever = ContextRetriever(targeted)
160
+ return retriever.retrieve_context(user_prompt, exact_tables=exact, fuzzy_tables=fuzzy, max_tables=8)
161
+
162
+ # 3. Fallback: inspect targeted
163
+ targeted = self.connector.inspect_targeted(user_prompt, max_tables=8)
164
+ retriever = ContextRetriever(targeted)
165
+ return retriever.retrieve_context(user_prompt, exact_tables=exact, fuzzy_tables=fuzzy, max_tables=8)
124
166
 
125
167
  def run(self, user_prompt: str, force: bool = False, bypass_cache: bool = False) -> PipelineResult:
126
168
  """
@@ -130,10 +172,16 @@ class QueryPipeline:
130
172
  result = PipelineResult()
131
173
  result.intent = ScriptGenerator.classify_intent(user_prompt)
132
174
 
133
- # --- 1. Conversational Chat Fast-Path ---
175
+ # --- 1. Conversational Chat Fast-Path (<10ms schema load) ---
134
176
  if result.intent == "chat":
135
177
  try:
136
- schema = self.connector.inspect_schema(include_samples=False, max_samples=0, include_row_counts=False)
178
+ cached = self.cache_mgr.load_cached_schema(self.db_id)
179
+ if cached:
180
+ schema = cached
181
+ else:
182
+ quick_tables = self.connector.get_table_names()[:25]
183
+ schema = self.connector.inspect_schema(table_names=quick_tables, include_samples=False, max_samples=0)
184
+
137
185
  result.schema = schema
138
186
  result.message = self.generator.generate_chat_response(
139
187
  schema=schema,
@@ -169,16 +217,9 @@ class QueryPipeline:
169
217
  result.error = "Ambiguous table reference — please specify table name."
170
218
  return result
171
219
 
172
- # --- 3. Fast AI Context Retrieval (Minimal sub-graph) ---
220
+ # --- 3. Fast AI Context Retrieval (Minimal sub-graph < 5ms) ---
173
221
  try:
174
- full_schema = self.connector.inspect_schema(include_samples=False, max_samples=0, include_row_counts=False)
175
- retriever = ContextRetriever(full_schema)
176
- sub_schema, join_steps = retriever.retrieve_context(
177
- user_prompt=user_prompt,
178
- exact_tables=exact,
179
- fuzzy_tables=fuzzy,
180
- max_tables=8,
181
- )
222
+ sub_schema, join_steps = self._get_fast_schema_context(user_prompt, exact, fuzzy)
182
223
  result.schema = sub_schema
183
224
  result.join_steps = join_steps
184
225
  except Exception as e:
@@ -272,10 +313,44 @@ class QueryPipeline:
272
313
  elif self.mode == OperationMode.ADMIN:
273
314
  should_execute = True
274
315
 
275
- # --- 11. Execute with Error Retry ---
316
+ # --- 11. Pre-Write Recovery Snapshot Capture & Execution ---
317
+ snapshot_data = None
318
+ target_table_name = exact[0] if exact else (sub_schema.tables[0].name if sub_schema.tables else "unknown")
319
+ table_model = sub_schema.get_table(target_table_name) if sub_schema else None
320
+
321
+ if should_execute and result.query_type in ("write", "ddl") and table_model:
322
+ can_snap, snap, snap_err = RowSnapshotRecoveryStrategy.capture_pre_write_snapshot(
323
+ connector=self.connector,
324
+ table_model=table_model,
325
+ sql=result.sql,
326
+ operation=result.query_type.upper(),
327
+ )
328
+ if can_snap and snap:
329
+ snapshot_data = snap
330
+
276
331
  if should_execute:
277
332
  self._execute_with_retry(result, augmented_prompt, sub_schema)
278
333
 
334
+ # --- Record ChangeSet / Recovery Record for successful write operations ---
335
+ if result.success and result.query_type in ("write", "ddl"):
336
+ affected = len(snapshot_data.get("rows", [])) if snapshot_data else 1
337
+ meta = self.recovery_mgr.create_change_record(
338
+ database=self.db_name,
339
+ sql=result.sql,
340
+ table_name=target_table_name,
341
+ operation=result.query_type.upper(),
342
+ affected_rows=affected,
343
+ snapshot_data=snapshot_data,
344
+ user=self.user_id,
345
+ )
346
+ result.change_id = meta["change_id"]
347
+ result.recovery_status = meta["recovery_status"]
348
+ result.expires_at = meta["expires_at"]
349
+ if meta["recovery_status"] == RecoveryStatus.AVAILABLE.value:
350
+ result.recovery_message = f"Available until {meta['expires_at']} ({self.recovery_mgr.retention_days} days retention)"
351
+ else:
352
+ result.recovery_message = "Not available (unsupported operation or missing primary key)"
353
+
279
354
  # --- 12. Sensitive Data Masking ---
280
355
  if result.success and self.mask_sensitive_data and result.rows:
281
356
  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")
@@ -338,11 +357,15 @@ def chat_command(
338
357
  console.print(f"[bold cyan]AI:[/bold cyan] {llm.name} | [bold yellow]Mode:[/bold yellow] {current_mode.value}")
339
358
  console.print("[dim]Type your question in plain English or use colon commands (e.g. :help, :tables, :doctor).[/dim]")
340
359
 
341
- # Display AI recommendations on connect
360
+ # Display AI recommendations on connect (instant via cache or top tables)
342
361
  try:
343
- with console.status("[dim]Analyzing schema recommendations...[/dim]"):
344
- full_schema = connector.inspect_schema(include_samples=False, max_samples=0, include_row_counts=False)
345
- recs = QueryRecommender.generate_recommendations(full_schema, max_suggestions=3)
362
+ cached_schema = cache_mgr.load_cached_schema(db_url)
363
+ if cached_schema:
364
+ recs = QueryRecommender.generate_recommendations(cached_schema, max_suggestions=3)
365
+ else:
366
+ quick_tables = table_names[:10]
367
+ quick_schema = connector.inspect_schema(table_names=quick_tables, include_samples=False, max_samples=0)
368
+ recs = QueryRecommender.generate_recommendations(quick_schema, max_suggestions=3)
346
369
  if recs:
347
370
  console.print("\n[bold cyan]💡 Try asking:[/bold cyan]")
348
371
  for r in recs:
@@ -424,8 +447,12 @@ def chat_command(
424
447
  continue
425
448
  exact, fuzzy = connector.resolve_tables(cmd_arg)
426
449
  target = exact[0] if exact else (fuzzy[0] if fuzzy else cmd_arg)
427
- with console.status(f"[bold cyan]Finding relationships for {target}...[/bold cyan]"):
428
- full_schema = connector.inspect_schema(include_samples=False, max_samples=0)
450
+ cached = cache_mgr.load_cached_schema(db_url)
451
+ if cached:
452
+ full_schema = cached
453
+ else:
454
+ with console.status(f"[bold cyan]Finding relationships for {target}...[/bold cyan]"):
455
+ full_schema = connector.inspect_schema(include_samples=False, max_samples=0)
429
456
 
430
457
  tmodel = full_schema.get_table(target)
431
458
  outgoing = tmodel.foreign_keys if tmodel else []
@@ -454,10 +481,14 @@ def chat_command(
454
481
  if len(args) < 2:
455
482
  print_warning("Usage: :path <start_table> <end_table> (e.g. :path users events)")
456
483
  continue
457
- with console.status("[bold cyan]Traversing relationship graph...[/bold cyan]"):
458
- full_schema = connector.inspect_schema(include_samples=False, max_samples=0)
459
- graph = RelationshipGraph(full_schema)
460
- path = graph.find_shortest_path(args[0], args[1])
484
+ cached = cache_mgr.load_cached_schema(db_url)
485
+ if cached:
486
+ full_schema = cached
487
+ else:
488
+ with console.status("[bold cyan]Traversing relationship graph...[/bold cyan]"):
489
+ full_schema = connector.inspect_schema(include_samples=False, max_samples=0)
490
+ graph = RelationshipGraph(full_schema)
491
+ path = graph.find_shortest_path(args[0], args[1])
461
492
  if path:
462
493
  console.print(f"\n[bold cyan]Shortest Relationship Path:[/bold cyan]\n{path.format_visual_path()}\n")
463
494
  else:
@@ -931,22 +962,34 @@ def schema_status_command(
931
962
  def changes_command(
932
963
  limit: int = typer.Option(20, "--limit", "-l", help="Number of change sets to show"),
933
964
  ):
934
- """List recorded write change sets."""
965
+ """List recorded write change sets and recovery availability."""
935
966
  print_banner()
936
- sets = change_mgr.list_change_sets(limit=limit)
937
- if not sets:
967
+ records = recovery_mgr.list_changes(limit=limit)
968
+ if not records:
938
969
  print_info("No recorded write change sets found.")
939
970
  return
940
971
 
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")
972
+ table = Table(title="[bold cyan]Recorded Write ChangeSets & Recovery Status[/bold cyan]", show_header=True, header_style="bold cyan")
973
+ table.add_column("Change ID", style="bold white")
943
974
  table.add_column("Timestamp", style="dim")
944
975
  table.add_column("Operation", style="yellow")
945
976
  table.add_column("Table", style="cyan")
946
- table.add_column("Affected Rows", justify="right")
977
+ table.add_column("Rows", justify="right")
978
+ table.add_column("Recovery Status", style="bold")
979
+ table.add_column("Expires", style="dim")
947
980
 
948
- for s in sets:
949
- table.add_row(s.change_set_id, s.timestamp, s.operation, s.table_name, str(s.affected_rows))
981
+ for s in records:
982
+ st = s.get("recovery_status", "AVAILABLE")
983
+ st_style = "[green]AVAILABLE[/green]" if st == "AVAILABLE" else ("[yellow]EXPIRED[/yellow]" if st == "EXPIRED" else f"[dim]{st}[/dim]")
984
+ table.add_row(
985
+ s.get("change_id", ""),
986
+ s.get("executed_at", ""),
987
+ s.get("operation", ""),
988
+ s.get("table_name", ""),
989
+ str(s.get("affected_rows", 0)),
990
+ st_style,
991
+ s.get("expires_at", "-"),
992
+ )
950
993
 
951
994
  console.print(table)
952
995
 
@@ -955,64 +998,165 @@ def changes_command(
955
998
 
956
999
  @app.command(name="changes-show")
957
1000
  def changes_show_command(
958
- change_set_id: str = typer.Argument(..., help="ChangeSet ID (e.g. CS-20260828-001)"),
1001
+ change_id: str = typer.Argument(..., help="Change ID (e.g. CHG-20260828-001)"),
959
1002
  ):
960
1003
  """Show details and before-snapshot of a recorded ChangeSet."""
961
1004
  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.")
1005
+ meta = recovery_mgr.get_change_record(change_id)
1006
+ if not meta:
1007
+ print_error(f"Change record '{change_id}' not found.")
965
1008
  raise typer.Exit(1)
966
1009
 
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")
1010
+ console.print(f"\n[bold cyan]Change Record: {meta['change_id']}[/bold cyan]")
1011
+ console.print(f" * Database: {meta.get('database')} | Executed: {meta.get('executed_at')}")
1012
+ console.print(f" * Operation: [{meta.get('operation')}] on Table `{meta.get('table_name')}` ({meta.get('affected_rows')} rows)")
1013
+ console.print(f" * Recovery Status: [bold green]{meta.get('recovery_status')}[/bold green] (Expires: {meta.get('expires_at')})")
1014
+ if meta.get("rollback_of"):
1015
+ console.print(f" * Rollback Of: [bold yellow]{meta['rollback_of']}[/bold yellow]")
1016
+ print_code(meta.get("sql", ""), "sql", "Executed Statement")
971
1017
 
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")
1018
+ snap = recovery_mgr.get_snapshot(change_id)
1019
+ if snap and "rows" in snap and snap["rows"]:
1020
+ rows_disp = [r.get("before", {}) for r in snap["rows"][:20]]
1021
+ if rows_disp and rows_disp[0]:
1022
+ console.print(f"\n[bold yellow]Captured Before-Snapshot ({len(snap['rows'])} rows captured):[/bold yellow]")
1023
+ cols = list(rows_disp[0].keys())
1024
+ print_results_table(cols, rows_disp, title="Snapshot Sample (First 20)")
976
1025
 
977
1026
 
978
1027
  # ---- rollback ----
979
1028
 
980
1029
  @app.command(name="rollback")
981
1030
  def rollback_command(
982
- change_set_id: str = typer.Argument(..., help="ChangeSet ID to roll back"),
1031
+ change_id: Optional[str] = typer.Argument(None, help="Change ID to roll back (e.g. CHG-20260828-001) or use --last"),
983
1032
  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"),
1033
+ preview: bool = typer.Option(False, "--preview", "-p", help="Preview rollback and detect conflicts without executing"),
1034
+ force: bool = typer.Option(False, "--force", "-f", help="Bypass confirmation and overwrite conflicts"),
1035
+ policy: str = typer.Option("skip", "--conflict-policy", help="Conflict resolution policy: skip, abort, force"),
1036
+ last: bool = typer.Option(False, "--last", "-l", help="Roll back the most recent committed change"),
1037
+ mode: str = typer.Option("CONFIRM", "--mode", help="Operation mode (SAFE, READ_ONLY, CONFIRM, ADMIN)"),
985
1038
  ):
986
- """Generate and execute recovery rollback SQL for a ChangeSet."""
1039
+ """Safely revert and recover committed database modifications with conflict protection."""
987
1040
  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.")
1041
+
1042
+ target_id = change_id
1043
+ if last or target_id == "--last":
1044
+ records = recovery_mgr.list_changes(limit=1)
1045
+ if not records:
1046
+ print_error("No recorded change sets available to roll back.")
1047
+ raise typer.Exit(1)
1048
+ target_id = records[0]["change_id"]
1049
+
1050
+ if not target_id:
1051
+ records = recovery_mgr.list_changes(limit=10)
1052
+ if not records:
1053
+ print_info("No recorded change sets found.")
1054
+ raise typer.Exit(0)
1055
+ console.print("[bold cyan]Recent Change Sets:[/bold cyan]")
1056
+ for r in records:
1057
+ 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)")
1058
+ target_id = Prompt.ask("\nEnter Change ID to roll back")
1059
+
1060
+ meta = recovery_mgr.get_change_record(target_id)
1061
+ if not meta:
1062
+ print_error(f"Change record '{target_id}' not found.")
1063
+ raise typer.Exit(1)
1064
+
1065
+ if meta.get("recovery_status") == RecoveryStatus.EXPIRED.value:
1066
+ print_error(f"Rollback rejected: Recovery snapshot for {target_id} has EXPIRED ({meta.get('expires_at')}).")
1067
+ raise typer.Exit(1)
1068
+
1069
+ if meta.get("recovery_status") == RecoveryStatus.UNAVAILABLE.value:
1070
+ print_error(f"Rollback unavailable: Operation was executed without long-term recovery snapshot.")
1071
+ raise typer.Exit(1)
1072
+
1073
+ snap = recovery_mgr.get_snapshot(target_id)
1074
+ if not snap:
1075
+ print_error(f"Recovery snapshot corrupted or missing checksum verification for '{target_id}'.")
1076
+ raise typer.Exit(1)
1077
+
1078
+ db_url = resolve_db_url(db or meta.get("database"))
1079
+ connector = create_connector(db_url)
1080
+
1081
+ # Concurrency & Conflict Preview
1082
+ with console.status("[bold cyan]Evaluating row eligibility & detecting conflicts...[/bold cyan]"):
1083
+ prev = RowSnapshotRecoveryStrategy.preview_recovery(connector, meta, snap)
1084
+
1085
+ console.print(f"\n{prev.format_summary()}\n")
1086
+
1087
+ if not prev.can_proceed:
1088
+ connector.close()
991
1089
  raise typer.Exit(1)
992
1090
 
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}'.")
1091
+ # If preview only, stop here
1092
+ if preview:
1093
+ console.print("[dim]Preview complete. Use 'db-agent rollback " + target_id + "' to execute recovery.[/dim]\n")
1094
+ connector.close()
996
1095
  return
997
1096
 
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")
1097
+ # Check Safe Mode
1098
+ if mode.upper() in ("SAFE", "READ_ONLY"):
1099
+ print_error(f"Rollback execution blocked in {mode.upper()} mode. Rollback requires write permission. Run with --preview to inspect.")
1100
+ connector.close()
1101
+ raise typer.Exit(1)
1102
+
1103
+ # Conflict Policy enforcement
1104
+ pol_enum = ConflictPolicy(policy.upper()) if policy.upper() in ConflictPolicy.__members__ else ConflictPolicy.SKIP
1105
+ if force:
1106
+ pol_enum = ConflictPolicy.FORCE
1001
1107
 
1002
- if not force and not Confirm.ask("Execute these rollback statements on the database?"):
1108
+ if prev.conflicts_detected > 0 and pol_enum == ConflictPolicy.ABORT:
1109
+ print_error(f"Rollback aborted: {prev.conflicts_detected} concurrency conflict(s) detected with policy=ABORT.")
1110
+ connector.close()
1111
+ raise typer.Exit(1)
1112
+
1113
+ if not force and not Confirm.ask(f"Execute rollback for {target_id} ({prev.rows_eligible} rows eligible)?", default=False):
1003
1114
  console.print("[dim]Rollback cancelled.[/dim]")
1115
+ connector.close()
1004
1116
  return
1005
1117
 
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)}")
1118
+ with console.status("[bold cyan]Executing rollback recovery...[/bold cyan]"):
1119
+ res = RowSnapshotRecoveryStrategy.execute_recovery(connector, meta, snap, policy=pol_enum)
1120
+ connector.close()
1121
+
1122
+ # Record rollback audit chain
1123
+ recovery_mgr.create_change_record(
1124
+ database=meta.get("database", "database"),
1125
+ sql=f"-- Rollback of {target_id}",
1126
+ table_name=meta.get("table_name", "unknown"),
1127
+ operation="ROLLBACK",
1128
+ affected_rows=res.restored_count,
1129
+ snapshot_data=None,
1130
+ rollback_of=target_id,
1131
+ )
1132
+
1133
+ console.print(f"\n{res.format_summary()}\n")
1134
+
1135
+
1136
+ # ---- recovery ----
1137
+
1138
+ @app.command(name="recovery")
1139
+ def recovery_command(
1140
+ action: str = typer.Argument("list", help="Action: list, cleanup, status"),
1141
+ limit: int = typer.Option(20, "--limit", "-l", help="Number of records to list"),
1142
+ ):
1143
+ """Manage database recovery snapshots, retention, and storage."""
1144
+ print_banner()
1145
+ if action == "list":
1146
+ changes_command(limit=limit)
1147
+ elif action == "cleanup":
1148
+ with console.status("[bold cyan]Cleaning up expired recovery snapshots...[/bold cyan]"):
1149
+ cleaned = recovery_mgr.cleanup_expired()
1150
+ print_success(f"Recovery cleanup complete. Removed {cleaned} expired snapshot file(s).")
1151
+ elif action == "status":
1152
+ st = recovery_mgr.get_storage_status()
1153
+ console.print(f"\n[bold cyan]Recovery Storage & Retention Status:[/bold cyan]")
1154
+ console.print(f" * Total Change Records: [bold white]{st['total_change_records']}[/bold white]")
1155
+ console.print(f" * Storage Used: [bold green]{st['storage_used_mb']} MB[/bold green] ({st['storage_used_bytes']:,} bytes)")
1156
+ console.print(f" * Retention Period: [bold]{st['retention_days']} days[/bold]")
1157
+ console.print(f" * Storage Directory: [dim]{st['directory']}[/dim]\n")
1158
+ else:
1159
+ print_error(f"Unknown recovery action '{action}'. Supported: list, cleanup, status")
1016
1160
 
1017
1161
 
1018
1162
  # ---- 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
+