This is page 24 of 45. Use http://codebase.md/dicklesworthstone/llm_gateway_mcp_server?lines=true&page={x} to view the full context.
# Directory Structure
```
├── .cursorignore
├── .env.example
├── .envrc
├── .gitignore
├── additional_features.md
├── check_api_keys.py
├── completion_support.py
├── comprehensive_test.py
├── docker-compose.yml
├── Dockerfile
├── empirically_measured_model_speeds.json
├── error_handling.py
├── example_structured_tool.py
├── examples
│ ├── __init__.py
│ ├── advanced_agent_flows_using_unified_memory_system_demo.py
│ ├── advanced_extraction_demo.py
│ ├── advanced_unified_memory_system_demo.py
│ ├── advanced_vector_search_demo.py
│ ├── analytics_reporting_demo.py
│ ├── audio_transcription_demo.py
│ ├── basic_completion_demo.py
│ ├── cache_demo.py
│ ├── claude_integration_demo.py
│ ├── compare_synthesize_demo.py
│ ├── cost_optimization.py
│ ├── data
│ │ ├── sample_event.txt
│ │ ├── Steve_Jobs_Introducing_The_iPhone_compressed.md
│ │ └── Steve_Jobs_Introducing_The_iPhone_compressed.mp3
│ ├── docstring_refiner_demo.py
│ ├── document_conversion_and_processing_demo.py
│ ├── entity_relation_graph_demo.py
│ ├── filesystem_operations_demo.py
│ ├── grok_integration_demo.py
│ ├── local_text_tools_demo.py
│ ├── marqo_fused_search_demo.py
│ ├── measure_model_speeds.py
│ ├── meta_api_demo.py
│ ├── multi_provider_demo.py
│ ├── ollama_integration_demo.py
│ ├── prompt_templates_demo.py
│ ├── python_sandbox_demo.py
│ ├── rag_example.py
│ ├── research_workflow_demo.py
│ ├── sample
│ │ ├── article.txt
│ │ ├── backprop_paper.pdf
│ │ ├── buffett.pdf
│ │ ├── contract_link.txt
│ │ ├── legal_contract.txt
│ │ ├── medical_case.txt
│ │ ├── northwind.db
│ │ ├── research_paper.txt
│ │ ├── sample_data.json
│ │ └── text_classification_samples
│ │ ├── email_classification.txt
│ │ ├── news_samples.txt
│ │ ├── product_reviews.txt
│ │ └── support_tickets.txt
│ ├── sample_docs
│ │ └── downloaded
│ │ └── attention_is_all_you_need.pdf
│ ├── sentiment_analysis_demo.py
│ ├── simple_completion_demo.py
│ ├── single_shot_synthesis_demo.py
│ ├── smart_browser_demo.py
│ ├── sql_database_demo.py
│ ├── sse_client_demo.py
│ ├── test_code_extraction.py
│ ├── test_content_detection.py
│ ├── test_ollama.py
│ ├── text_classification_demo.py
│ ├── text_redline_demo.py
│ ├── tool_composition_examples.py
│ ├── tournament_code_demo.py
│ ├── tournament_text_demo.py
│ ├── unified_memory_system_demo.py
│ ├── vector_search_demo.py
│ ├── web_automation_instruction_packs.py
│ └── workflow_delegation_demo.py
├── LICENSE
├── list_models.py
├── marqo_index_config.json.example
├── mcp_protocol_schema_2025-03-25_version.json
├── mcp_python_lib_docs.md
├── mcp_tool_context_estimator.py
├── model_preferences.py
├── pyproject.toml
├── quick_test.py
├── README.md
├── resource_annotations.py
├── run_all_demo_scripts_and_check_for_errors.py
├── storage
│ └── smart_browser_internal
│ ├── locator_cache.db
│ ├── readability.js
│ └── storage_state.enc
├── test_client.py
├── test_connection.py
├── TEST_README.md
├── test_sse_client.py
├── test_stdio_client.py
├── tests
│ ├── __init__.py
│ ├── conftest.py
│ ├── integration
│ │ ├── __init__.py
│ │ └── test_server.py
│ ├── manual
│ │ ├── test_extraction_advanced.py
│ │ └── test_extraction.py
│ └── unit
│ ├── __init__.py
│ ├── test_cache.py
│ ├── test_providers.py
│ └── test_tools.py
├── TODO.md
├── tool_annotations.py
├── tools_list.json
├── ultimate_mcp_banner.webp
├── ultimate_mcp_logo.webp
├── ultimate_mcp_server
│ ├── __init__.py
│ ├── __main__.py
│ ├── cli
│ │ ├── __init__.py
│ │ ├── __main__.py
│ │ ├── commands.py
│ │ ├── helpers.py
│ │ └── typer_cli.py
│ ├── clients
│ │ ├── __init__.py
│ │ ├── completion_client.py
│ │ └── rag_client.py
│ ├── config
│ │ └── examples
│ │ └── filesystem_config.yaml
│ ├── config.py
│ ├── constants.py
│ ├── core
│ │ ├── __init__.py
│ │ ├── evaluation
│ │ │ ├── base.py
│ │ │ └── evaluators.py
│ │ ├── providers
│ │ │ ├── __init__.py
│ │ │ ├── anthropic.py
│ │ │ ├── base.py
│ │ │ ├── deepseek.py
│ │ │ ├── gemini.py
│ │ │ ├── grok.py
│ │ │ ├── ollama.py
│ │ │ ├── openai.py
│ │ │ └── openrouter.py
│ │ ├── server.py
│ │ ├── state_store.py
│ │ ├── tournaments
│ │ │ ├── manager.py
│ │ │ ├── tasks.py
│ │ │ └── utils.py
│ │ └── ums_api
│ │ ├── __init__.py
│ │ ├── ums_database.py
│ │ ├── ums_endpoints.py
│ │ ├── ums_models.py
│ │ └── ums_services.py
│ ├── exceptions.py
│ ├── graceful_shutdown.py
│ ├── services
│ │ ├── __init__.py
│ │ ├── analytics
│ │ │ ├── __init__.py
│ │ │ ├── metrics.py
│ │ │ └── reporting.py
│ │ ├── cache
│ │ │ ├── __init__.py
│ │ │ ├── cache_service.py
│ │ │ ├── persistence.py
│ │ │ ├── strategies.py
│ │ │ └── utils.py
│ │ ├── cache.py
│ │ ├── document.py
│ │ ├── knowledge_base
│ │ │ ├── __init__.py
│ │ │ ├── feedback.py
│ │ │ ├── manager.py
│ │ │ ├── rag_engine.py
│ │ │ ├── retriever.py
│ │ │ └── utils.py
│ │ ├── prompts
│ │ │ ├── __init__.py
│ │ │ ├── repository.py
│ │ │ └── templates.py
│ │ ├── prompts.py
│ │ └── vector
│ │ ├── __init__.py
│ │ ├── embeddings.py
│ │ └── vector_service.py
│ ├── tool_token_counter.py
│ ├── tools
│ │ ├── __init__.py
│ │ ├── audio_transcription.py
│ │ ├── base.py
│ │ ├── completion.py
│ │ ├── docstring_refiner.py
│ │ ├── document_conversion_and_processing.py
│ │ ├── enhanced-ums-lookbook.html
│ │ ├── entity_relation_graph.py
│ │ ├── excel_spreadsheet_automation.py
│ │ ├── extraction.py
│ │ ├── filesystem.py
│ │ ├── html_to_markdown.py
│ │ ├── local_text_tools.py
│ │ ├── marqo_fused_search.py
│ │ ├── meta_api_tool.py
│ │ ├── ocr_tools.py
│ │ ├── optimization.py
│ │ ├── provider.py
│ │ ├── pyodide_boot_template.html
│ │ ├── python_sandbox.py
│ │ ├── rag.py
│ │ ├── redline-compiled.css
│ │ ├── sentiment_analysis.py
│ │ ├── single_shot_synthesis.py
│ │ ├── smart_browser.py
│ │ ├── sql_databases.py
│ │ ├── text_classification.py
│ │ ├── text_redline_tools.py
│ │ ├── tournament.py
│ │ ├── ums_explorer.html
│ │ └── unified_memory_system.py
│ ├── utils
│ │ ├── __init__.py
│ │ ├── async_utils.py
│ │ ├── display.py
│ │ ├── logging
│ │ │ ├── __init__.py
│ │ │ ├── console.py
│ │ │ ├── emojis.py
│ │ │ ├── formatter.py
│ │ │ ├── logger.py
│ │ │ ├── panels.py
│ │ │ ├── progress.py
│ │ │ └── themes.py
│ │ ├── parse_yaml.py
│ │ ├── parsing.py
│ │ ├── security.py
│ │ └── text.py
│ └── working_memory_api.py
├── unified_memory_system_technical_analysis.md
└── uv.lock
```
# Files
--------------------------------------------------------------------------------
/examples/sql_database_demo.py:
--------------------------------------------------------------------------------
```python
1 | #!/usr/bin/env python
2 | """Demonstration script for SQLTool in Ultimate MCP Server."""
3 |
4 | import asyncio
5 | import datetime as dt
6 | import os
7 | import sqlite3
8 | import sys
9 | import tempfile
10 | from pathlib import Path
11 | from typing import Any, Dict, Optional
12 |
13 | # Add project root to path for imports when running as script
14 | sys.path.insert(0, str(Path(__file__).parent.parent))
15 |
16 | # Rich imports for nice UI
17 | import pandas as pd
18 | import pandera as pa
19 | from rich import box
20 | from rich.console import Console
21 | from rich.markup import escape
22 | from rich.panel import Panel
23 | from rich.progress import BarColumn, Progress, TextColumn
24 | from rich.rule import Rule
25 | from rich.syntax import Syntax
26 | from rich.table import Table
27 | from rich.traceback import install as install_rich_traceback
28 | from rich.tree import Tree
29 |
30 | from ultimate_mcp_server.core.server import Gateway # Import the actual Gateway
31 | from ultimate_mcp_server.exceptions import ToolError, ToolInputError
32 |
33 | # Import the SQLTool class from our module
34 | from ultimate_mcp_server.tools.sql_databases import SQLTool
35 | from ultimate_mcp_server.utils import get_logger
36 |
37 | # Initialize Rich console and logger
38 | console = Console()
39 | logger = get_logger("demo.sql_tool")
40 |
41 | # Install rich tracebacks for better error display
42 | install_rich_traceback(show_locals=False, width=console.width)
43 |
44 | # --- Configuration ---
45 | DEFAULT_CONNECTION_STRING = "sqlite:///:memory:" # In-memory SQLite for demo
46 | # You can replace with a connection string like:
47 | # "postgresql://username:password@localhost:5432/demo_db"
48 | # "mysql+pymysql://username:password@localhost:3306/demo_db"
49 | # "mssql+pyodbc://username:password@localhost:1433/demo_db?driver=ODBC+Driver+17+for+SQL+Server"
50 |
51 | # --- Demo Helper Functions ---
52 |
53 | def display_result(title: str, result: Dict[str, Any], query_str: Optional[str] = None) -> None:
54 | """Display query result with enhanced formatting."""
55 | console.print(Rule(f"[bold cyan]{escape(title)}[/bold cyan]"))
56 |
57 | if query_str:
58 | console.print(Panel(
59 | Syntax(query_str.strip(), "sql", theme="default", line_numbers=False, word_wrap=True),
60 | title="Executed Query",
61 | border_style="blue",
62 | padding=(1, 2)
63 | ))
64 |
65 | if not result.get("success", False):
66 | error_msg = result.get("error", "Unknown error")
67 | console.print(Panel(
68 | f"[bold red]:x: Operation Failed:[/]\n{escape(error_msg)}",
69 | title="Error",
70 | border_style="red",
71 | padding=(1, 2),
72 | expand=False
73 | ))
74 | return
75 |
76 | # Handle different result types based on content
77 | if "rows" in result:
78 | # Query result with rows
79 | rows = result.get("rows", [])
80 | columns = result.get("columns", [])
81 | row_count = result.get("row_count", len(rows))
82 |
83 | if not rows:
84 | console.print(Panel("[yellow]No results returned for this operation.", padding=(0, 1), border_style="yellow"))
85 | return
86 |
87 | table_title = f"Results ({row_count} row{'s' if row_count != 1 else ''} returned)"
88 | if "pagination" in result:
89 | pagination = result["pagination"]
90 | table_title += f" - Page {pagination.get('page', '?')}"
91 |
92 | table = Table(title=table_title, box=box.ROUNDED, show_header=True, padding=(0, 1), border_style="bright_blue")
93 |
94 | # Add columns
95 | for name in columns:
96 | justify = "right" if any(k in name.lower() for k in ['id', 'count', 'price', 'amount', 'quantity', 'total']) else "left"
97 | style = "cyan" if justify == "left" else "magenta"
98 | table.add_column(name, style=style, justify=justify, header_style=f"bold {style}")
99 |
100 | # Add data rows
101 | for row in rows:
102 | table.add_row(*[escape(str(row.get(col_name, ''))) for col_name in columns])
103 |
104 | console.print(table)
105 |
106 | # Display pagination info if available
107 | if "pagination" in result:
108 | pagination = result["pagination"]
109 | pagination_info = Table(title="Pagination Info", show_header=False, box=box.SIMPLE, padding=(0, 1))
110 | pagination_info.add_column("Metric", style="cyan", justify="right")
111 | pagination_info.add_column("Value", style="white")
112 | pagination_info.add_row("Page", str(pagination.get("page")))
113 | pagination_info.add_row("Page Size", str(pagination.get("page_size")))
114 | pagination_info.add_row("Has Next", "[green]:heavy_check_mark:[/]" if pagination.get("has_next_page") else "[dim]:x:[/]")
115 | pagination_info.add_row("Has Previous", "[green]:heavy_check_mark:[/]" if pagination.get("has_previous_page") else "[dim]:x:[/]")
116 | console.print(pagination_info)
117 |
118 | # Show if truncated
119 | if result.get("truncated"):
120 | console.print("[yellow]⚠ Results truncated (reached max_rows limit)[/yellow]")
121 |
122 | elif "documentation" in result:
123 | # Documentation result
124 | doc_content = result.get("documentation", "")
125 | format_type = result.get("format", "markdown")
126 |
127 | console.print(Panel(
128 | Syntax(doc_content, format_type, theme="default", line_numbers=False, word_wrap=True),
129 | title=f"Documentation ({format_type.upper()})",
130 | border_style="magenta",
131 | padding=(1, 2)
132 | ))
133 |
134 | else:
135 | # Generic success result, display as is
136 | console.print(Panel(
137 | "\n".join([f"[cyan]{k}:[/] {escape(str(v))}" for k, v in result.items() if k != "success"]),
138 | title="Operation Result",
139 | border_style="green",
140 | padding=(1, 2)
141 | ))
142 |
143 | console.print() # Add spacing
144 |
145 | # Add setup functionality directly to avoid import issues
146 | def init_demo_database(db_path):
147 | """Set up a demo database with sample tables and data."""
148 | logger.info(f"Setting up demo database at: {db_path}")
149 |
150 | # Connect to SQLite database
151 | conn = sqlite3.connect(db_path)
152 | cursor = conn.cursor()
153 |
154 | # Create tables
155 | setup_queries = [
156 | """
157 | CREATE TABLE IF NOT EXISTS customers (
158 | customer_id INTEGER PRIMARY KEY,
159 | name TEXT NOT NULL,
160 | email TEXT UNIQUE,
161 | signup_date TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
162 | status TEXT CHECK(status IN ('active', 'inactive', 'pending')) DEFAULT 'pending',
163 | ssn TEXT,
164 | credit_card TEXT
165 | )
166 | """,
167 | """
168 | CREATE TABLE IF NOT EXISTS products (
169 | product_id INTEGER PRIMARY KEY,
170 | name TEXT NOT NULL,
171 | description TEXT,
172 | price DECIMAL(10,2) NOT NULL,
173 | category TEXT,
174 | in_stock BOOLEAN DEFAULT 1
175 | )
176 | """,
177 | """
178 | CREATE TABLE IF NOT EXISTS orders (
179 | order_id INTEGER PRIMARY KEY,
180 | customer_id INTEGER NOT NULL,
181 | order_date TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
182 | total_amount DECIMAL(10,2) NOT NULL,
183 | status TEXT DEFAULT 'pending',
184 | FOREIGN KEY (customer_id) REFERENCES customers(customer_id)
185 | )
186 | """,
187 | """
188 | CREATE TABLE IF NOT EXISTS order_items (
189 | item_id INTEGER PRIMARY KEY,
190 | order_id INTEGER NOT NULL,
191 | product_id INTEGER NOT NULL,
192 | quantity INTEGER NOT NULL,
193 | price_per_unit DECIMAL(10,2) NOT NULL,
194 | FOREIGN KEY (order_id) REFERENCES orders(order_id),
195 | FOREIGN KEY (product_id) REFERENCES products(product_id)
196 | )
197 | """
198 | ]
199 |
200 | # Insert sample data
201 | sample_data_queries = [
202 | # Insert customers with PII data already included
203 | """
204 | INSERT INTO customers (customer_id, name, email, status, ssn, credit_card) VALUES
205 | (1, 'Alice Johnson', '[email protected]', 'active', '123-45-6789', '4111-1111-1111-1111'),
206 | (2, 'Bob Smith', '[email protected]', 'active', '234-56-7890', '4222-2222-2222-2222'),
207 | (3, 'Charlie Davis', '[email protected]', 'inactive', '345-67-8901', '4333-3333-3333-3333'),
208 | (4, 'Diana Miller', '[email protected]', 'active', '456-78-9012', '4444-4444-4444-4444'),
209 | (5, 'Ethan Garcia', '[email protected]', 'pending', '567-89-0123', '4555-5555-5555-5555')
210 | """,
211 | # Insert products
212 | """
213 | INSERT INTO products (product_id, name, description, price, category, in_stock) VALUES
214 | (1, 'Laptop Pro X', 'High-performance laptop with 16GB RAM', 1499.99, 'Electronics', 1),
215 | (2, 'Smartphone Z', 'Latest flagship smartphone', 999.99, 'Electronics', 1),
216 | (3, 'Wireless Earbuds', 'Noise-cancelling earbuds', 179.99, 'Audio', 1),
217 | (4, 'Smart Coffee Maker', 'WiFi-enabled coffee machine', 119.99, 'Kitchen', 0),
218 | (5, 'Fitness Tracker', 'Waterproof fitness band with GPS', 79.99, 'Wearables', 1)
219 | """,
220 | # Insert orders
221 | """
222 | INSERT INTO orders (order_id, customer_id, total_amount, status) VALUES
223 | (1, 1, 1499.98, 'completed'),
224 | (2, 2, 89.99, 'processing'),
225 | (3, 1, 249.99, 'completed'),
226 | (4, 3, 1099.98, 'completed'),
227 | (5, 4, 49.99, 'processing')
228 | """,
229 | # Insert order items
230 | """
231 | INSERT INTO order_items (item_id, order_id, product_id, quantity, price_per_unit) VALUES
232 | (1, 1, 1, 1, 1499.99),
233 | (2, 2, 5, 1, 79.99),
234 | (3, 3, 3, 1, 179.99),
235 | (4, 3, 4, 1, 119.99),
236 | (5, 4, 2, 1, 999.99),
237 | (6, 4, 5, 1, 79.99),
238 | (7, 5, 4, 1, 119.99)
239 | """
240 | ]
241 |
242 | try:
243 | # Execute each query to set up schema
244 | for query in setup_queries:
245 | cursor.execute(query)
246 | logger.info(f"Created table: {query.strip().split()[2]}")
247 |
248 | # Execute each query to insert data
249 | for query in sample_data_queries:
250 | cursor.execute(query)
251 | table_name = query.strip().split()[2]
252 | row_count = cursor.rowcount
253 | logger.info(f"Inserted {row_count} rows into {table_name}")
254 |
255 | # Commit the changes
256 | conn.commit()
257 | logger.info("Database setup complete")
258 |
259 | except sqlite3.Error as e:
260 | logger.error(f"SQLite error: {e}")
261 | conn.rollback()
262 | raise
263 | finally:
264 | conn.close()
265 |
266 | return db_path
267 |
268 |
269 | # --- Demo Functions ---
270 |
271 | async def connection_demo(sql_tool: SQLTool, conn_string: Optional[str] = None) -> Optional[str]:
272 | """Demonstrate database connection and status checking."""
273 | console.print(Rule("[bold green]1. Database Connection Demo[/bold green]", style="green"))
274 | logger.info("Starting database connection demo")
275 |
276 | connection_id = None
277 | connection_string = conn_string or DEFAULT_CONNECTION_STRING
278 |
279 | with console.status("[bold cyan]Connecting to database...", spinner="earth"):
280 | try:
281 | # Connect to database
282 | connection_result = await sql_tool.manage_database(
283 | action="connect",
284 | connection_string=connection_string,
285 | echo=False # Disable SQLAlchemy logging for cleaner output
286 | )
287 |
288 | if connection_result.get("success"):
289 | connection_id = connection_result.get("connection_id")
290 | db_type = connection_result.get("database_type", "Unknown")
291 |
292 | logger.success(f"Connected to database with ID: {connection_id}")
293 | console.print(Panel(
294 | f"Connection ID: [bold cyan]{escape(connection_id)}[/]\n"
295 | f"Database Type: [blue]{escape(db_type)}[/]",
296 | title="[bold green]:link: Connected[/]",
297 | border_style="green",
298 | padding=(1, 2),
299 | expand=False
300 | ))
301 |
302 | # Test the connection
303 | console.print("[cyan]Testing connection health...[/]")
304 | test_result = await sql_tool.manage_database(
305 | action="test",
306 | connection_id=connection_id
307 | )
308 |
309 | if test_result.get("success"):
310 | resp_time = test_result.get("response_time_seconds", 0)
311 | version = test_result.get("version", "N/A")
312 | console.print(Panel(
313 | f"[green]:heavy_check_mark: Connection test OK\n"
314 | f"Response time: {resp_time:.4f}s\n"
315 | f"DB Version: {version}",
316 | border_style="green",
317 | padding=(1, 2)
318 | ))
319 | else:
320 | console.print(Panel(
321 | f"[bold red]:x: Connection test failed:[/]\n{escape(test_result.get('error', 'Unknown error'))}",
322 | border_style="red",
323 | padding=(1, 2)
324 | ))
325 |
326 | # Get connection status
327 | console.print("[cyan]Fetching database status...[/]")
328 | status_result = await sql_tool.manage_database(
329 | action="status",
330 | connection_id=connection_id
331 | )
332 |
333 | if status_result.get("success"):
334 | status_table = Table(title="Active Connections", box=box.HEAVY, padding=(0, 1), border_style="blue")
335 | status_table.add_column("Connection ID", style="cyan")
336 | status_table.add_column("Database", style="blue")
337 | status_table.add_column("Last Accessed", style="dim")
338 | status_table.add_column("Idle Time", style="yellow")
339 |
340 | connections = status_result.get("connections", {})
341 | for conn_id, conn_info in connections.items():
342 | status_table.add_row(
343 | conn_id,
344 | conn_info.get("dialect", "unknown"),
345 | conn_info.get("last_accessed", "N/A"),
346 | f"{conn_info.get('idle_time_seconds', 0):.1f}s"
347 | )
348 |
349 | console.print(status_table)
350 | else:
351 | console.print(Panel(
352 | f"[bold red]:x: Failed to get database status:[/]\n{escape(status_result.get('error', 'Unknown error'))}",
353 | border_style="red",
354 | padding=(1, 2)
355 | ))
356 | else:
357 | error_msg = connection_result.get('error', 'Unknown error')
358 | logger.error(f"Failed to connect to database: {error_msg}")
359 | console.print(Panel(
360 | f"[bold red]:x: Connection failed:[/]\n{escape(error_msg)}",
361 | border_style="red",
362 | padding=(1, 2)
363 | ))
364 |
365 | except Exception as e:
366 | logger.error(f"Unexpected error in connection demo: {e}")
367 | console.print(f"[bold red]:x: Unexpected Error:[/]\n{escape(str(e))}")
368 |
369 | console.print() # Spacing
370 | return connection_id
371 |
372 | async def schema_discovery_demo(sql_tool: SQLTool, connection_id: str) -> None:
373 | """Demonstrate database schema discovery."""
374 | console.print(Rule("[bold green]2. Schema Discovery Demo[/bold green]", style="green"))
375 | logger.info("Starting schema discovery demo")
376 |
377 | with console.status("[bold cyan]Discovering database schema...", spinner="dots"):
378 | try:
379 | schema_result = await sql_tool.explore_database(
380 | connection_id=connection_id,
381 | action="schema",
382 | include_indexes=True,
383 | include_foreign_keys=True,
384 | detailed=True
385 | )
386 |
387 | if schema_result.get("success"):
388 | tables = schema_result.get("tables", [])
389 | views = schema_result.get("views", [])
390 | relationships = schema_result.get("relationships", [])
391 |
392 | logger.success(f"Schema discovered: {len(tables)} tables, {len(views)} views, {len(relationships)} relationships")
393 |
394 | # Create a tree visualization
395 | tree = Tree(
396 | f"[bold bright_blue]:database: Database Schema ({len(tables)} Tables, {len(views)} Views)[/]",
397 | guide_style="bright_blue"
398 | )
399 |
400 | # Add Tables branch
401 | if tables:
402 | tables_branch = tree.add("[bold cyan]:page_facing_up: Tables[/]")
403 | for table in tables:
404 | table_name = table.get("name", "Unknown")
405 | table_node = tables_branch.add(f"[cyan]{escape(table_name)}[/]")
406 |
407 | # Add columns
408 | cols = table.get("columns", [])
409 | if cols:
410 | cols_branch = table_node.add("[bold yellow]:heavy_minus_sign: Columns[/]")
411 | for col in cols:
412 | col_name = col.get("name", "?")
413 | col_type = col.get("type", "?")
414 | is_pk = col.get("primary_key", False)
415 | is_nullable = col.get("nullable", True)
416 | pk_str = " [bold magenta](PK)[/]" if is_pk else ""
417 | null_str = "" if is_nullable else " [dim]NOT NULL[/]"
418 | cols_branch.add(f"[yellow]{escape(col_name)}[/]: {escape(col_type)}{pk_str}{null_str}")
419 |
420 | # Add foreign keys
421 | fks = table.get("foreign_keys", [])
422 | if fks:
423 | fks_branch = table_node.add("[bold blue]:link: Foreign Keys[/]")
424 | for fk in fks:
425 | ref_table = fk.get("referred_table", "?")
426 | con_cols = ', '.join(fk.get("constrained_columns", []))
427 | ref_cols = ', '.join(fk.get("referred_columns", []))
428 | fks_branch.add(f"[blue]({escape(con_cols)})[/] -> [cyan]{escape(ref_table)}[/]({escape(ref_cols)})")
429 |
430 | # Add Views branch
431 | if views:
432 | views_branch = tree.add("[bold magenta]:scroll: Views[/]")
433 | for view in views:
434 | view_name = view.get("name", "Unknown")
435 | views_branch.add(f"[magenta]{escape(view_name)}[/]")
436 |
437 | console.print(Panel(tree, title="Schema Overview", border_style="bright_blue", padding=(1, 2)))
438 |
439 | # Show schema hash if available
440 | if schema_hash := schema_result.get("schema_hash"):
441 | console.print(f"[dim]Schema Hash: {schema_hash}[/dim]")
442 | else:
443 | error_msg = schema_result.get('error', 'Unknown error')
444 | logger.error(f"Failed to discover schema: {error_msg}")
445 | console.print(Panel(
446 | f"[bold red]:x: Schema discovery failed:[/]\n{escape(error_msg)}",
447 | border_style="red",
448 | padding=(1, 2)
449 | ))
450 |
451 | except Exception as e:
452 | logger.error(f"Unexpected error in schema discovery demo: {e}")
453 | console.print(f"[bold red]:x: Unexpected Error:[/]\n{escape(str(e))}")
454 |
455 | console.print() # Spacing
456 |
457 | async def table_details_demo(sql_tool: SQLTool, connection_id: str, table_name: str) -> None:
458 | """Demonstrate getting detailed information about a specific table."""
459 | console.print(Rule(f"[bold green]3. Table Details: [cyan]{escape(table_name)}[/cyan][/bold green]", style="green"))
460 | logger.info(f"Getting details for table: {table_name}")
461 |
462 | try:
463 | table_result = await sql_tool.explore_database(
464 | connection_id=connection_id,
465 | action="table",
466 | table_name=table_name,
467 | include_sample_data=True,
468 | sample_size=3,
469 | include_statistics=True
470 | )
471 |
472 | if table_result.get("success"):
473 | logger.success(f"Successfully retrieved details for table: {table_name}")
474 | console.print(Panel(f"[green]:heavy_check_mark: Details retrieved for [cyan]{escape(table_name)}[/]", border_style="green", padding=(0, 1)))
475 |
476 | # Display columns
477 | columns = table_result.get("columns", [])
478 | if columns:
479 | cols_table = Table(title="Columns", box=box.ROUNDED, show_header=True, padding=(0, 1), border_style="yellow")
480 | cols_table.add_column("Name", style="yellow", header_style="bold yellow")
481 | cols_table.add_column("Type", style="white")
482 | cols_table.add_column("Nullable", style="dim")
483 | cols_table.add_column("PK", style="magenta")
484 | cols_table.add_column("Default", style="dim")
485 |
486 | for column in columns:
487 | cols_table.add_row(
488 | escape(column.get("name", "?")),
489 | escape(column.get("type", "?")),
490 | ":heavy_check_mark:" if column.get("nullable", False) else ":x:",
491 | "[bold magenta]:key:[/]" if column.get("primary_key", False) else "",
492 | escape(str(column.get("default", "")))
493 | )
494 | console.print(cols_table)
495 |
496 | # Display sample data
497 | sample_data = table_result.get("sample_data", {})
498 | sample_rows = sample_data.get("rows", [])
499 | sample_cols = sample_data.get("columns", [])
500 |
501 | if sample_rows:
502 | sample_table = Table(title="Sample Data (first 3 rows)", box=box.ROUNDED, show_header=True, padding=(0, 1), border_style="green")
503 | for col_name in sample_cols:
504 | sample_table.add_column(col_name, style="dim cyan", header_style="bold cyan")
505 |
506 | for row in sample_rows:
507 | sample_table.add_row(*[escape(str(row.get(col, ""))) for col in sample_cols])
508 |
509 | console.print(sample_table)
510 |
511 | # Display row count
512 | row_count = table_result.get("row_count", "N/A")
513 | console.print(f"[cyan]Total Rows:[/] [yellow]{row_count}[/yellow]")
514 |
515 | # Display statistics if available
516 | statistics = table_result.get("statistics", {})
517 | if statistics:
518 | stats_table = Table(title="Column Statistics", box=box.SIMPLE, show_header=True, padding=(0, 1), border_style="magenta")
519 | stats_table.add_column("Column", style="cyan")
520 | stats_table.add_column("Null Count", style="yellow", justify="right")
521 | stats_table.add_column("Distinct Count", style="blue", justify="right")
522 |
523 | for col_name, stats in statistics.items():
524 | if isinstance(stats, dict) and "error" not in stats:
525 | null_count = stats.get("null_count", "N/A")
526 | distinct_count = stats.get("distinct_count", "N/A")
527 | stats_table.add_row(escape(col_name), str(null_count), str(distinct_count))
528 |
529 | console.print(stats_table)
530 | else:
531 | error_msg = table_result.get('error', 'Unknown error')
532 | logger.error(f"Failed to get table details: {error_msg}")
533 | console.print(Panel(
534 | f"[bold red]:x: Failed to get table details:[/]\n{escape(error_msg)}",
535 | border_style="red",
536 | padding=(1, 2)
537 | ))
538 |
539 | except Exception as e:
540 | logger.error(f"Unexpected error in table details demo: {e}")
541 | console.print(f"[bold red]:x: Unexpected Error:[/]\n{escape(str(e))}")
542 |
543 | console.print() # Spacing
544 |
545 | async def find_related_tables_demo(sql_tool: SQLTool, connection_id: str, table_name: str) -> None:
546 | """Demonstrate finding tables related to a specific table."""
547 | console.print(Rule(f"[bold green]4. Related Tables: [cyan]{escape(table_name)}[/cyan][/bold green]", style="green"))
548 | logger.info(f"Finding tables related to {table_name}")
549 |
550 | try:
551 | relations_result = await sql_tool.explore_database(
552 | connection_id=connection_id,
553 | action="relationships",
554 | table_name=table_name,
555 | depth=2 # Explore relationships to depth 2
556 | )
557 |
558 | if relations_result.get("success"):
559 | rel_graph = relations_result.get("relationship_graph", {})
560 | parents = rel_graph.get("parents", [])
561 | children = rel_graph.get("children", [])
562 |
563 | if parents or children:
564 | logger.success(f"Found relationships for table: {table_name}")
565 |
566 | # Create tree visualization
567 | rel_tree = Tree(f"[bold blue]:link: Relationships for [cyan]{escape(table_name)}[/][/]", guide_style="blue")
568 |
569 | # Add parent relationships (tables referenced by this table)
570 | if parents:
571 | parent_branch = rel_tree.add("[bold green]:arrow_up: References (Parents)[/]")
572 | for parent in parents:
573 | relationship = parent.get("relationship", "")
574 | target = parent.get("target", {})
575 | target_table = target.get("table", "?")
576 | parent_branch.add(f"[blue]{escape(relationship)}[/] -> [green]{escape(target_table)}[/]")
577 |
578 | # Add child relationships (tables that reference this table)
579 | if children:
580 | child_branch = rel_tree.add("[bold magenta]:arrow_down: Referenced By (Children)[/]")
581 | for child in children:
582 | relationship = child.get("relationship", "")
583 | source = child.get("source", {})
584 | source_table = source.get("table", "?")
585 | child_branch.add(f"[magenta]{escape(source_table)}[/] -> [blue]{escape(relationship)}[/]")
586 |
587 | console.print(Panel(rel_tree, title="Table Relationships", border_style="blue", padding=(1, 2)))
588 | else:
589 | logger.info(f"No direct relationships found for {table_name}")
590 | console.print(Panel(f"[yellow]No direct relationships found for '{escape(table_name)}'", border_style="yellow", padding=(0, 1)))
591 | else:
592 | error_msg = relations_result.get('error', 'Unknown error')
593 | logger.error(f"Failed to find relationships: {error_msg}")
594 | console.print(Panel(
595 | f"[bold red]:x: Failed to find relationships:[/]\n{escape(error_msg)}",
596 | border_style="red",
597 | padding=(1, 2)
598 | ))
599 |
600 | except Exception as e:
601 | logger.error(f"Unexpected error in relationship discovery demo: {e}")
602 | console.print(f"[bold red]:x: Unexpected Error:[/]\n{escape(str(e))}")
603 |
604 | console.print() # Spacing
605 |
606 | async def column_statistics_demo(sql_tool: SQLTool, connection_id: str, table_name: str, column_name: str) -> None:
607 | """Demonstrate detailed column statistics."""
608 | console.print(Rule(f"[bold green]5. Column Statistics: [cyan]{escape(table_name)}.[yellow]{escape(column_name)}[/yellow][/cyan][/bold green]", style="green"))
609 | logger.info(f"Analyzing statistics for column {table_name}.{column_name}")
610 |
611 | try:
612 | stats_result = await sql_tool.explore_database(
613 | connection_id=connection_id,
614 | action="column",
615 | table_name=table_name,
616 | column_name=column_name,
617 | histogram=True,
618 | num_buckets=8
619 | )
620 |
621 | if stats_result.get("success"):
622 | logger.success(f"Successfully analyzed statistics for {table_name}.{column_name}")
623 |
624 | # Display basic statistics
625 | statistics = stats_result.get("statistics", {})
626 | if statistics:
627 | stats_table = Table(title=f"Statistics for {column_name}", box=box.ROUNDED, show_header=False, padding=(1, 1), border_style="cyan")
628 | stats_table.add_column("Metric", style="cyan", justify="right")
629 | stats_table.add_column("Value", style="white")
630 |
631 | for key, value in statistics.items():
632 | stats_table.add_row(key.replace("_", " ").title(), str(value))
633 |
634 | console.print(stats_table)
635 |
636 | # Display histogram if available
637 | histogram = stats_result.get("histogram", {})
638 | buckets = histogram.get("buckets", [])
639 |
640 | if buckets:
641 | console.print("[bold cyan]Value Distribution:[/]")
642 |
643 | # Find the max count for scaling
644 | max_count = max(bucket.get("count", 0) for bucket in buckets)
645 |
646 | # Create a progress bar visualization for the histogram
647 | progress = Progress(
648 | TextColumn("[cyan]{task.description}", justify="right"),
649 | BarColumn(bar_width=40),
650 | TextColumn("[magenta]{task.fields[count]} ({task.percentage:>3.1f}%)")
651 | )
652 |
653 | with progress:
654 | for bucket in buckets:
655 | label = bucket.get("range", "?")
656 | count = bucket.get("count", 0)
657 | percentage = (count / max_count) * 100 if max_count > 0 else 0
658 |
659 | # Add a task for this bucket
660 | progress.add_task(
661 | description=escape(str(label)),
662 | total=100,
663 | completed=percentage,
664 | count=count
665 | )
666 | else:
667 | error_msg = stats_result.get('error', 'Unknown error')
668 | logger.error(f"Failed to analyze column statistics: {error_msg}")
669 | console.print(Panel(
670 | f"[bold red]:x: Failed to analyze column statistics:[/]\n{escape(error_msg)}",
671 | border_style="red",
672 | padding=(1, 2)
673 | ))
674 |
675 | except Exception as e:
676 | logger.error(f"Unexpected error in column statistics demo: {e}")
677 | console.print(f"[bold red]:x: Unexpected Error:[/]\n{escape(str(e))}")
678 |
679 | console.print() # Spacing
680 |
681 | async def query_execution_demo(sql_tool: SQLTool, connection_id: str) -> None:
682 | """Demonstrate query execution capabilities."""
683 | console.print(Rule("[bold green]6. Query Execution Demo[/bold green]", style="green"))
684 | logger.info("Demonstrating query execution capabilities")
685 |
686 | try:
687 | # Simple SELECT query
688 | simple_query = "SELECT customer_id, name, email, status FROM customers WHERE status = 'active'"
689 | logger.info("Executing simple query...")
690 |
691 | with console.status("[cyan]Running simple query...[/]"):
692 | query_result = await sql_tool.execute_sql(
693 | connection_id=connection_id,
694 | query=simple_query,
695 | read_only=True,
696 | max_rows=10
697 | )
698 |
699 | display_result("Simple Query: Active Customers", query_result, query_str=simple_query)
700 |
701 | # Parameterized query
702 | param_query = "SELECT product_id, name, price FROM products WHERE category = :category AND price < :max_price ORDER BY price DESC"
703 | params = {"category": "Electronics", "max_price": 1000.00}
704 |
705 | logger.info(f"Executing parameterized query with params: {params}")
706 |
707 | with console.status("[cyan]Running parameterized query...[/]"):
708 | param_result = await sql_tool.execute_sql(
709 | connection_id=connection_id,
710 | query=param_query,
711 | parameters=params,
712 | read_only=True
713 | )
714 |
715 | display_result("Parameterized Query: Electronics under $1000", param_result, query_str=param_query)
716 |
717 | # Pagination query
718 | pagination_query = "SELECT product_id, name, category, price FROM products ORDER BY price DESC"
719 | logger.info("Executing query with pagination (Page 1)")
720 |
721 | with console.status("[cyan]Running paginated query (Page 1)...[/]"):
722 | pagination_result_p1 = await sql_tool.execute_sql(
723 | connection_id=connection_id,
724 | query=pagination_query,
725 | pagination={"page": 1, "page_size": 2},
726 | read_only=True
727 | )
728 |
729 | display_result("Paginated Query: Products by Price (Page 1)", pagination_result_p1, query_str=pagination_query)
730 |
731 | # Pagination page 2
732 | logger.info("Executing query with pagination (Page 2)")
733 |
734 | with console.status("[cyan]Running paginated query (Page 2)...[/]"):
735 | pagination_result_p2 = await sql_tool.execute_sql(
736 | connection_id=connection_id,
737 | query=pagination_query,
738 | pagination={"page": 2, "page_size": 2},
739 | read_only=True
740 | )
741 |
742 | display_result("Paginated Query: Products by Price (Page 2)", pagination_result_p2)
743 |
744 | # Join query with multiple tables
745 | join_query = """
746 | SELECT c.name AS customer_name, o.order_id, o.order_date, o.total_amount, o.status
747 | FROM customers c
748 | JOIN orders o ON c.customer_id = o.customer_id
749 | WHERE c.status = 'active'
750 | ORDER BY o.order_date DESC
751 | """
752 |
753 | logger.info("Executing join query")
754 |
755 | with console.status("[cyan]Running join query...[/]"):
756 | join_result = await sql_tool.execute_sql(
757 | connection_id=connection_id,
758 | query=join_query,
759 | read_only=True
760 | )
761 |
762 | display_result("Join Query: Orders by Active Customers", join_result, query_str=join_query)
763 |
764 | except Exception as e:
765 | logger.error(f"Unexpected error in query execution demo: {e}")
766 | console.print(f"[bold red]:x: Unexpected Error:[/]\n{escape(str(e))}")
767 |
768 | console.print() # Spacing
769 |
770 | async def nl_to_sql_demo(sql_tool: SQLTool, connection_id: str) -> None:
771 | """Demonstrate natural language to SQL conversion."""
772 | console.print(Rule("[bold green]7. Natural Language to SQL Demo[/bold green]", style="green"))
773 | logger.info("Demonstrating natural language to SQL conversion")
774 |
775 | try:
776 | # Example NL query
777 | natural_language = "Show me all active customers and their total order value"
778 |
779 | logger.info(f"Converting natural language to SQL: '{natural_language}'")
780 |
781 | with console.status("[cyan]Converting natural language to SQL...[/]"):
782 | nl_result = await sql_tool.execute_sql(
783 | connection_id=connection_id,
784 | natural_language=natural_language,
785 | read_only=True
786 | )
787 |
788 | if nl_result.get("success"):
789 | generated_sql = nl_result.get("generated_sql", "")
790 | confidence = nl_result.get("confidence", 0.0)
791 |
792 | # Display the generated SQL and confidence
793 | console.print(Panel(
794 | Syntax(generated_sql, "sql", theme="default", line_numbers=False, word_wrap=True),
795 | title=f"Generated SQL (Confidence: {confidence:.2f})",
796 | border_style="green",
797 | padding=(1, 2)
798 | ))
799 |
800 | # Display the query results
801 | display_result("Natural Language Query Results", nl_result)
802 | else:
803 | error_msg = nl_result.get('error', 'Unknown error')
804 | logger.error(f"Failed to convert natural language to SQL: {error_msg}")
805 | console.print(Panel(
806 | f"[bold red]:x: Natural language conversion failed:[/]\n{escape(error_msg)}",
807 | border_style="red",
808 | padding=(1, 2)
809 | ))
810 |
811 | # Try another more complex example
812 | complex_nl = "What's the average price of products by category?"
813 |
814 | logger.info(f"Converting complex natural language to SQL: '{complex_nl}'")
815 |
816 | with console.status("[cyan]Converting complex natural language to SQL...[/]"):
817 | complex_result = await sql_tool.execute_sql(
818 | connection_id=connection_id,
819 | natural_language=complex_nl,
820 | read_only=True
821 | )
822 |
823 | if complex_result.get("success"):
824 | generated_sql = complex_result.get("generated_sql", "")
825 | confidence = complex_result.get("confidence", 0.0)
826 |
827 | console.print(Panel(
828 | Syntax(generated_sql, "sql", theme="default", line_numbers=False, word_wrap=True),
829 | title=f"Generated SQL for complex query (Confidence: {confidence:.2f})",
830 | border_style="green",
831 | padding=(1, 2)
832 | ))
833 |
834 | display_result("Complex Natural Language Query Results", complex_result)
835 | else:
836 | error_msg = complex_result.get('error', 'Unknown error')
837 | logger.error(f"Failed to convert complex natural language to SQL: {error_msg}")
838 | console.print(Panel(
839 | f"[bold red]:x: Complex natural language conversion failed:[/]\n{escape(error_msg)}",
840 | border_style="red",
841 | padding=(1, 2)
842 | ))
843 |
844 | except Exception as e:
845 | logger.error(f"Unexpected error in NL to SQL demo: {e}")
846 | console.print(f"[bold red]:x: Unexpected Error:[/]\n{escape(str(e))}")
847 |
848 | console.print() # Spacing
849 |
850 | async def documentation_demo(sql_tool: SQLTool, connection_id: str) -> None:
851 | """Demonstrate database documentation generation."""
852 | console.print(Rule("[bold green]8. Database Documentation Demo[/bold green]", style="green"))
853 | logger.info("Demonstrating database documentation generation")
854 |
855 | try:
856 | # Generate database documentation
857 | logger.info("Generating database documentation")
858 |
859 | with console.status("[cyan]Generating database documentation...[/]"):
860 | doc_result = await sql_tool.explore_database(
861 | connection_id=connection_id,
862 | action="documentation",
863 | output_format="markdown"
864 | )
865 |
866 | if doc_result.get("success"):
867 | logger.success("Successfully generated database documentation")
868 |
869 | # Display the documentation
870 | display_result("Database Documentation", doc_result)
871 |
872 | # Optionally save to file
873 | documentation = doc_result.get("documentation", "")
874 | if documentation:
875 | # Create a temporary file to save the documentation
876 | fd, doc_path = tempfile.mkstemp(suffix=".md", prefix="db_doc_")
877 | os.close(fd)
878 |
879 | with open(doc_path, "w") as f:
880 | f.write(documentation)
881 |
882 | console.print(f"[green]Documentation saved to: [cyan]{doc_path}[/cyan][/green]")
883 | else:
884 | error_msg = doc_result.get('error', 'Unknown error')
885 | logger.error(f"Failed to generate documentation: {error_msg}")
886 | console.print(Panel(
887 | f"[bold red]:x: Documentation generation failed:[/]\n{escape(error_msg)}",
888 | border_style="red",
889 | padding=(1, 2)
890 | ))
891 |
892 | except Exception as e:
893 | logger.error(f"Unexpected error in documentation demo: {e}")
894 | console.print(f"[bold red]:x: Unexpected Error:[/]\n{escape(str(e))}")
895 |
896 | console.print() # Spacing
897 |
898 | async def security_features_demo(sql_tool: SQLTool, connection_id: str) -> None:
899 | """Demonstrate security features of the SQLTool."""
900 | console.print(Rule("[bold green]10. Security Features Demo[/bold green]", style="green"))
901 | logger.info("Demonstrating security features")
902 |
903 | # --- PII MASKING DEMO ---
904 | console.print(Rule("[bold blue]10.1 PII Data Masking[/bold blue]", style="blue"))
905 | logger.info("Demonstrating PII data masking")
906 |
907 | try:
908 | console.print("[green]PII test data added successfully.[/]")
909 | # Now run a query to show masked PII data
910 | pii_select_query = """
911 | SELECT customer_id, name, email, ssn, credit_card
912 | FROM customers
913 | ORDER BY customer_id
914 | """
915 | with console.status("[cyan]Executing query with PII data...[/]"):
916 | pii_result = await sql_tool.execute_sql(
917 | connection_id=connection_id,
918 | query=pii_select_query,
919 | read_only=True
920 | )
921 |
922 | display_result("PII Masking Demo: Automatically Masked Sensitive Data", pii_result, pii_select_query)
923 |
924 | console.print(Panel(
925 | "Notice how the [bold]SSN[/bold], [bold]credit card numbers[/bold], and [bold]email addresses[/bold] are "
926 | "automatically masked according to SQLTool's masking rules, protecting sensitive information.",
927 | title="PII Masking Explanation",
928 | border_style="cyan",
929 | padding=(1, 2)
930 | ))
931 |
932 | # --- PROHIBITED STATEMENT DETECTION DEMO ---
933 | console.print(Rule("[bold blue]10.2 Prohibited Statement Detection[/bold blue]", style="blue"))
934 | logger.info("Demonstrating prohibited statement detection")
935 |
936 | # List of prohibited statements to test
937 | prohibited_queries = [
938 | "DROP TABLE customers",
939 | "DELETE FROM products",
940 | "TRUNCATE TABLE orders",
941 | "ALTER TABLE customers DROP COLUMN name",
942 | "GRANT ALL PRIVILEGES ON products TO user",
943 | "CREATE USER hacker WITH PASSWORD 'password'"
944 | ]
945 |
946 | prohibited_table = Table(title="Prohibited Statement Detection", box=box.ROUNDED, show_header=True, padding=(0, 1), border_style="red")
947 | prohibited_table.add_column("Prohibited SQL", style="yellow")
948 | prohibited_table.add_column("Result", style="green")
949 |
950 | for query in prohibited_queries:
951 | try:
952 | with console.status(f"[cyan]Testing: {query}[/]"):
953 | await sql_tool.execute_sql(
954 | connection_id=connection_id,
955 | query=query,
956 | read_only=True
957 | )
958 | # If we get here, protection failed (no exception was raised)
959 | prohibited_table.add_row(query, "[red]FAILED - Statement was allowed![/]")
960 | except ToolError as e:
961 | # This is expected behavior - statement should be blocked
962 | prohibited_table.add_row(query, f"[green]SUCCESS - Blocked: {str(e).split(':')[0]}[/]")
963 | except Exception as e:
964 | prohibited_table.add_row(query, f"[yellow]ERROR: {str(e)[:50]}...[/]")
965 |
966 | console.print(prohibited_table)
967 |
968 | # --- ACL CONTROLS DEMO ---
969 | console.print(Rule("[bold blue]10.3 Access Control Lists (ACL)[/bold blue]", style="blue"))
970 | logger.info("Demonstrating ACL controls")
971 |
972 | # Set up ACL restrictions
973 | console.print("[cyan]Setting up ACL restrictions...[/]")
974 | # We'll restrict access to the 'customers' table and the 'credit_card' column
975 | sql_tool.update_acl(tables=["customers"], columns=["credit_card", "ssn"])
976 |
977 | console.print(Panel(
978 | "Access control lists configured:\n"
979 | "- Restricted tables: [red]customers[/]\n"
980 | "- Restricted columns: [red]credit_card, ssn[/]",
981 | title="ACL Configuration",
982 | border_style="yellow",
983 | padding=(1, 2)
984 | ))
985 |
986 | # Try to access restricted table
987 | restricted_table_query = "SELECT * FROM customers"
988 | console.print("\n[cyan]Attempting to query restricted table:[/]")
989 | console.print(Syntax(restricted_table_query, "sql", theme="default"))
990 |
991 | try:
992 | with console.status("[cyan]Executing query on restricted table...[/]"):
993 | await sql_tool.execute_sql(
994 | connection_id=connection_id,
995 | query=restricted_table_query,
996 | read_only=True
997 | )
998 | console.print("[red]ACL FAILURE: Query was allowed on restricted table![/]")
999 | except ToolError as e:
1000 | console.print(Panel(
1001 | f"[green]✅ ACL WORKING: Access denied as expected:[/]\n{escape(str(e))}",
1002 | border_style="green",
1003 | padding=(1, 2)
1004 | ))
1005 |
1006 | # Try to access restricted column
1007 | restricted_column_query = "SELECT customer_id, name, credit_card FROM products JOIN customers USING(customer_id)"
1008 | console.print("\n[cyan]Attempting to query restricted column:[/]")
1009 | console.print(Syntax(restricted_column_query, "sql", theme="default"))
1010 |
1011 | try:
1012 | with console.status("[cyan]Executing query with restricted column...[/]"):
1013 | await sql_tool.execute_sql(
1014 | connection_id=connection_id,
1015 | query=restricted_column_query,
1016 | read_only=True
1017 | )
1018 | console.print("[red]ACL FAILURE: Query was allowed with restricted column![/]")
1019 | except ToolError as e:
1020 | console.print(Panel(
1021 | f"[green]✅ ACL WORKING: Access denied as expected:[/]\n{escape(str(e))}",
1022 | border_style="green",
1023 | padding=(1, 2)
1024 | ))
1025 |
1026 | # Clear ACL restrictions for further demos
1027 | sql_tool.update_acl(tables=[], columns=[])
1028 | console.print("[cyan]ACL restrictions cleared for following demos.[/]")
1029 |
1030 | # --- SCHEMA DRIFT DETECTION ---
1031 | console.print(Rule("[bold blue]10.4 Schema Drift Detection[/bold blue]", style="blue"))
1032 | logger.info("Demonstrating schema drift detection")
1033 |
1034 | # First run schema discovery to capture initial state
1035 | console.print("[cyan]Capturing initial schema state...[/]")
1036 |
1037 | with console.status("[cyan]Performing initial schema discovery...[/]"):
1038 | initial_schema = await sql_tool.explore_database(
1039 | connection_id=connection_id,
1040 | action="schema",
1041 | include_indexes=True,
1042 | include_foreign_keys=True
1043 | )
1044 |
1045 | initial_hash = initial_schema.get("schema_hash", "unknown")
1046 | console.print(f"[green]Initial schema captured with hash: [bold]{initial_hash[:16]}...[/][/]")
1047 |
1048 | # Now make a schema change
1049 | schema_change_query = "ALTER TABLE products ADD COLUMN last_updated TIMESTAMP"
1050 |
1051 | console.print("[cyan]Making a schema change...[/]")
1052 | console.print(Syntax(schema_change_query, "sql", theme="default"))
1053 |
1054 | with console.status("[cyan]Executing schema change...[/]"):
1055 | # Execute the schema change
1056 | await sql_tool.execute_sql(
1057 | connection_id=connection_id,
1058 | query=schema_change_query,
1059 | read_only=False # Need to disable read-only for ALTER TABLE
1060 | )
1061 |
1062 | # Now run schema discovery again to detect the change
1063 | with console.status("[cyan]Performing follow-up schema discovery to detect changes...[/]"):
1064 | new_schema = await sql_tool.explore_database(
1065 | connection_id=connection_id,
1066 | action="schema",
1067 | include_indexes=True,
1068 | include_foreign_keys=True
1069 | )
1070 |
1071 | new_hash = new_schema.get("schema_hash", "unknown")
1072 | schema_changed = new_schema.get("schema_change_detected", False)
1073 |
1074 | if initial_hash != new_hash:
1075 | console.print(Panel(
1076 | f"[green]✅ SCHEMA DRIFT DETECTED:[/]\n"
1077 | f"- Initial hash: [dim]{initial_hash[:16]}...[/]\n"
1078 | f"- New hash: [bold]{new_hash[:16]}...[/]\n"
1079 | f"- Change detected by system: {'[green]Yes[/]' if schema_changed else '[red]No[/]'}",
1080 | title="Schema Drift Detection Result",
1081 | border_style="green",
1082 | padding=(1, 2)
1083 | ))
1084 | else:
1085 | console.print(Panel(
1086 | "[red]Schema drift detection did not identify a change in hash even though schema was modified.[/]",
1087 | border_style="red",
1088 | padding=(1, 2)
1089 | ))
1090 |
1091 | except Exception as e:
1092 | logger.error(f"Error in security features demo: {e}", exc_info=True)
1093 | console.print(Panel(
1094 | f"[bold red]Error in security features demo:[/]\n{escape(str(e))}",
1095 | border_style="red",
1096 | padding=(1, 2)
1097 | ))
1098 |
1099 | console.print() # Spacing
1100 |
1101 | async def advanced_export_demo(sql_tool: SQLTool, connection_id: str) -> None:
1102 | """Demonstrate advanced export options."""
1103 | console.print(Rule("[bold green]11. Advanced Export Options Demo[/bold green]", style="green"))
1104 | logger.info("Demonstrating advanced export options")
1105 |
1106 | # Query to export
1107 | export_query = """
1108 | SELECT p.product_id, p.name AS product_name, p.category, p.price,
1109 | SUM(oi.quantity) AS units_sold,
1110 | SUM(oi.quantity * oi.price_per_unit) AS total_revenue
1111 | FROM products p
1112 | LEFT JOIN order_items oi ON p.product_id = oi.product_id
1113 | GROUP BY p.product_id, p.name, p.category, p.price
1114 | ORDER BY total_revenue DESC
1115 | """
1116 |
1117 | try:
1118 | # --- PANDAS DATAFRAME EXPORT ---
1119 | console.print(Rule("[bold blue]11.1 Pandas DataFrame Export[/bold blue]", style="blue"))
1120 | logger.info("Demonstrating Pandas DataFrame export")
1121 |
1122 | console.print(Syntax(export_query, "sql", theme="default", line_numbers=False))
1123 |
1124 | with console.status("[cyan]Executing query and exporting to Pandas DataFrame...[/]"):
1125 | df_result = await sql_tool.execute_sql(
1126 | connection_id=connection_id,
1127 | query=export_query,
1128 | read_only=True,
1129 | export={"format": "pandas"}
1130 | )
1131 |
1132 | if df_result.get("success") and "dataframe" in df_result:
1133 | df = df_result["dataframe"]
1134 |
1135 | # Display DataFrame info
1136 | df_info = [
1137 | f"Shape: {df.shape[0]} rows × {df.shape[1]} columns",
1138 | f"Memory usage: {df.memory_usage(deep=True).sum() / 1024:.2f} KB",
1139 | f"Column dtypes: {', '.join([f'{col}: {dtype}' for col, dtype in df.dtypes.items()])}"
1140 | ]
1141 |
1142 | console.print(Panel(
1143 | "\n".join(df_info),
1144 | title="Pandas DataFrame Export Result",
1145 | border_style="green",
1146 | padding=(1, 2)
1147 | ))
1148 |
1149 | # Show DataFrame operations
1150 | console.print("[cyan]Demonstrating DataFrame operations:[/]")
1151 |
1152 | # Create a summary statistics table
1153 | stats_table = Table(title="DataFrame Statistics", box=box.ROUNDED, padding=(0, 1), border_style="blue")
1154 | stats_table.add_column("Statistic", style="cyan")
1155 | stats_table.add_column("Value", style="yellow")
1156 |
1157 | # Add some sample statistics
1158 | stats_table.add_row("Average Price", f"${df['price'].mean():.2f}")
1159 | stats_table.add_row("Max Price", f"${df['price'].max():.2f}")
1160 | stats_table.add_row("Min Price", f"${df['price'].min():.2f}")
1161 | stats_table.add_row("Total Revenue", f"${df['total_revenue'].sum():.2f}")
1162 | stats_table.add_row("Highest Revenue Product", df.loc[df['total_revenue'].idxmax()]['product_name'])
1163 |
1164 | console.print(stats_table)
1165 |
1166 | # Create a simple DataFrame transformation
1167 | console.print("\n[cyan]Demonstrating DataFrame transformation - Adding discount column:[/]")
1168 | df['discount_price'] = df['price'] * 0.9
1169 |
1170 | # Display the first few rows of the transformed DataFrame
1171 | table = Table(title="Transformed DataFrame (First 3 Rows)", box=box.ROUNDED, show_header=True)
1172 |
1173 | # Add columns based on the DataFrame
1174 | for col in df.columns:
1175 | justify = "right" if df[col].dtype.kind in 'ifc' else "left"
1176 | table.add_column(col, style="cyan", justify=justify)
1177 |
1178 | # Add the first 3 rows
1179 | for _, row in df.head(3).iterrows():
1180 | # Format numeric values nicely
1181 | formatted_row = []
1182 | for col in df.columns:
1183 | val = row[col]
1184 | if pd.api.types.is_numeric_dtype(df[col].dtype): # Check column dtype, not row value
1185 | if 'price' in col or 'revenue' in col:
1186 | formatted_row.append(f"${val:.2f}")
1187 | else:
1188 | formatted_row.append(f"{val:,.2f}" if isinstance(val, float) else f"{val:,}")
1189 | else:
1190 | formatted_row.append(str(val))
1191 |
1192 | table.add_row(*formatted_row)
1193 |
1194 | console.print(table)
1195 | else:
1196 | console.print(Panel(
1197 | f"[red]Failed to export to DataFrame: {df_result.get('error', 'Unknown error')}[/]",
1198 | border_style="red",
1199 | padding=(1, 2)
1200 | ))
1201 |
1202 | # --- EXCEL EXPORT WITH FORMATTING ---
1203 | console.print(Rule("[bold blue]11.2 Excel Export with Formatting[/bold blue]", style="blue"))
1204 | logger.info("Demonstrating Excel export with formatting")
1205 |
1206 | excel_fd, excel_path = tempfile.mkstemp(suffix=".xlsx", prefix="sql_demo_export_")
1207 | os.close(excel_fd) # Close file descriptor, as we only need the path
1208 |
1209 | with console.status("[cyan]Executing query and exporting to formatted Excel...[/]"):
1210 | excel_result = await sql_tool.execute_sql(
1211 | connection_id=connection_id,
1212 | query=export_query,
1213 | read_only=True,
1214 | export={
1215 | "format": "excel",
1216 | "path": excel_path,
1217 | # Note: Additional formatting options might be available in your implementation
1218 | }
1219 | )
1220 |
1221 | if excel_result.get("success") and "excel_path" in excel_result:
1222 | export_path = excel_result["excel_path"]
1223 | file_size = os.path.getsize(export_path) / 1024 # Size in KB
1224 |
1225 | console.print(Panel(
1226 | f"[green]✅ Successfully exported to Excel:[/]\n"
1227 | f"Path: [cyan]{export_path}[/]\n"
1228 | f"Size: [yellow]{file_size:.2f} KB[/]",
1229 | title="Excel Export Result",
1230 | border_style="green",
1231 | padding=(1, 2)
1232 | ))
1233 | else:
1234 | console.print(Panel(
1235 | f"[red]Failed to export to Excel: {excel_result.get('error', 'Unknown error')}[/]",
1236 | border_style="red",
1237 | padding=(1, 2)
1238 | ))
1239 |
1240 | # --- CUSTOM EXPORT PATH (CSV) ---
1241 | console.print(Rule("[bold blue]11.3 Custom Export Path (CSV)[/bold blue]", style="blue"))
1242 | logger.info("Demonstrating custom export path")
1243 |
1244 | # Create a custom path in the user's home directory
1245 | user_home = os.path.expanduser("~")
1246 | custom_dir = os.path.join(user_home, "sql_demo_exports")
1247 | os.makedirs(custom_dir, exist_ok=True)
1248 |
1249 | timestamp = dt.datetime.now().strftime("%Y%m%d_%H%M%S")
1250 | custom_path = os.path.join(custom_dir, f"product_sales_{timestamp}.csv")
1251 |
1252 | console.print(f"[cyan]Exporting to custom path: [/][yellow]{custom_path}[/]")
1253 |
1254 | with console.status("[cyan]Executing query and exporting to custom CSV path...[/]"):
1255 | csv_result = await sql_tool.execute_sql(
1256 | connection_id=connection_id,
1257 | query=export_query,
1258 | read_only=True,
1259 | export={
1260 | "format": "csv",
1261 | "path": custom_path
1262 | }
1263 | )
1264 |
1265 | if csv_result.get("success") and "csv_path" in csv_result:
1266 | export_path = csv_result["csv_path"]
1267 | file_size = os.path.getsize(export_path) / 1024 # Size in KB
1268 |
1269 | # Read first few lines to show content
1270 | with open(export_path, 'r') as f:
1271 | first_lines = [next(f) for _ in range(3)]
1272 |
1273 | console.print(Panel(
1274 | f"[green]✅ Successfully exported to custom CSV path:[/]\n"
1275 | f"Path: [cyan]{export_path}[/]\n"
1276 | f"Size: [yellow]{file_size:.2f} KB[/]\n\n"
1277 | f"[dim]Preview (first 3 lines):[/]\n"
1278 | f"[white]{escape(''.join(first_lines))}[/]",
1279 | title="Custom CSV Export Result",
1280 | border_style="green",
1281 | padding=(1, 2)
1282 | ))
1283 | else:
1284 | console.print(Panel(
1285 | f"[red]Failed to export to custom CSV path: {csv_result.get('error', 'Unknown error')}[/]",
1286 | border_style="red",
1287 | padding=(1, 2)
1288 | ))
1289 |
1290 | except Exception as e:
1291 | logger.error(f"Error in advanced export demo: {e}", exc_info=True)
1292 | console.print(Panel(
1293 | f"[bold red]Error in advanced export demo:[/]\n{escape(str(e))}",
1294 | border_style="red",
1295 | padding=(1, 2)
1296 | ))
1297 |
1298 | console.print() # Spacing
1299 |
1300 | async def schema_validation_demo(sql_tool: SQLTool, connection_id: str) -> None:
1301 | """Demonstrate Pandera schema validation for query results."""
1302 | console.print(Rule("[bold green]12. Schema Validation Demo[/bold green]", style="green"))
1303 | logger.info("Demonstrating Pandera schema validation")
1304 |
1305 | try:
1306 | # Query to validate
1307 | validation_query = """
1308 | SELECT
1309 | product_id,
1310 | name AS product_name,
1311 | price,
1312 | category,
1313 | in_stock
1314 | FROM products
1315 | """
1316 |
1317 | console.print("[cyan]We'll validate that query results conform to a specified schema:[/]")
1318 | console.print(Syntax(validation_query, "sql", theme="default"))
1319 |
1320 | # Define a Pandera schema
1321 | schema_code = """
1322 | # Define a Pandera schema for validation using DataFrameSchema
1323 | product_schema = pa.DataFrameSchema({
1324 | "product_id": pa.Column(int, checks=pa.Check.greater_than(0)),
1325 | "product_name": pa.Column(str, nullable=False),
1326 | "price": pa.Column(float, checks=[
1327 | pa.Check.greater_than(0, error="price must be positive"),
1328 | pa.Check.less_than(2000.0, error="price must be under $2000")
1329 | ]),
1330 | "category": pa.Column(
1331 | str,
1332 | checks=pa.Check.isin(["Electronics", "Audio", "Kitchen", "Wearables"]),
1333 | nullable=False
1334 | ),
1335 | "in_stock": pa.Column(bool)
1336 | })
1337 | """
1338 |
1339 | console.print(Panel(
1340 | Syntax(schema_code, "python", theme="default"),
1341 | title="Pandera Validation Schema",
1342 | border_style="cyan",
1343 | padding=(1, 2)
1344 | ))
1345 |
1346 | # Check pandera version
1347 | version = getattr(pa, '__version__', 'unknown')
1348 | console.print(f"[dim]Using pandera version: {version}")
1349 |
1350 | # Define the actual schema
1351 | product_schema = pa.DataFrameSchema({
1352 | "product_id": pa.Column(int, checks=pa.Check.greater_than(0)),
1353 | "product_name": pa.Column(str, nullable=False),
1354 | "price": pa.Column(float, checks=[
1355 | pa.Check.greater_than(0, error="price must be positive"),
1356 | pa.Check.less_than(2000.0, error="price must be under $2000")
1357 | ]),
1358 | "category": pa.Column(
1359 | str,
1360 | checks=pa.Check.isin(["Electronics", "Audio", "Kitchen", "Wearables"]),
1361 | nullable=False
1362 | ),
1363 | "in_stock": pa.Column(bool)
1364 | })
1365 |
1366 | # WORKAROUND: Instead of using built-in validation (which has an error),
1367 | # we'll fetch the data first, then validate it manually
1368 | console.print("[cyan]Executing query to fetch data...[/]")
1369 |
1370 | with console.status("[cyan]Running query...[/]"):
1371 | query_result = await sql_tool.execute_sql(
1372 | connection_id=connection_id,
1373 | query=validation_query,
1374 | read_only=True
1375 | )
1376 |
1377 | if query_result.get("success"):
1378 | # Show the data
1379 | display_result("Data Retrieved for Validation", query_result)
1380 |
1381 | # Now manually validate with Pandera
1382 | console.print("[cyan]Now validating results with Pandera...[/]")
1383 |
1384 | if pd is not None:
1385 | try:
1386 | # Create DataFrame from results
1387 | df = pd.DataFrame(query_result.get("rows", []), columns=query_result.get("columns", []))
1388 |
1389 | # Fix type issues - convert in_stock to boolean if needed
1390 | if "in_stock" in df.columns and df["in_stock"].dtype != bool:
1391 | df["in_stock"] = df["in_stock"].astype(bool)
1392 |
1393 | console.print(f"[dim]Created DataFrame with shape {df.shape} for validation")
1394 |
1395 | # Validate the data
1396 | with console.status("[cyan]Validating against schema...[/]"):
1397 | try:
1398 | product_schema.validate(df)
1399 | console.print(Panel(
1400 | "[green]✅ Schema validation passed![/]\n"
1401 | "All data meets the requirements defined in the schema.",
1402 | title="Validation Result",
1403 | border_style="green",
1404 | padding=(1, 2)
1405 | ))
1406 | except Exception as val_err:
1407 | console.print(Panel(
1408 | f"[yellow]⚠ Schema validation failed![/]\n"
1409 | f"Error: {str(val_err)}",
1410 | title="Validation Result",
1411 | border_style="yellow",
1412 | padding=(1, 2)
1413 | ))
1414 | except Exception as df_err:
1415 | console.print(f"[red]Error creating DataFrame: {df_err}[/]")
1416 | else:
1417 | console.print("[yellow]Pandas is not available, cannot perform validation.[/]")
1418 | else:
1419 | console.print(Panel(
1420 | f"[red]Failed to execute query: {query_result.get('error', 'Unknown error')}[/]",
1421 | border_style="red",
1422 | padding=(1, 2)
1423 | ))
1424 |
1425 | # Simulate a failing validation case
1426 | console.print("\n[cyan]Simulating validation failure with invalid data...[/]")
1427 |
1428 | if pd is not None:
1429 | # Create a DataFrame with valid and invalid data
1430 | test_data = [
1431 | # Valid data
1432 | {"product_id": 1, "product_name": "Laptop Pro X", "price": 1499.99, "category": "Electronics", "in_stock": True},
1433 | {"product_id": 2, "product_name": "Smartphone Z", "price": 999.99, "category": "Electronics", "in_stock": True},
1434 | # Invalid data (negative price)
1435 | {"product_id": 6, "product_name": "Invalid Product", "price": -10.0, "category": "Electronics", "in_stock": True},
1436 | # Invalid data (unknown category)
1437 | {"product_id": 7, "product_name": "Test Product", "price": 50.0, "category": "Invalid Category", "in_stock": True}
1438 | ]
1439 |
1440 | test_df = pd.DataFrame(test_data)
1441 |
1442 | # Display the test data
1443 | test_table = Table(title="Test Data for Validation", box=box.ROUNDED, show_header=True, padding=(0, 1), border_style="yellow")
1444 | for col in test_df.columns:
1445 | test_table.add_column(str(col), style="cyan")
1446 |
1447 | for _, row in test_df.iterrows():
1448 | test_table.add_row(*[str(val) for val in row])
1449 |
1450 | console.print(test_table)
1451 |
1452 | # Test validation
1453 | console.print("[cyan]Attempting to validate this data...[/]")
1454 | try:
1455 | # Try to validate the DataFrame directly
1456 | product_schema.validate(test_df, lazy=True)
1457 | console.print(Panel(
1458 | "[red]Unexpected result: Validation passed when it should have failed![/]",
1459 | border_style="red",
1460 | padding=(1, 2)
1461 | ))
1462 | except Exception as val_err:
1463 | console.print(Panel(
1464 | f"[green]✅ Validation correctly failed as expected![/]\n"
1465 | f"Error: {str(val_err)}",
1466 | title="Expected Validation Failure (Simulated)",
1467 | border_style="green",
1468 | padding=(1, 2)
1469 | ))
1470 | else:
1471 | console.print("[yellow]Pandas not available, cannot demonstrate validation failure.[/]")
1472 |
1473 | except Exception as e:
1474 | logger.error(f"Error in schema validation demo: {e}", exc_info=True)
1475 | console.print(Panel(
1476 | f"[bold red]Error in schema validation demo:[/]\n{escape(str(e))}",
1477 | border_style="red",
1478 | padding=(1, 2)
1479 | ))
1480 |
1481 | console.print() # Spacing
1482 |
1483 | async def audit_log_demo(sql_tool: SQLTool, connection_id: str) -> None:
1484 | """Demonstrate audit log functionality."""
1485 | console.print(Rule("[bold green]9. Audit Log Demo[/bold green]", style="green"))
1486 | logger.info("Demonstrating audit log functionality")
1487 |
1488 | try:
1489 | # View the audit log
1490 | logger.info("Viewing audit log")
1491 |
1492 | with console.status("[cyan]Retrieving audit log...[/]"):
1493 | audit_result = await sql_tool.access_audit_log(
1494 | action="view",
1495 | limit=10
1496 | )
1497 |
1498 | if audit_result.get("success"):
1499 | logger.success("Successfully retrieved audit log")
1500 |
1501 | records = audit_result.get("records", [])
1502 | if records:
1503 | audit_table = Table(title="Audit Log", box=box.ROUNDED, show_header=True, padding=(0, 1), border_style="blue")
1504 | audit_table.add_column("ID", style="dim")
1505 | audit_table.add_column("Timestamp", style="cyan")
1506 | audit_table.add_column("Tool", style="green")
1507 | audit_table.add_column("Action", style="yellow")
1508 | audit_table.add_column("Connection ID", style="magenta")
1509 | audit_table.add_column("Success", style="cyan")
1510 |
1511 | for record in records:
1512 | audit_table.add_row(
1513 | record.get("audit_id", "?"),
1514 | record.get("timestamp", "?"),
1515 | record.get("tool_name", "?"),
1516 | record.get("action", "?"),
1517 | record.get("connection_id", "?"),
1518 | "[green]:heavy_check_mark:[/]" if record.get("success") else "[red]:x:[/]"
1519 | )
1520 |
1521 | console.print(audit_table)
1522 |
1523 | # Show details of one specific audit record
1524 | if records:
1525 | sample_record = records[0]
1526 | console.print(Panel(
1527 | "\n".join([f"[cyan]{k}:[/] {escape(str(v))}" for k, v in sample_record.items() if k not in ["audit_id", "timestamp", "tool_name", "action", "connection_id", "success"]]),
1528 | title=f"Audit Record Details: {sample_record.get('audit_id', '?')}",
1529 | border_style="dim",
1530 | padding=(1, 2)
1531 | ))
1532 | else:
1533 | console.print(Panel("[yellow]No audit records found.", border_style="yellow", padding=(0, 1)))
1534 | else:
1535 | error_msg = audit_result.get('error', 'Unknown error')
1536 | logger.error(f"Failed to retrieve audit log: {error_msg}")
1537 | console.print(Panel(
1538 | f"[bold red]:x: Audit log retrieval failed:[/]\n{escape(error_msg)}",
1539 | border_style="red",
1540 | padding=(1, 2)
1541 | ))
1542 |
1543 | # Export the audit log
1544 | logger.info("Exporting audit log")
1545 |
1546 | with console.status("[cyan]Exporting audit log to CSV...[/]"):
1547 | export_result = await sql_tool.access_audit_log(
1548 | action="export",
1549 | export_format="csv"
1550 | )
1551 |
1552 | if export_result.get("success"):
1553 | export_path = export_result.get("path", "")
1554 | record_count = export_result.get("record_count", 0)
1555 | logger.success(f"Successfully exported {record_count} audit records to CSV")
1556 |
1557 | console.print(Panel(
1558 | f"[green]:heavy_check_mark: Exported {record_count} audit records to:[/]\n[cyan]{export_path}[/]",
1559 | border_style="green",
1560 | padding=(1, 2)
1561 | ))
1562 | else:
1563 | error_msg = export_result.get('error', 'Unknown error')
1564 | logger.error(f"Failed to export audit log: {error_msg}")
1565 | console.print(Panel(
1566 | f"[bold red]:x: Audit log export failed:[/]\n{escape(error_msg)}",
1567 | border_style="red",
1568 | padding=(1, 2)
1569 | ))
1570 |
1571 | except Exception as e:
1572 | logger.error(f"Unexpected error in audit log demo: {e}")
1573 | console.print(f"[bold red]:x: Unexpected Error:[/]\n{escape(str(e))}")
1574 |
1575 | console.print() # Spacing
1576 |
1577 | async def cleanup_demo(sql_tool: SQLTool, connection_id: str) -> None:
1578 | """Demonstrate disconnecting from the database."""
1579 | console.print(Rule("[bold green]Database Cleanup and Disconnection[/bold green]", style="green"))
1580 | logger.info("Disconnecting from database")
1581 |
1582 | try:
1583 | # Disconnect from the database
1584 | disconnect_result = await sql_tool.manage_database(
1585 | action="disconnect",
1586 | connection_id=connection_id
1587 | )
1588 |
1589 | if disconnect_result.get("success"):
1590 | logger.success(f"Successfully disconnected from database (ID: {connection_id})")
1591 | console.print(Panel(
1592 | f"[green]:heavy_check_mark: Successfully disconnected from database. Connection ID: [dim]{connection_id}[/dim][/]",
1593 | border_style="green",
1594 | padding=(0, 1)
1595 | ))
1596 | else:
1597 | logger.error(f"Failed to disconnect: {disconnect_result.get('error')}")
1598 | console.print(Panel(
1599 | f"[bold red]:x: Failed to disconnect:[/]\n{escape(disconnect_result.get('error', 'Unknown error'))}",
1600 | border_style="red",
1601 | padding=(1, 2)
1602 | ))
1603 |
1604 | except Exception as e:
1605 | logger.error(f"Error in cleanup demo: {e}")
1606 | console.print(f"[bold red]:x: Error in cleanup:[/]\n{escape(str(e))}")
1607 |
1608 | console.print()
1609 |
1610 | async def verify_demo_database(sql_tool, connection_id: str) -> None:
1611 | """Verify the demo database has been set up correctly."""
1612 | logger.info("Verifying database setup...")
1613 |
1614 | # For consistency, we'll still display the setup status
1615 | console.print(Panel("[green]:heavy_check_mark: Using prepared sample database.", padding=(0, 1), border_style="green"))
1616 |
1617 | # Check the tables to ensure the database was set up correctly
1618 | try:
1619 | # Execute a simple query to check if the tables have data
1620 | result = await sql_tool.execute_sql(
1621 | connection_id=connection_id,
1622 | query="SELECT COUNT(*) as count FROM customers",
1623 | read_only=True
1624 | )
1625 |
1626 | count = result.get("rows", [{}])[0].get("count", 0)
1627 | if count > 0:
1628 | logger.info(f"Verified database setup: {count} customers found")
1629 | console.print(Panel(f"[green]:heavy_check_mark: Sample database verified with {count} customer records.", padding=(0, 1), border_style="green"))
1630 | else:
1631 | logger.warning("Database tables found but they appear to be empty")
1632 | console.print(Panel("[yellow]⚠ Database tables found but they appear to be empty.", padding=(0, 1), border_style="yellow"))
1633 |
1634 | except (ToolError, ToolInputError) as e:
1635 | logger.error(f"Error checking database setup: {e}")
1636 | console.print(Panel(f"[bold red]:x: Database Setup Error:[/]\n{escape(str(e))}", padding=(1, 2), border_style="red"))
1637 |
1638 | # --- Main Function ---
1639 |
1640 | async def main() -> int:
1641 | """Run the SQL database tools demo."""
1642 | console.print(Rule("[bold magenta]SQL Database Tools Demo[/bold magenta]"))
1643 |
1644 | exit_code = 0
1645 | connection_id = None
1646 |
1647 | # Get path to the pre-initialized database
1648 | db_file = os.path.join(os.path.dirname(__file__), "demo.db")
1649 |
1650 | # Force recreate the demo database
1651 | if os.path.exists(db_file):
1652 | try:
1653 | os.remove(db_file)
1654 | console.print("[yellow]Removed existing database file to ensure correct schema.[/]")
1655 | except OSError as e:
1656 | console.print(f"[yellow]Warning: Could not remove existing database: {e}[/]")
1657 |
1658 | # Check if the demo database exists, and create it if not
1659 | if not os.path.exists(db_file):
1660 | console.print("[yellow]Demo database not found. Creating it now...[/]")
1661 | try:
1662 | # Initialize the database directly
1663 | init_demo_database(db_file)
1664 | console.print("[green]Demo database created successfully.[/]")
1665 | except Exception as e:
1666 | console.print(f"[red]Failed to create demo database: {e}[/]")
1667 | return 1
1668 |
1669 | gateway = Gateway("sql-database-demo", register_tools=False)
1670 |
1671 | # Connection string for file-based SQLite database instead of memory
1672 | file_connection_string = f"sqlite:///{db_file}"
1673 |
1674 | # Create an instance of the SQLTool
1675 | try:
1676 | sql_tool = SQLTool(gateway)
1677 |
1678 | # Run the demonstrations
1679 | connection_id = await connection_demo(sql_tool, file_connection_string)
1680 |
1681 | if connection_id:
1682 | await verify_demo_database(sql_tool, connection_id)
1683 | await schema_discovery_demo(sql_tool, connection_id)
1684 | await table_details_demo(sql_tool, connection_id, "customers")
1685 | await find_related_tables_demo(sql_tool, connection_id, "orders")
1686 | await column_statistics_demo(sql_tool, connection_id, "products", "price")
1687 | await query_execution_demo(sql_tool, connection_id)
1688 | await nl_to_sql_demo(sql_tool, connection_id)
1689 | await documentation_demo(sql_tool, connection_id)
1690 | await audit_log_demo(sql_tool, connection_id)
1691 |
1692 | # Add the new demos
1693 | await security_features_demo(sql_tool, connection_id)
1694 | await advanced_export_demo(sql_tool, connection_id)
1695 | await schema_validation_demo(sql_tool, connection_id)
1696 |
1697 | await cleanup_demo(sql_tool, connection_id)
1698 | else:
1699 | logger.error("Skipping demonstrations due to connection failure")
1700 | exit_code = 1
1701 |
1702 | except Exception as e:
1703 | logger.critical(f"Demo failed with unexpected error: {e}")
1704 | console.print(f"[bold red]CRITICAL ERROR: {escape(str(e))}[/]")
1705 | exit_code = 1
1706 | finally:
1707 | # Ensure we shutdown the SQLTool if it was created
1708 | if 'sql_tool' in locals():
1709 | try:
1710 | await sql_tool.shutdown()
1711 | logger.info("SQLTool shut down successfully")
1712 | except Exception as shutdown_err:
1713 | logger.error(f"Error during SQLTool shutdown: {shutdown_err}")
1714 | # Clean up the demo database file
1715 | try:
1716 | if os.path.exists(db_file) and 'sql_demo_export' in db_file:
1717 | os.remove(db_file)
1718 | logger.info(f"Cleaned up demo database file: {db_file}")
1719 | except Exception as clean_err:
1720 | logger.warning(f"Could not clean up demo database: {clean_err}")
1721 |
1722 | return exit_code
1723 |
1724 | if __name__ == "__main__":
1725 | # Setup logging
1726 | import logging
1727 | logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(name)s - %(levelname)s - %(message)s')
1728 |
1729 | # Run the demo
1730 | exit_code = asyncio.run(main())
1731 | sys.exit(exit_code)
```