cursor.py

"""A module for housing the Cursor class.

(C) Copyright 2013-2025 Dassault Systemes SE.  All Rights Reserved.

This software is licensed under a BSD 3-Clause License.
See the LICENSE file provided with this software.

Exported Classes:

Cursor -- Class for representing a database cursor.
"""

from collections import deque

# pylint: disable=unused-import
try:
    from typing import Any, Collection, Deque, Dict, Iterable, Mapping, List, Optional
except ImportError:
    pass

from .exception import Error, NotSupportedError, ProgrammingError

from . import statement       # pylint: disable=unused-import
from . import encodedsession  # pylint: disable=unused-import
from . import result_set      # pylint: disable=unused-import


class Cursor(object):
    """A database cursor.

    Public Functions:
    close -- Closes the cursor into the database.
    callproc -- Currently not supported.
    execute -- Executes an SQL operation.
    executemany -- Executes the operation for each list of paramaters passed in.
    fetchone -- Fetches the first row of results generated by the previous execute.
    fetchmany -- Fetches the number of rows that are passed in.
    fetchall -- Fetches everything generated by the previous execute.
    nextset -- Currently not supported.
    setinputsizes -- Currently not supported.
    setoutputsize -- Currently not supported.
    """

    description = None  # type: Optional[List[Any]]

    _result_set = None  # type: Optional[result_set.ResultSet]
    __query = None      # type: Optional[str]

    @property
    def query(self):
        # type: () -> Optional[str]
        """Return the most recent query."""
        return self.__query

    def __init__(self, session, cache_size):
        # type: (encodedsession.EncodedSession, int) -> None
        """Create a Cursor object.

        :param session: The session to use with this Cursor.
        """
        self.session = session
        self._statement_cache = StatementCache(session, cache_size)
        self._result_set = None
        self.closed = False
        self.arraysize = 1
        self.rowcount = -1
        self.colcount = -1
        self.rownumber = 0

    def __iter__(self):
        # type: () -> Cursor
        return self

    def next(self):
        # type: () -> result_set.Row
        """Return the next row of results from the previous SQL operation."""
        row = self.fetchone()
        if row is None:
            raise StopIteration
        return row

    def close(self):
        # type: () -> None
        """Close this cursor."""
        self._check_closed()
        self._statement_cache.shutdown()
        self._close_result_set()
        self.closed = True

    def _check_closed(self):
        # type: () -> None
        """Check if the cursor is closed.

        :raises Error: If the cursor is closed.
        """
        if self.closed:
            raise Error("cursor is closed")
        if self.session.closed:
            raise Error("connection is closed")

    def _close_result_set(self):
        # type: () -> None
        """Close current ResultSet on client and server side."""
        if self._result_set:
            self.session.close_result_set(self._result_set)
            self._result_set = None

    def _reset(self):
        # type: () -> None
        """Reset the Cursor to a default state."""
        self.description = None
        self.rowcount = -1
        self.colcount = -1
        self._close_result_set()

    def callproc(self, procname, parameters=None):  # pylint: disable=no-self-use
        # type: (str, Optional[Mapping[str, str]]) -> None
        """Invoke a stored procedure.

        Currently not supported.
        """
        if procname is not None or parameters is not None:
            raise NotSupportedError("Currently unsupported")

    def execute(self, operation, parameters=None):
        # type: (str, Optional[Collection[Any]]) -> None
        """Execute a SQL operation.

        :param operation: The SQL operation to be executed.
        :param parameters: Statement parameters.
        """
        self._check_closed()
        self._reset()
        self.__query = operation

        if parameters is None:
            exec_result = self._execute(operation)
        else:
            exec_result = self._executeprepared(operation, parameters)

        self.rowcount = exec_result.row_count
        if exec_result.result > 0:
            self._result_set = self.session.fetch_result_set(exec_result.statement)
            self.description = self.session.fetch_result_set_description(self._result_set)

        # TODO: ???
        if self.rowcount < 0:
            self.rowcount = -1
        self.rownumber = 0

    def _execute(self, operation):
        # type: (str) -> statement.ExecutionResult
        """Execute an operation without parameters.

        :param operation: SQL operation to execute.
        """
        return self.session.execute_statement(
            self._statement_cache.get_statement(), operation)

    def _executeprepared(self, operation, parameters):
        # type: (str, Collection[result_set.Value]) -> statement.ExecutionResult
        """Execute an operation with parameters.

        :param operation: SQL operation to execute.
        :raises ProgrammingError: Incorrect number of parameters
        """
        p_statement = self._statement_cache.get_prepared_statement(operation)
        if p_statement.parameter_count != len(parameters):
            raise ProgrammingError(
                "Incorrect number of parameters: expected %d, got %d" %
                (p_statement.parameter_count, len(parameters)))

        # Use handle to query
        return self.session.execute_prepared_statement(p_statement, parameters)

    def executemany(self, operation, seq_of_parameters):
        # type: (str, Collection[Collection[result_set.Value]]) -> List[int]
        """Execute the operation for each list of paramaters passed in."""
        self._check_closed()
        p_statement = self._statement_cache.get_prepared_statement(operation)
        return self.session.execute_batch_prepared_statement(
            p_statement, seq_of_parameters)

    def fetchone(self):
        # type: () -> Optional[result_set.Row]
        """Return the next row of results from the previous SQL operation."""
        # Inline _check_closed to avoid per-row function-call overhead.
        if self.closed:
            raise Error("cursor is closed")
        if self.session.closed:
            raise Error("connection is closed")
        if self._result_set is None:
            raise Error("Previous execute did not produce any results or no call was issued yet")
        self.rownumber += 1
        if not self._result_set.is_complete():
            self.session.fetch_result_set_next(self._result_set)
        return self._result_set.fetchone()

    def fetchmany(self, size=None):
        # type: (Optional[int]) -> List[result_set.Row]
        """Return SIZE rows from the previous SQL operation.

        If size is None, uses the default size for this Cursor.
        """
        self._check_closed()

        if size is None:
            size = self.arraysize

        fetched_rows = []
        num_fetched_rows = 0
        while num_fetched_rows < size:
            row = self.fetchone()
            if row is None:
                break
            else:
                fetched_rows.append(row)
                num_fetched_rows += 1
        return fetched_rows

    def fetchall(self):
        # type: () -> List[result_set.Row]
        """Return all rows generated by the previous SQL operation."""
        self._check_closed()
        if self._result_set is None:
            raise Error("Previous execute did not produce any results or no call was issued yet")

        fetched_rows = []  # type: List[result_set.Row]
        while True:
            # Drain the current in-memory batch in one shot instead of calling
            # fetchone() per row.
            idx = self._result_set.results_idx
            batch = self._result_set.results
            if idx < len(batch):
                fetched_rows.extend(batch[idx:])
                self._result_set.results_idx = len(batch)
            if self._result_set.complete:
                break
            self.session.fetch_result_set_next(self._result_set)
        self.rownumber += len(fetched_rows)
        return fetched_rows

    def nextset(self):  # pylint: disable=no-self-use
        # type: () -> None
        """Not supported."""
        raise NotSupportedError("Currently unsupported")

    def setinputsizes(self, sizes):
        # type: (int) -> None
        """Not supported."""
        pass

    def setoutputsize(self, size, column=None):
        # type: (int, Optional[int]) -> None
        """Not supported."""
        pass


