Spaces:

amaye15
/

DuckDB-UI

Running

+<!DOCTYPE html>
+<html lang="en">
+<head>
+    <meta charset="UTF-8">
+    <meta name="viewport" content="width=device-width, initial-scale=1.0">
+    <title>DuckDB Explorer</title>
+    <style>
+        body {
+            font-family: -apple-system, BlinkMacSystemFont, "Segoe UI", Roboto, Helvetica, Arial, sans-serif;
+            margin: 0;
+            padding: 0;
+            display: flex;
+            flex-direction: column;
+            height: 100vh;
+            background-color: #f4f7f6;
+            color: #333;
+        }
+        header {
+            background-color: #4CAF50;
+            color: white;
+            padding: 10px 20px;
+            display: flex;
+            align-items: center;
+            box-shadow: 0 2px 4px rgba(0,0,0,0.1);
+        }
+        header input[type="text"] {
+            flex-grow: 1;
+            margin-right: 10px;
+            padding: 8px;
+            border: 1px solid #ccc;
+            border-radius: 4px;
+        }
+        header button {
+            padding: 8px 15px;
+            background-color: #367c39;
+            color: white;
+            border: none;
+            border-radius: 4px;
+            cursor: pointer;
+            transition: background-color 0.2s;
+        }
+        header button:hover {
+            background-color: #2a622d;
+        }
+        .container {
+            display: flex;
+            flex: 1;
+            overflow: hidden; /* Prevent overall container scroll */
+        }
+        #sidebar {
+            width: 200px;
+            background-color: #e9ecef;
+            padding: 15px;
+            overflow-y: auto;
+            border-right: 1px solid #dee2e6;
+        }
+        #sidebar h3 {
+            margin-top: 0;
+            color: #495057;
+        }
+        #tableList {
+            list-style: none;
+            padding: 0;
+            margin: 0;
+        }
+        #tableList li {
+            padding: 8px 5px;
+            cursor: pointer;
+            border-radius: 4px;
+            margin-bottom: 5px;
+            transition: background-color 0.2s;
+        }
+        #tableList li:hover, #tableList li.active {
+            background-color: #d4dadf;
+        }
+        #mainContent {
+            flex: 1;
+            display: flex;
+            flex-direction: column;
+            padding: 20px;
+            overflow: hidden; /* Prevent main content scroll */
+        }
+        .content-area {
+            flex: 1;
+            display: flex;
+            flex-direction: column;
+            overflow: hidden; /* Child takes scroll */
+        }
+        #schemaDisplay, #dataDisplayContainer, #queryResultContainer {
+            background-color: #fff;
+            border: 1px solid #dee2e6;
+            border-radius: 4px;
+            padding: 15px;
+            margin-bottom: 20px;
+            overflow: auto; /* Allow scrolling within these areas */
+        }
+        #dataDisplayContainer {
+            flex: 1; /* Takes remaining space */
+        }
+        #queryArea {
+            margin-top: auto; /* Push query area to bottom if space */
+            padding-top: 20px;
+            border-top: 1px solid #dee2e6;
+        }
+        #queryArea textarea {
+            width: 100%;
+            min-height: 80px;
+            padding: 10px;
+            border: 1px solid #ccc;
+            border-radius: 4px;
+            box-sizing: border-box;
+            font-family: monospace;
+            margin-bottom: 10px;
+        }
+        #queryArea button {
+            padding: 10px 20px;
+            background-color: #007bff;
+            color: white;
+            border: none;
+            border-radius: 4px;
+            cursor: pointer;
+            transition: background-color 0.2s;
+        }
+        #queryArea button:hover {
+            background-color: #0056b3;
+        }
+        table {
+            width: 100%;
+            border-collapse: collapse;
+            margin-top: 10px;
+        }
+        th, td {
+            border: 1px solid #ddd;
+            padding: 8px;
+            text-align: left;
+            white-space: nowrap;
+        }
+        th {
+            background-color: #f2f2f2;
+            font-weight: bold;
+        }
+        tr:nth-child(even) {
+            background-color: #f9f9f9;
+        }
+        #statusMessage {
+            padding: 10px;
+            margin-top: 10px;
+            border-radius: 4px;
+            display: none; /* Hidden by default */
+        }
+        #statusMessage.success {
+            background-color: #d4edda;
+            color: #155724;
+            border: 1px solid #c3e6cb;
+        }
+        #statusMessage.error {
+            background-color: #f8d7da;
+            color: #721c24;
+            border: 1px solid #f5c6cb;
+        }
+         .loader {
+             border: 4px solid #f3f3f3; /* Light grey */
+             border-top: 4px solid #3498db; /* Blue */
+             border-radius: 50%;
+             width: 20px;
+             height: 20px;
+             animation: spin 1s linear infinite;
+             display: none; /* Hidden by default */
+             margin-left: 10px;
+         }
+         @keyframes spin {
+             0% { transform: rotate(0deg); }
+             100% { transform: rotate(360deg); }
+         }
+    </style>
+</head>
+<body>
+    <header>
+        <label for="apiUrl" style="margin-right: 10px; font-weight: bold; color: white;">API URL:</label>
+        <input type="text" id="apiUrl" value="http://localhost:8000" placeholder="Enter API Base URL">
+        <button id="connectButton">Connect</button>
+        <div class="loader" id="loadingIndicator"></div>
+    </header>
+    <div class="container">
+        <aside id="sidebar">
+            <h3>Tables</h3>
+            <ul id="tableList">
+                <!-- Table names will be loaded here -->
+            </ul>
+        </aside>
+        <main id="mainContent">
+            <div class="content-area">
+                <div id="schemaDisplay">
+                    <h4>Schema</h4>
+                    <p>Select a table from the list to view its schema.</p>
+                    <table id="schemaTable"></table>
+                </div>
+                <div id="dataDisplayContainer">
+                     <h4>Data <span id="tableDataHeader"></span></h4>
+                     <p>Select a table from the list to view its data (limited rows).</p>
+                    <div id="dataDisplay" style="max-height: 300px; overflow-y: auto;">
+                        <table id="dataTable"></table>
+                    </div>
+                </div>
+                <div id="queryResultContainer" style="display: none;">
+                     <h4>Query Result</h4>
+                     <div id="queryResultDisplay" style="max-height: 300px; overflow-y: auto;">
+                        <table id="queryResultTable"></table>
+                    </div>
+                </div>
+            </div>
+            <div id="queryArea">
+                <h4>Custom SQL Query (SELECT only)</h4>
+                <textarea id="sqlInput" placeholder="Enter your SELECT query here..."></textarea>
+                <button id="runSqlButton">Run SQL</button>
+            </div>
+             <div id="statusMessage"></div>
+        </main>
+    </div>
+    <script>
+        const apiUrlInput = document.getElementById('apiUrl');
+        const connectButton = document.getElementById('connectButton');
+        const tableList = document.getElementById('tableList');
+        const schemaDisplay = document.getElementById('schemaDisplay');
+        const schemaTable = document.getElementById('schemaTable');
+        const dataDisplayContainer = document.getElementById('dataDisplayContainer');
+        const dataDisplay = document.getElementById('dataDisplay');
+        const dataTable = document.getElementById('dataTable');
+        const tableDataHeader = document.getElementById('tableDataHeader');
+        const sqlInput = document.getElementById('sqlInput');
+        const runSqlButton = document.getElementById('runSqlButton');
+        const queryResultContainer = document.getElementById('queryResultContainer');
+        const queryResultDisplay = document.getElementById('queryResultDisplay');
+        const queryResultTable = document.getElementById('queryResultTable');
+        const statusMessage = document.getElementById('statusMessage');
+        const loadingIndicator = document.getElementById('loadingIndicator');
+        let API_BASE_URL = '';
+        let currentTables = [];
+        let selectedTable = null;
+        // --- Utility Functions ---
+        function showLoader(show) {
+            loadingIndicator.style.display = show ? 'inline-block' : 'none';
+        }
+        function showStatus(message, isError = false) {
+            statusMessage.textContent = message;
+            statusMessage.className = isError ? 'error' : 'success';
+            statusMessage.style.display = 'block';
+            // Automatically hide after a few seconds
+            setTimeout(() => { statusMessage.style.display = 'none'; }, 5000);
+        }
+        function clearStatus() {
+            statusMessage.textContent = '';
+            statusMessage.style.display = 'none';
+        }
+        async function fetchAPI(endpoint, options = {}) {
+            showLoader(true);
+            clearStatus();
+            const url = `${API_BASE_URL}${endpoint}`;
+            try {
+                const response = await fetch(url, options);
+                if (!response.ok) {
+                    let errorDetail = `HTTP error! status: ${response.status}`;
+                    try {
+                        const errorJson = await response.json();
+                        errorDetail += ` - ${errorJson.detail || JSON.stringify(errorJson)}`;
+                    } catch (e) { /* Ignore if response is not JSON */ }
+                    throw new Error(errorDetail);
+                }
+                // Handle empty responses for non-JSON endpoints if necessary
+                 if (response.headers.get("content-type")?.includes("application/json")) {
+                    return await response.json();
+                 }
+                 return await response.text(); // Or handle other types
+            } catch (error) {
+                console.error('API Fetch Error:', error);
+                showStatus(`Error: ${error.message}`, true);
+                throw error; // Re-throw to stop further processing
+            } finally {
+                showLoader(false);
+            }
+        }
+        function renderTable(data, tableElement) {
+            tableElement.innerHTML = ''; // Clear previous content
+            if (!data || data.length === 0) {
+                tableElement.innerHTML = '<tbody><tr><td>No data available.</td></tr></tbody>';
+                return;
+            }
+            const headers = Object.keys(data[0]);
+            const thead = tableElement.createTHead();
+            const headerRow = thead.insertRow();
+            headers.forEach(headerText => {
+                const th = document.createElement('th');
+                th.textContent = headerText;
+                headerRow.appendChild(th);
+            });
+            const tbody = tableElement.createTBody();
+            data.forEach(rowData => {
+                const row = tbody.insertRow();
+                headers.forEach(header => {
+                    const cell = row.insertCell();
+                    // Handle null or undefined gracefully
+                    cell.textContent = rowData[header] === null || rowData[header] === undefined ? 'NULL' : String(rowData[header]);
+                });
+            });
+        }
+        function renderSchema(schemaData) {
+             const tableElement = schemaTable;
+             tableElement.innerHTML = ''; // Clear previous
+            if (!schemaData || !schemaData.columns || schemaData.columns.length === 0) {
+                schemaDisplay.innerHTML = '<h4>Schema</h4><p>No schema information available.</p>';
+                return;
+            }
+            schemaDisplay.innerHTML = '<h4>Schema</h4>'; // Reset header
+            const thead = tableElement.createTHead();
+            const headerRow = thead.insertRow();
+            ['Name', 'Type'].forEach(headerText => {
+                const th = document.createElement('th');
+                th.textContent = headerText;
+                headerRow.appendChild(th);
+            });
+            const tbody = tableElement.createTBody();
+            schemaData.columns.forEach(column => {
+                const row = tbody.insertRow();
+                row.insertCell().textContent = column.name;
+                row.insertCell().textContent = column.type;
+            });
+         }
+        // --- Event Handlers ---
+        async function loadTables() {
+            API_BASE_URL = apiUrlInput.value.trim().replace(/\/$/, ''); // Remove trailing slash
+            if (!API_BASE_URL) {
+                showStatus("API URL cannot be empty.", true);
+                return;
+            }
+            try {
+                // Optional: Ping root or health endpoint first
+                // await fetchAPI('/');
+                currentTables = await fetchAPI('/tables');
+                displayTables(currentTables);
+                showStatus("Connected. Tables loaded.", false);
+                // Clear previous displays
+                schemaDisplay.innerHTML = '<h4>Schema</h4><p>Select a table from the list.</p>';
+                dataTable.innerHTML = '';
+                tableDataHeader.textContent = '';
+                queryResultContainer.style.display = 'none';
+            } catch (error) {
+                tableList.innerHTML = '<li>Error loading tables.</li>';
+            }
+        }
+        function displayTables(tables) {
+            tableList.innerHTML = ''; // Clear list
+            if (tables.length === 0) {
+                tableList.innerHTML = '<li>No tables found.</li>';
+                return;
+            }
+            tables.sort().forEach(tableName => {
+                const li = document.createElement('li');
+                li.textContent = tableName;
+                li.dataset.tableName = tableName; // Store table name
+                li.onclick = () => handleTableSelection(li);
+                tableList.appendChild(li);
+            });
+        }
+        async function handleTableSelection(listItem) {
+            // Remove active class from previously selected
+            const currentActive = tableList.querySelector('.active');
+            if (currentActive) {
+                currentActive.classList.remove('active');
+            }
+            // Add active class to newly selected
+            listItem.classList.add('active');
+            selectedTable = listItem.dataset.tableName;
+            if (!selectedTable) return;
+            queryResultContainer.style.display = 'none'; // Hide query results
+            dataDisplayContainer.style.display = 'flex'; // Show table data area
+            tableDataHeader.textContent = `for table "${selectedTable}"`;
+             schemaTable.innerHTML = '<tbody><tr><td>Loading schema...</td></tr></tbody>';
+             dataTable.innerHTML = '<tbody><tr><td>Loading data...</td></tr></tbody>';
+            try {
+                const [schemaData, tableData] = await Promise.all([
+                    fetchAPI(`/tables/${selectedTable}/schema`),
+                    fetchAPI(`/tables/${selectedTable}?limit=100`) // Load first 100 rows
+                ]);
+                renderSchema(schemaData);
+                renderTable(tableData, dataTable);
+            } catch (error) {
+                 schemaTable.innerHTML = '<tbody><tr><td>Error loading schema.</td></tr></tbody>';
+                 dataTable.innerHTML = '<tbody><tr><td>Error loading data.</td></tr></tbody>';
+            }
+        }
+         async function runCustomQuery() {
+             const sql = sqlInput.value.trim();
+             if (!sql) {
+                 showStatus("SQL query cannot be empty.", true);
+                 return;
+             }
+             if (!API_BASE_URL) {
+                 showStatus("Connect to the API first (enter URL and press Connect).", true);
+                 return;
+            }
+            dataDisplayContainer.style.display = 'none'; // Hide table data
+            queryResultContainer.style.display = 'block'; // Show query results area
+            queryResultTable.innerHTML = '<tbody><tr><td>Running query...</td></tr></tbody>';
+             try {
+                 const resultData = await fetchAPI('/query', {
+                     method: 'POST',
+                     headers: {
+                         'Content-Type': 'application/json',
+                     },
+                     body: JSON.stringify({ sql: sql }),
+                 });
+                 renderTable(resultData, queryResultTable);
+                 showStatus("Query executed successfully.", false);
+             } catch (error) {
+                queryResultTable.innerHTML = '<tbody><tr><td>Error executing query.</td></tr></tbody>';
+                 // fetchAPI already shows the status
+             }
+         }
+        // --- Initial Setup ---
+        connectButton.onclick = loadTables;
+        runSqlButton.onclick = runCustomQuery;
+        // Optional: Load tables on page load if API URL is preset
+        // if (apiUrlInput.value) {
+        //     loadTables();
+        // }
+    </script>
+</body>
+</html>

