main.py

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")
DATA_DIR = "data"

# Ensure data directory exists
os.makedirs(DATA_DIR, exist_ok=True)

# --- Logging ---
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

# --- FastAPI App ---
app = FastAPI(
    title="DuckDB API",
    description="An API to interact with a DuckDB database.",
    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 ---
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
    details: Optional[Any] = None

# --- 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)
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"),
    schema: CreateTableRequest = ...,
):
    """Creates a new table with the specified schema."""
    table_name_safe = safe_identifier(table_name)
    if not schema.columns:
        raise HTTPException(status_code=400, detail="Table must have at least one column.")

    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}")
        raise HTTPException(status_code=400, detail=f"Error creating table: {e}")
    except Exception as e:
        logger.error(f"Unexpected error creating table '{table_name}': {e}")
        raise HTTPException(status_code=500, detail="An unexpected error occurred.")

@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.")
    except duckdb.Error as e:
        logger.error(f"Error reading table '{table_name}': {e}")
        raise HTTPException(status_code=400, detail=f"Error reading table: {e}")
    except Exception as e:
        logger.error(f"Unexpected error reading table '{table_name}': {e}")
        raise HTTPException(status_code=500, detail="An unexpected error occurred.")


@app.post("/tables/{table_name}/rows", summary="Create Rows", response_model=ApiResponse, status_code=201)
async def create_rows(
    table_name: str = FastPath(..., description="Name of the table to insert into"),
    request: CreateRowRequest = ...,
):
    """Inserts one or more rows into the specified table."""
    table_name_safe = safe_identifier(table_name)
    if not request.rows:
        raise HTTPException(status_code=400, detail="No rows provided to insert.")

    # Assume all rows have the same columns based on the first row
    columns = list(request.rows[0].keys())
    columns_safe = [safe_identifier(col) for col in columns]
    placeholders = ", ".join(["?"] * len(columns))
    columns_sql = ", ".join(columns_safe)

    sql = f"INSERT INTO {table_name_safe} ({columns_sql}) VALUES ({placeholders});"

    # Convert list of dicts to list of lists/tuples for executemany
    params_list = []
    for row_dict in request.rows:
        if list(row_dict.keys()) != columns:
             raise HTTPException(status_code=400, detail="All rows must have the same columns in the same order.")
        params_list.append(list(row_dict.values()))

    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}")
        raise HTTPException(status_code=500, detail="An unexpected error occurred.")


@app.put("/tables/{table_name}/rows", summary="Update Rows", response_model=ApiResponse)
async def update_rows(
    table_name: str = FastPath(..., description="Name of the table to update"),
    request: UpdateRowRequest = ...,
):
    """Updates rows in the table based on a condition."""
    table_name_safe = safe_identifier(table_name)
    if not request.updates:
        raise HTTPException(status_code=400, detail="No updates provided.")
    if not request.condition:
         raise HTTPException(status_code=400, detail="Update condition (WHERE clause) is required.")

    set_clauses = []
    params = []
    for col, value in request.updates.items():
        set_clauses.append(f"{safe_identifier(col)} = ?")
        params.append(value)

    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.")
    except duckdb.Error as e:
        logger.error(f"Error updating rows in '{table_name}': {e}")
        raise HTTPException(status_code=400, detail=f"Error updating rows: {e}")
    except Exception as e:
        logger.error(f"Unexpected error updating rows in '{table_name}': {e}")
        raise HTTPException(status_code=500, detail="An unexpected error occurred.")

@app.delete("/tables/{table_name}/rows", summary="Delete Rows", response_model=ApiResponse)
async def delete_rows(
    table_name: str = FastPath(..., description="Name of the table to delete from"),
    request: DeleteRowRequest = ...,
):
    """Deletes rows from the table based on a condition."""
    table_name_safe = safe_identifier(table_name)
    if not request.condition:
         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.")
    except duckdb.Error as e:
        logger.error(f"Error deleting rows from '{table_name}': {e}")
        raise HTTPException(status_code=400, detail=f"Error deleting rows: {e}")
    except Exception as e:
        logger.error(f"Unexpected error deleting rows from '{table_name}': {e}")
        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}")
        except Exception as 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",
        headers={"Content-Disposition": f"attachment; filename={table_name}.csv"},
    )


@app.get("/download/database", summary="Download Database File")
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
    )


# --- Health Check ---
@app.get("/health", summary="Health Check", response_model=ApiResponse)
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.")