class StatementCache(object):
    """Keep a cache of prepared statements."""

    def __init__(self, session, prepared_statement_cache_size):
        # type: (encodedsession.EncodedSession, int) -> None
        self._session = session
        self._statement = self._session.create_statement()
        self._ps_cache = dict()  # type: Dict[str, statement.PreparedStatement]
        self._ps_key_queue = deque()  # type: Deque[str]
        self._ps_cache_size = prepared_statement_cache_size

    def get_statement(self):
        # type; () -> statement.Statement
        """Return the Statement object for this cache."""
        return self._statement

    def get_prepared_statement(self, query):
        # type: (str) -> statement.PreparedStatement
        """Return a PreparedStatement for the provided query.

        If we don't have a cached PreparedStatement then create one and add it
        to the cache.  If we do have one move it to the front of the queue and
        return it.
        :returns: A PreparedStatement for the given query.
        """
        stmt = self._ps_cache.get(query)
        if stmt is not None:
            self._ps_key_queue.remove(query)
            self._ps_key_queue.append(query)
            return stmt

        stmt = self._session.create_prepared_statement(query)

        while len(self._ps_cache) >= self._ps_cache_size:
            lru_statement_key = self._ps_key_queue.popleft()
            statement_to_remove = self._ps_cache[lru_statement_key]
            self._session.close_statement(statement_to_remove)
            del self._ps_cache[lru_statement_key]

        self._ps_key_queue.append(query)
        self._ps_cache[query] = stmt

        return stmt

    def shutdown(self):
        # type: () -> None
        """Close the connection and clear the cursor cache."""
        self._session.close_statement(self._statement)

        for key in self._ps_cache:
            statement_to_remove = self._ps_cache[key]
            self._session.close_statement(statement_to_remove)

        self._ps_cache.clear()
        self._ps_key_queue.clear()