main.py CHANGED Viewed

@@ -1,12 +1,14 @@
 import duckdb
 import os
-from fastapi import FastAPI, HTTPException, Request, Path as FastPath
 from fastapi.responses import FileResponse, StreamingResponse
 from pydantic import BaseModel, Field
 from typing import List, Dict, Any, Optional
 import logging
 import io
 import asyncio
 # --- Configuration ---
 DATABASE_PATH = os.environ.get("DUCKDB_PATH", "data/mydatabase.db")
@@ -26,24 +28,24 @@ app = FastAPI(
     version="0.1.0"
 )
-# --- Database Connection ---
-# For simplicity in this example, we connect within each request.
-# For production, consider dependency injection or connection pooling.
-def get_db():
     try:
         # Check if the database file needs initialization
         initialize = not os.path.exists(DATABASE_PATH) or os.path.getsize(DATABASE_PATH) == 0
-        conn = duckdb.connect(DATABASE_PATH, read_only=False)
         if initialize:
             logger.info(f"Database file not found or empty at {DATABASE_PATH}. Initializing.")
-            # You could add initial schema setup here if needed
-            # conn.execute("CREATE TABLE IF NOT EXISTS initial_table (id INTEGER, name VARCHAR);")
         yield conn
     except duckdb.Error as e:
         logger.error(f"Database connection error: {e}")
         raise HTTPException(status_code=500, detail=f"Database connection error: {e}")
     finally:
