import logging
from typing import Optional
from queryservice_client.queryclient import QueryClient
from queryservice_client.stream.segments.statement import Statement

logger = logging.getLogger(__name__)

class QuerySession:
    def __init__(self, session_id: int, client: QueryClient):
        self.session_id = session_id
        self.client = client

    def query(self, query: str, limit: Optional[int] = None, timeout_in_seconds: Optional[int] = None) -> Statement:
        logger.info(f"Query timeout set to: {timeout_in_seconds} seconds")
        return self.client.query(self.session_id, query, limit, timeout_in_seconds)

    def cancel_query(self):
        self.client.cancel_query(self.session_id)

    def close(self):
        self.client.end_session(self.session_id)