bigquery_client.py

"""
BigQuery Client Wrapper
Provides simplified query interface with connection pooling
"""

import logging
from typing import Any, Dict, List, Optional

from google.cloud import bigquery

logger = logging.getLogger(__name__)


class BigQueryClient:
    def __init__(self, project_id: str, dataset_id: str):
        self.project_id = project_id
        self.dataset_id = dataset_id
        self.client = bigquery.Client(project=project_id)
        logger.info(f"BigQuery client initialized: {project_id}.{dataset_id}")

    def _format_query(self, query: str) -> str:
        """Replace {dataset} placeholder with actual dataset ID"""
        return query.format(dataset=f"{self.project_id}.{self.dataset_id}")

    def fetch_one(self, query: str) -> Optional[Dict[str, Any]]:
        """
        Execute query and return first row as dict.
        Returns None if no results.
        """
        formatted_query = self._format_query(query)
        logger.debug(f"Executing query: {formatted_query[:100]}...")

        try:
            result = self.client.query(formatted_query).result()
            rows = list(result)

            if not rows:
                return None

            # Convert Row to dict
            row = rows[0]
            return {key: row[key] for key in row.keys()}

        except Exception as e:
            logger.error(f"Query failed: {e}")
            raise

    def fetch_all(self, query: str) -> List[Dict[str, Any]]:
        """
        Execute query and return all rows as list of dicts.
        """
        formatted_query = self._format_query(query)
        logger.debug(f"Executing query: {formatted_query[:100]}...")

        try:
            result = self.client.query(formatted_query).result()

            # Convert Rows to dicts
            return [{key: row[key] for key in row.keys()} for row in result]

        except Exception as e:
            logger.error(f"Query failed: {e}")
            raise

    def execute(self, query: str) -> int:
        """
        Execute query without returning results.
        Returns number of rows affected.
        """
        formatted_query = self._format_query(query)
        logger.debug(f"Executing query: {formatted_query[:100]}...")

        try:
            job = self.client.query(formatted_query)
            job.result()  # Wait for completion
            return job.total_rows or 0

        except Exception as e:
            logger.error(f"Query failed: {e}")
            raise

    def stream_insert(
        self, table_id: str, rows: List[Dict[str, Any]], skip_invalid: bool = False
    ) -> List[Dict[str, Any]]:
        """
        Stream insert rows into table.
        Returns list of errors (empty if successful).
        """
        full_table_id = f"{self.project_id}.{self.dataset_id}.{table_id}"

        try:
            errors = self.client.insert_rows_json(
                full_table_id, rows, skip_invalid_rows=skip_invalid
            )

            if errors:
                logger.error(f"Insert errors: {errors}")
            else:
                logger.info(f"Inserted {len(rows)} rows into {table_id}")

            return errors

        except Exception as e:
            logger.error(f"Stream insert failed: {e}")
            raise

    def get_table_info(self, table_id: str) -> Dict[str, Any]:
        """Get table metadata"""
        full_table_id = f"{self.project_id}.{self.dataset_id}.{table_id}"

        try:
            table = self.client.get_table(full_table_id)
            return {
                "num_rows": table.num_rows,
                "num_bytes": table.num_bytes,
                "created": table.created,
                "modified": table.modified,
                "schema": [
                    {"name": field.name, "type": field.field_type}
                    for field in table.schema
                ],
            }
        except Exception as e:
            logger.error(f"Failed to get table info: {e}")
            raise