-        if 'conn' in locals() and conn:
             conn.close()
 # --- Pydantic Models ---
@@ -51,19 +53,24 @@ class ColumnDefinition(BaseModel):
     name: str
     type: str
 class CreateTableRequest(BaseModel):
     columns: List[ColumnDefinition]
 class CreateRowRequest(BaseModel):
-    # List of rows, where each row is a dict of column_name: value
     rows: List[Dict[str, Any]]
 class UpdateRowRequest(BaseModel):
-    updates: Dict[str, Any] # Column value pairs to set
-    condition: str         # SQL WHERE clause string to identify rows
 class DeleteRowRequest(BaseModel):
-    condition: str         # SQL WHERE clause string to identify rows
 class ApiResponse(BaseModel):
     message: str
@@ -71,35 +78,53 @@ class ApiResponse(BaseModel):
 # --- Helper Functions ---
 def safe_identifier(name: str) -> str:
-    """Quotes an identifier safely."""
-    if not name.isidentifier():
-        # Basic check, consider more robust validation/sanitization if needed
-         # Use DuckDB's quoting
-        try:
-            conn = duckdb.connect(':memory:')
-            quoted = conn.execute(f"SELECT '{name}'::IDENTIFIER").fetchone()[0]
-            conn.close()
-            return quoted
-        except duckdb.Error:
-             raise HTTPException(status_code=400, detail=f"Invalid identifier: {name}")
-    # Also quote standard identifiers to be safe
-    return f'"{name}"'
 def generate_column_sql(columns: List[ColumnDefinition]) -> str:
     """Generates the column definition part of a CREATE TABLE statement."""
     defs = []
     for col in columns:
         col_name_safe = safe_identifier(col.name)
