pizzasql/clients/python/pizzasql.py

"""
PizzaSQL Client for Python

A simple, Pythonic client for PizzaSQL.
"""

from typing import Any, Dict, List, Optional, Union
from dataclasses import dataclass
from urllib.parse import urlparse
import json

try:
    import httpx
    _client_class = httpx.Client
    _async_client_class = httpx.AsyncClient
except ImportError:
    import urllib.request
    import urllib.error
    _client_class = None
    _async_client_class = None


@dataclass
class Column:
    """Represents a column in a query result."""
    name: str
    type: str


@dataclass
class QueryResult:
    """Result of a SQL query."""
    columns: List[Column]
    rows: List[Dict[str, Any]]
    rows_affected: int
    last_insert_id: int
    execution_time: str

    def __iter__(self):
        return iter(self.rows)

    def __len__(self):
        return len(self.rows)

    def __getitem__(self, index):
        return self.rows[index]


class PizzaSQLError(Exception):
    """Exception raised for PizzaSQL errors."""
    def __init__(self, code: str, message: str, details: Optional[Dict] = None):
        self.code = code
        self.message = message
        self.details = details
        super().__init__(f"[{code}] {message}")


class PizzaSQL:
    """
    PizzaSQL client for Python.

    Usage:
        db = PizzaSQL('http://localhost:8080/mydb', api_key='your-key')
        rows = db.sql('SELECT * FROM users')
    """

    def __init__(
        self,
        uri: str,
        api_key: Optional[str] = None,
        timeout: float = 30.0
    ):
        """
        Create a new PizzaSQL connection.

        Args:
            uri: Database URI (e.g., 'http://localhost:8080/mydb')
            api_key: Optional API key for authentication
            timeout: Request timeout in seconds
        """
        parsed = urlparse(uri)
        self._base_url = f"{parsed.scheme}://{parsed.netloc}"
        self._database = parsed.path.lstrip('/') or None
        self._api_key = api_key
        self._timeout = timeout

        if _client_class:
            self._client = _client_class(timeout=timeout)
        else:
            self._client = None

    def _headers(self) -> Dict[str, str]:
        headers = {'Content-Type': 'application/json'}
        if self._api_key:
            headers['Authorization'] = f'Bearer {self._api_key}'
        if self._database:
            headers['X-Database'] = self._database
        return headers

    def _request(self, method: str, path: str, data: Optional[Dict] = None) -> Dict:
        url = f"{self._base_url}{path}"
        headers = self._headers()

        if self._client:
            # Use httpx
            if method == 'GET':
                response = self._client.get(url, headers=headers)
            else:
                response = self._client.post(url, headers=headers, json=data)

            result = response.json()
            if response.status_code >= 400:
                error = result.get('error', {})
                raise PizzaSQLError(
                    error.get('code', 'UNKNOWN'),
                    error.get('message', 'Unknown error'),
                    error.get('details')
                )
            return result
        else:
            # Fallback to urllib
            req = urllib.request.Request(url, headers=headers)
            if data:
                req.data = json.dumps(data).encode('utf-8')

            try:
                with urllib.request.urlopen(req, timeout=self._timeout) as response:
                    return json.loads(response.read().decode('utf-8'))
            except urllib.error.HTTPError as e:
                result = json.loads(e.read().decode('utf-8'))
                error = result.get('error', {})
                raise PizzaSQLError(
                    error.get('code', 'UNKNOWN'),
                    error.get('message', str(e)),
                    error.get('details')
                )

    def _transform_rows(self, columns: List[Dict], rows: List[List]) -> List[Dict[str, Any]]:
        """Transform array rows to dictionaries."""
        return [
            {col['name']: row[i] for i, col in enumerate(columns)}
            for row in rows
        ]

    def query(self, sql: str, params: Optional[List] = None) -> QueryResult:
        """
        Execute a SQL query and return full result.

        Args:
            sql: SQL query string
            params: Optional list of parameters

        Returns:
            QueryResult with columns, rows, and metadata
        """
        result = self._request('POST', '/query', {
            'sql': sql,
            'params': params or []
        })

        columns = [Column(**col) for col in result.get('columns', [])]
        rows = self._transform_rows(result.get('columns', []), result.get('rows', []))

        return QueryResult(
            columns=columns,
            rows=rows,
            rows_affected=result.get('rowsAffected', 0),
            last_insert_id=result.get('lastInsertId', 0),
            execution_time=result.get('executionTime', '')
        )

    def sql(self, sql: str, params: Optional[List] = None) -> List[Dict[str, Any]]:
        """
        Execute a SQL query and return rows.

        Args:
            sql: SQL query string
            params: Optional list of parameters

        Returns:
            List of row dictionaries
        """
        return self.query(sql, params).rows

    def execute(
        self,
        statements: List[Dict[str, Any]],
        transaction: bool = True
    ) -> Dict[str, Any]:
        """
        Execute multiple statements in a batch.

        Args:
            statements: List of {'sql': ..., 'params': [...]} dicts
            transaction: Whether to wrap in a transaction

        Returns:
            Execution result with affected rows
        """
        return self._request('POST', '/execute', {
            'statements': [
                {'sql': s['sql'], 'params': s.get('params', [])}
                for s in statements
            ],
            'transaction': transaction
        })

    def tables(self) -> List[str]:
        """
        List all tables in the database.

        Returns:
            List of table names
        """
        result = self._request('GET', '/schema/tables')
        return result.get('tables', [])

    def schema(self, table_name: str) -> List[Column]:
        """
        Get schema for a specific table.

        Args:
            table_name: Name of the table

        Returns:
            List of Column objects
        """
        result = self._request('GET', f'/schema/tables/{table_name}')
        return [Column(**col) for col in result.get('columns', [])]

    def health(self) -> Dict[str, str]:
        """
        Check database health.

        Returns:
            Health status dict
        """
        return self._request('GET', '/health')

    def use(self, database: str) -> 'PizzaSQL':
        """
        Create a new client for a different database.

        Args:
            database: Database name

        Returns:
            New PizzaSQL client
        """
        client = PizzaSQL(self._base_url, self._api_key, self._timeout)
        client._database = database
        return client

    def close(self):
        """Close the underlying HTTP client."""
        if self._client and hasattr(self._client, 'close'):
            self._client.close()

    def __enter__(self):
        return self

    def __exit__(self, *args):
        self.close()


# Convenience function
def connect(uri: str, api_key: Optional[str] = None) -> PizzaSQL:
    """
    Create a new PizzaSQL connection.

    Args:
        uri: Database URI (e.g., 'http://localhost:8080/mydb')
        api_key: Optional API key for authentication

    Returns:
        PizzaSQL client instance
    """
    return PizzaSQL(uri, api_key)