-        # Basic type validation (can be expanded)
-        allowed_types = ['INTEGER', 'VARCHAR', 'TEXT', 'BOOLEAN', 'FLOAT', 'DOUBLE', 'DATE', 'TIMESTAMP', 'BLOB', 'BIGINT', 'DECIMAL']
         type_upper = col.type.strip().upper()
-        # Allow DECIMAL(p,s) syntax
-        if not (type_upper.startswith('DECIMAL(') and type_upper.endswith(')')) and \
-           not any(base_type in type_upper for base_type in allowed_types):
-             raise HTTPException(status_code=400, detail=f"Unsupported or invalid data type: {col.type}")
-        defs.append(f"{col_name_safe} {col.type}")
     return ", ".join(defs)
 # --- API Endpoints ---
 @app.get("/", summary="API Root", response_model=ApiResponse)
@@ -107,6 +132,74 @@ async def read_root():
     """Provides a welcome message for the API."""
     return {"message": "Welcome to the DuckDB API!"}
 @app.post("/tables/{table_name}", summary="Create Table", response_model=ApiResponse, status_code=201)
 async def create_table(
     table_name: str = FastPath(..., description="Name of the table to create"),
@@ -119,12 +212,12 @@ async def create_table(
     try:
         columns_sql = generate_column_sql(schema.columns)
-        sql = f"CREATE TABLE {table_name_safe} ({columns_sql});"
         logger.info(f"Executing SQL: {sql}")
-        for conn in get_db():
             conn.execute(sql)
-        return {"message": f"Table '{table_name}' created successfully."}
-    except HTTPException as e: # Re-raise validation errors
         raise e
     except duckdb.Error as e:
         logger.error(f"Error creating table '{table_name}': {e}")
@@ -136,28 +229,27 @@ async def create_table(
 @app.get("/tables/{table_name}", summary="Read Table Data")
 async def read_table(
     table_name: str = FastPath(..., description="Name of the table to read from"),
-    limit: Optional[int] = None,
-    offset: Optional[int] = None
 ):
-    """Reads and returns all rows from a specified table. Supports limit and offset."""
     table_name_safe = safe_identifier(table_name)
     sql = f"SELECT * FROM {table_name_safe}"
     params = []
-    if limit is not None:
         sql += " LIMIT ?"
         params.append(limit)
-    if offset is not None:
         sql += " OFFSET ?"
         params.append(offset)
     sql += ";"
     try:
         logger.info(f"Executing SQL: {sql} with params: {params}")
-        for conn in get_db():
-            result = conn.execute(sql, params).fetchall()
-            # Convert rows to dictionaries for JSON serialization
-            column_names = [desc[0] for desc in conn.description]
-            data = [dict(zip(column_names, row)) for row in result]
         return data
     except duckdb.CatalogException as e:
          raise HTTPException(status_code=404, detail=f"Table '{table_name}' not found.")
@@ -196,16 +288,14 @@ async def create_rows(
     try:
         logger.info(f"Executing SQL: {sql} for {len(params_list)} rows")
-        for conn in get_db():
             conn.executemany(sql, params_list)
-            conn.commit() # Explicit commit after potential bulk insert
         return {"message": f"Successfully inserted {len(params_list)} rows into '{table_name}'."}
     except duckdb.CatalogException as e:
          raise HTTPException(status_code=404, detail=f"Table '{table_name}' not found.")
     except duckdb.Error as e:
         logger.error(f"Error inserting rows into '{table_name}': {e}")
-        # Rollback on error might be needed depending on transaction behavior
-        # For get_db creating connection per request, this is less critical
         raise HTTPException(status_code=400, detail=f"Error inserting rows: {e}")
     except Exception as e:
         logger.error(f"Unexpected error inserting rows into '{table_name}': {e}")
@@ -232,15 +322,14 @@ async def update_rows(
     set_sql = ", ".join(set_clauses)
     # WARNING: Injecting request.condition directly is a security risk.
-    # In a real app, use query parameters or a safer way to build the WHERE clause.
-    sql = f"UPDATE {table_name_safe} SET {set_sql} WHERE {request.condition};"
     try:
         logger.info(f"Executing SQL: {sql} with params: {params}")
-        for conn in get_db():
-            # Use execute for safety with parameters
             conn.execute(sql, params)
-            conn.commit()
         return {"message": f"Rows in '{table_name}' updated successfully based on condition."}
     except duckdb.CatalogException as e:
          raise HTTPException(status_code=404, detail=f"Table '{table_name}' not found.")
@@ -262,15 +351,13 @@ async def delete_rows(
          raise HTTPException(status_code=400, detail="Delete condition (WHERE clause) is required.")
     # WARNING: Injecting request.condition directly is a security risk.
-    # In a real app, use query parameters or a safer way to build the WHERE clause.
-    sql = f"DELETE FROM {table_name_safe} WHERE {request.condition};"
     try:
         logger.info(f"Executing SQL: {sql}")
-        for conn in get_db():
-            # Execute does not directly support parameters for WHERE in DELETE like this easily
             conn.execute(sql)
-            conn.commit()
         return {"message": f"Rows from '{table_name}' deleted successfully based on condition."}
     except duckdb.CatalogException as e:
          raise HTTPException(status_code=404, detail=f"Table '{table_name}' not found.")
@@ -282,49 +369,38 @@ async def delete_rows(
         raise HTTPException(status_code=500, detail="An unexpected error occurred.")
 # --- Download Endpoints ---
 @app.get("/download/table/{table_name}", summary="Download Table as CSV")
 async def download_table_csv(
     table_name: str = FastPath(..., description="Name of the table to download")
 ):
     """Downloads the entire content of a table as a CSV file."""
     table_name_safe = safe_identifier(table_name)
-    # Use COPY TO STDOUT for efficient streaming
     sql = f"COPY (SELECT * FROM {table_name_safe}) TO STDOUT (FORMAT CSV, HEADER)"
     async def stream_csv_data():
-        # We need a non-blocking way to stream data from DuckDB.
-        # DuckDB's Python API is blocking. A simple approach for this demo
-        # is to fetch all data first, then stream it.
-        # A more advanced approach would involve running the DuckDB query
-        # in a separate thread or process pool managed by asyncio.
         try:
             all_data_io = io.StringIO()
-            # This COPY TO variant isn't directly available in Python API for streaming to a buffer easily.
-            # Let's fetch data and format as CSV manually or use Pandas.
-            for conn in get_db():
-                df = conn.execute(f"SELECT * FROM {table_name_safe}").df() # Use pandas for CSV conversion
-            # Use an in-memory text buffer
             df.to_csv(all_data_io, index=False)
             all_data_io.seek(0)
-            # Stream the content chunk by chunk
             chunk_size = 8192
             while True:
                 chunk = all_data_io.read(chunk_size)
                 if not chunk:
                     break
-                yield chunk
-                # Allow other tasks to run
                 await asyncio.sleep(0)
             all_data_io.close()
         except duckdb.CatalogException as e:
-            # Stream an error message if the table doesn't exist
             yield f"Error: Table '{table_name}' not found.".encode('utf-8')
-            logger.error(f"Error downloading table '{table_name}': {e}")
         except duckdb.Error as e:
             yield f"Error: Could not export table '{table_name}'. {e}".encode('utf-8')
             logger.error(f"Error downloading table '{table_name}': {e}")
@@ -332,7 +408,6 @@ async def download_table_csv(
              yield f"Error: An unexpected error occurred.".encode('utf-8')
              logger.error(f"Unexpected error downloading table '{table_name}': {e}")
     return StreamingResponse(
         stream_csv_data(),
         media_type="text/csv",
@@ -345,16 +420,11 @@ async def download_database_file():
     """Downloads the entire DuckDB database file."""
     if not os.path.exists(DATABASE_PATH):
         raise HTTPException(status_code=404, detail="Database file not found.")
-    # Ensure connections are closed before downloading to avoid partial writes/locking issues.
-    # This is tricky with the current get_db pattern. A proper app stop/start or
-    # dedicated maintenance mode would be better. For this demo, we hope for the best.
     logger.warning("Attempting to download database file. Ensure no active writes are occurring.")
     return FileResponse(
         path=DATABASE_PATH,
         filename=os.path.basename(DATABASE_PATH),
-        media_type="application/octet-stream" # Generic binary file type
     )
@@ -363,20 +433,9 @@ async def download_database_file():
 async def health_check():
     """Checks if the API and database connection are working."""
     try:
-        for conn in get_db():
             conn.execute("SELECT 1")
         return {"message": "API is healthy and database connection is successful."}
     except Exception as e:
         logger.error(f"Health check failed: {e}")
-        raise HTTPException(status_code=503, detail=f"Health check failed: {e}")
-# --- Optional: Add Startup/Shutdown events if needed ---
-# @app.on_event("startup")
-# async def startup_event():
-#     # Initialize database connection pool, etc.
-#     logger.info("Application startup.")
-# @app.on_event("shutdown")
-# async def shutdown_event():
-#     # Clean up resources, close connections, etc.
-#     logger.info("Application shutdown.")

+# ... (keep existing imports and setup) ...
 import duckdb
 import os
+from fastapi import FastAPI, HTTPException, Request, Path as FastPath, Body
 from fastapi.responses import FileResponse, StreamingResponse
 from pydantic import BaseModel, Field
 from typing import List, Dict, Any, Optional
 import logging
 import io
 import asyncio
+from contextlib import contextmanager # <--- Add contextlib
 # --- Configuration ---
 DATABASE_PATH = os.environ.get("DUCKDB_PATH", "data/mydatabase.db")
     version="0.1.0"
 )
+# --- Database Connection (using context manager for safety) ---
+@contextmanager
+def get_db_context():
+    conn = None
     try:
         # Check if the database file needs initialization
         initialize = not os.path.exists(DATABASE_PATH) or os.path.getsize(DATABASE_PATH) == 0
+        conn = duckdb.connect(DATABASE_PATH, read_only=False) # Allow writes for setup
         if initialize:
             logger.info(f"Database file not found or empty at {DATABASE_PATH}. Initializing.")
+            # Optionally create a default table if the DB is new
+            # conn.execute("CREATE TABLE IF NOT EXISTS example (id INTEGER, name VARCHAR);")
         yield conn
     except duckdb.Error as e:
         logger.error(f"Database connection error: {e}")
         raise HTTPException(status_code=500, detail=f"Database connection error: {e}")
     finally:
+        if conn:
             conn.close()
 # --- Pydantic Models ---
     name: str
     type: str
+class TableSchemaResponse(BaseModel):
+    columns: List[ColumnDefinition]
 class CreateTableRequest(BaseModel):
     columns: List[ColumnDefinition]
 class CreateRowRequest(BaseModel):
     rows: List[Dict[str, Any]]
 class UpdateRowRequest(BaseModel):
+    updates: Dict[str, Any]
+    condition: str
 class DeleteRowRequest(BaseModel):
+    condition: str
+class SQLQueryRequest(BaseModel):
+    sql: str
 class ApiResponse(BaseModel):
     message: str
 # --- Helper Functions ---
 def safe_identifier(name: str) -> str:
+    """Quotes an identifier safely using DuckDB."""
+    # Basic check
+    if not name or not isinstance(name, str):
+         raise HTTPException(status_code=400, detail=f"Invalid identifier provided: {name}")
+    # Use DuckDB's quoting mechanism
+    try:
+        # Use a temporary in-memory connection for quoting safely
+        with duckdb.connect(':memory:') as temp_conn:
+            # Use sql() which returns a relation, then fetch the result
+            quoted = temp_conn.sql(f"SELECT '{name}'::IDENTIFIER").fetchone()
+            if quoted:
+                return quoted[0]
+            else:
+                 raise HTTPException(status_code=500, detail="Failed to quote identifier")
+    except duckdb.Error as e:
+         logger.error(f"Error quoting identifier '{name}': {e}")
+         # Fallback or re-raise depending on policy, here we raise
+         raise HTTPException(status_code=400, detail=f"Invalid identifier '{name}': {e}")
 def generate_column_sql(columns: List[ColumnDefinition]) -> str:
     """Generates the column definition part of a CREATE TABLE statement."""
     defs = []
     for col in columns:
         col_name_safe = safe_identifier(col.name)
+        # More robust type validation needed for production
+        allowed_types_prefix = ['INTEGER', 'VARCHAR', 'TEXT', 'BOOLEAN', 'FLOAT', 'DOUBLE', 'DATE', 'TIMESTAMP', 'BLOB', 'BIGINT', 'DECIMAL', 'LIST', 'STRUCT', 'MAP', 'UNION']
         type_upper = col.type.strip().upper()
+        is_allowed = False
+        for prefix in allowed_types_prefix:
+             # Allow types like VARCHAR(255), DECIMAL(10,2), LIST<INT>, STRUCT<a INT> etc.
+            if type_upper.startswith(prefix):
+                is_allowed = True
+                break
+        if not is_allowed:
+            # Very basic check, expand as needed
+             raise HTTPException(status_code=400, detail=f"Unsupported or potentially invalid data type: {col.type}")
+        defs.append(f"{col_name_safe} {col.type}") # Pass type string directly
     return ", ".join(defs)
+def result_to_dict(cursor_description, rows):
+    """Converts cursor results (description + rows) to a list of dictionaries."""
+    column_names = [desc[0] for desc in cursor_description]
+    return [dict(zip(column_names, row)) for row in rows]
 # --- API Endpoints ---
 @app.get("/", summary="API Root", response_model=ApiResponse)
     """Provides a welcome message for the API."""
     return {"message": "Welcome to the DuckDB API!"}
+# --- NEW ENDPOINT ---
+@app.get("/tables", summary="List Tables", response_model=List[str])
+async def list_tables():
+    """Lists all tables in the default schema."""
+    try:
+        with get_db_context() as conn:
+            # Show user tables (excluding system tables)
+            tables = conn.execute("SELECT table_name FROM information_schema.tables WHERE table_schema = 'main'").fetchall()
+            return [table[0] for table in tables]
+    except duckdb.Error as e:
+        logger.error(f"Error listing tables: {e}")
+        raise HTTPException(status_code=500, detail=f"Error listing tables: {e}")
+# --- NEW ENDPOINT ---
+@app.get("/tables/{table_name}/schema", summary="Get Table Schema", response_model=TableSchemaResponse)
+async def get_table_schema(
+    table_name: str = FastPath(..., description="Name of the table")
+):
+    """Gets the schema (column names and types) for a specific table."""
+    table_name_safe = safe_identifier(table_name)
+    # Use PRAGMA for schema info
+    sql = f"PRAGMA table_info({table_name_safe});"
+    try:
+        with get_db_context() as conn:
+            result = conn.execute(sql).fetchall()
+            if not result:
+                raise HTTPException(status_code=404, detail=f"Table '{table_name}' not found or has no columns.")
+            # PRAGMA table_info columns: cid, name, type, notnull, dflt_value, pk
+            columns = [ColumnDefinition(name=row[1], type=row[2]) for row in result]
+        return TableSchemaResponse(columns=columns)
+    except duckdb.CatalogException as e:
+         raise HTTPException(status_code=404, detail=f"Table '{table_name}' not found.")
+    except duckdb.Error as e:
+        logger.error(f"Error getting schema for table '{table_name}': {e}")
+        raise HTTPException(status_code=400, detail=f"Error getting table schema: {e}")
+# --- NEW ENDPOINT ---
+@app.post("/query", summary="Execute Read-Only SQL Query")
+async def execute_query(query_request: SQLQueryRequest):
+    """Executes a provided SQL query (read-only enforced)."""
+    sql = query_request.sql.strip()
+    # **Security:** Basic check to prevent modification queries.
+    # This is NOT foolproof. A robust solution needs proper SQL parsing or
+    # database roles/permissions restricting the API user.
+    forbidden_keywords = ['INSERT', 'UPDATE', 'DELETE', 'DROP', 'CREATE', 'ALTER', 'ATTACH', 'DETACH', 'COPY', 'EXPORT', 'IMPORT']
+    sql_upper = sql.upper()
+    if any(keyword in sql_upper for keyword in forbidden_keywords):
+         raise HTTPException(status_code=403, detail="Only SELECT queries are allowed.")
+    if not sql_upper.startswith('SELECT') and not sql_upper.startswith('WITH'):
+         raise HTTPException(status_code=400, detail="Query must start with SELECT or WITH.")
+    try:
+        logger.info(f"Executing user SQL: {sql}")
+        with get_db_context() as conn:
+            description = conn.execute(sql).description
+            result = conn.fetchall()
+            # Convert rows to dictionaries for JSON serialization
+            data = result_to_dict(description, result)
+        return data
+    except duckdb.Error as e:
+        logger.error(f"Error executing user query: {e}")
+        raise HTTPException(status_code=400, detail=f"Error executing query: {e}")
+    except Exception as e:
+        logger.error(f"Unexpected error executing user query: {e}")
+        raise HTTPException(status_code=500, detail="An unexpected error occurred during query execution.")
+# --- Existing Endpoints (Keep or adapt as needed) ---
 @app.post("/tables/{table_name}", summary="Create Table", response_model=ApiResponse, status_code=201)
 async def create_table(
     table_name: str = FastPath(..., description="Name of the table to create"),
     try:
         columns_sql = generate_column_sql(schema.columns)
+        sql = f"CREATE OR REPLACE TABLE {table_name_safe} ({columns_sql});" # Use CREATE OR REPLACE for simplicity
         logger.info(f"Executing SQL: {sql}")
+        with get_db_context() as conn:
             conn.execute(sql)
+        return {"message": f"Table '{table_name}' created or replaced successfully."}
+    except HTTPException as e:
         raise e
     except duckdb.Error as e:
         logger.error(f"Error creating table '{table_name}': {e}")
 @app.get("/tables/{table_name}", summary="Read Table Data")
 async def read_table(
     table_name: str = FastPath(..., description="Name of the table to read from"),
+    limit: Optional[int] = 100, # Default limit
+    offset: Optional[int] = 0 # Default offset
 ):
+    """Reads and returns rows from a specified table. Supports limit and offset."""
     table_name_safe = safe_identifier(table_name)
     sql = f"SELECT * FROM {table_name_safe}"
     params = []
+    if limit is not None and limit >= 0:
         sql += " LIMIT ?"
         params.append(limit)
+    if offset is not None and offset >= 0:
         sql += " OFFSET ?"
         params.append(offset)
     sql += ";"
     try:
         logger.info(f"Executing SQL: {sql} with params: {params}")
+        with get_db_context() as conn:
+            description = conn.execute(sql, params).description
+            result = conn.fetchall()
+            data = result_to_dict(description, result)
         return data
     except duckdb.CatalogException as e:
          raise HTTPException(status_code=404, detail=f"Table '{table_name}' not found.")
     try:
         logger.info(f"Executing SQL: {sql} for {len(params_list)} rows")
+        with get_db_context() as conn:
             conn.executemany(sql, params_list)
+            # Removed commit - context manager handles it
         return {"message": f"Successfully inserted {len(params_list)} rows into '{table_name}'."}
     except duckdb.CatalogException as e:
          raise HTTPException(status_code=404, detail=f"Table '{table_name}' not found.")
     except duckdb.Error as e:
         logger.error(f"Error inserting rows into '{table_name}': {e}")
         raise HTTPException(status_code=400, detail=f"Error inserting rows: {e}")
     except Exception as e:
         logger.error(f"Unexpected error inserting rows into '{table_name}': {e}")
     set_sql = ", ".join(set_clauses)
     # WARNING: Injecting request.condition directly is a security risk.
+    # Use parameters for values, but condition structure still needs care.
+    sql = f"UPDATE {table_name_safe} SET {set_sql} WHERE {request.condition};" # Condition not parameterized here
     try:
         logger.info(f"Executing SQL: {sql} with params: {params}")
+        with get_db_context() as conn:
             conn.execute(sql, params)
+            # Removed commit
         return {"message": f"Rows in '{table_name}' updated successfully based on condition."}
     except duckdb.CatalogException as e:
          raise HTTPException(status_code=404, detail=f"Table '{table_name}' not found.")
          raise HTTPException(status_code=400, detail="Delete condition (WHERE clause) is required.")
     # WARNING: Injecting request.condition directly is a security risk.
+    sql = f"DELETE FROM {table_name_safe} WHERE {request.condition};" # Condition not parameterized here
     try:
         logger.info(f"Executing SQL: {sql}")
+        with get_db_context() as conn:
             conn.execute(sql)
+            # Removed commit
         return {"message": f"Rows from '{table_name}' deleted successfully based on condition."}
     except duckdb.CatalogException as e:
          raise HTTPException(status_code=404, detail=f"Table '{table_name}' not found.")
         raise HTTPException(status_code=500, detail="An unexpected error occurred.")
 # --- Download Endpoints ---
 @app.get("/download/table/{table_name}", summary="Download Table as CSV")
 async def download_table_csv(
     table_name: str = FastPath(..., description="Name of the table to download")
 ):
     """Downloads the entire content of a table as a CSV file."""
     table_name_safe = safe_identifier(table_name)
     sql = f"COPY (SELECT * FROM {table_name_safe}) TO STDOUT (FORMAT CSV, HEADER)"
     async def stream_csv_data():
         try:
+            # Use pandas for CSV conversion in-memory
+            with get_db_context() as conn:
+                # Check if table exists before fetching
+                conn.execute(f"SELECT 1 FROM {table_name_safe} LIMIT 0")
+                df = conn.execute(f"SELECT * FROM {table_name_safe}").df()
             all_data_io = io.StringIO()
             df.to_csv(all_data_io, index=False)
             all_data_io.seek(0)
             chunk_size = 8192
             while True:
                 chunk = all_data_io.read(chunk_size)
                 if not chunk:
                     break
+                yield chunk.encode('utf-8') # Encode to bytes for streaming response
                 await asyncio.sleep(0)
             all_data_io.close()
         except duckdb.CatalogException as e:
             yield f"Error: Table '{table_name}' not found.".encode('utf-8')
+            logger.error(f"Error downloading table '{table_name}': Table not found.")
         except duckdb.Error as e:
             yield f"Error: Could not export table '{table_name}'. {e}".encode('utf-8')
             logger.error(f"Error downloading table '{table_name}': {e}")
              yield f"Error: An unexpected error occurred.".encode('utf-8')
              logger.error(f"Unexpected error downloading table '{table_name}': {e}")
     return StreamingResponse(
         stream_csv_data(),
         media_type="text/csv",
     """Downloads the entire DuckDB database file."""
     if not os.path.exists(DATABASE_PATH):
         raise HTTPException(status_code=404, detail="Database file not found.")
     logger.warning("Attempting to download database file. Ensure no active writes are occurring.")
     return FileResponse(
         path=DATABASE_PATH,
         filename=os.path.basename(DATABASE_PATH),
+        media_type="application/vnd.duckdb.database" # More specific media type
     )
 async def health_check():
     """Checks if the API and database connection are working."""
     try:
+        with get_db_context() as conn:
             conn.execute("SELECT 1")
         return {"message": "API is healthy and database connection is successful."}
     except Exception as e:
         logger.error(f"Health check failed: {e}")
+        raise HTTPException(status_code=503, detail=f"Health check failed: {e}")