# -*- test-case-name: twisted.test.test_adbapi -*-
# Copyright (c) Twisted Matrix Laboratories.
# See LICENSE for details.

"""
An asynchronous mapping to U{DB-API
2.0<http://www.python.org/topics/database/DatabaseAPI-2.0.html>}.
"""


from twisted.internet import threads
from twisted.python import log, reflect


class ConnectionLost(Exception):
    """
    This exception means that a db connection has been lost.  Client code may
    try again.
    """


class Connection:
    """
    A wrapper for a DB-API connection instance.

    The wrapper passes almost everything to the wrapped connection and so has
    the same API. However, the L{Connection} knows about its pool and also
    handle reconnecting should when the real connection dies.
    """

    def __init__(self, pool):
        self._pool = pool
        self._connection = None
        self.reconnect()

    def close(self):
        # The way adbapi works right now means that closing a connection is
        # a really bad thing  as it leaves a dead connection associated with
        # a thread in the thread pool.
        # Really, I think closing a pooled connection should return it to the
        # pool but that's handled by the runWithConnection method already so,
        # rather than upsetting anyone by raising an exception, let's ignore
        # the request
        pass

    def rollback(self):
        if not self._pool.reconnect:
            self._connection.rollback()
            return

        try:
            self._connection.rollback()
            curs = self._connection.cursor()
            curs.execute(self._pool.good_sql)
            curs.close()
            self._connection.commit()
            return
        except BaseException:
            log.err(None, "Rollback failed")

        self._pool.disconnect(self._connection)

        if self._pool.noisy:
            log.msg("Connection lost.")

        raise ConnectionLost()

    def reconnect(self):
        if self._connection is not None:
            self._pool.disconnect(self._connection)
        self._connection = self._pool.connect()

    def __getattr__(self, name):
        return getattr(self._connection, name)


class Transaction:
    """
    A lightweight wrapper for a DB-API 'cursor' object.

    Relays attribute access to the DB cursor. That is, you can call
    C{execute()}, C{fetchall()}, etc., and they will be called on the
    underlying DB-API cursor object. Attributes will also be retrieved from
    there.
    """

    _cursor = None

    def __init__(self, pool, connection):
        self._pool = pool
        self._connection = connection
        self.reopen()

    def close(self):
        _cursor = self._cursor
        self._cursor = None
        _cursor.close()

    def reopen(self):
        if self._cursor is not None:
            self.close()

        try:
            self._cursor = self._connection.cursor()
            return
        except BaseException:
            if not self._pool.reconnect:
                raise
            else:
                log.err(None, "Cursor creation failed")

        if self._pool.noisy:
            log.msg("Connection lost, reconnecting")

        self.reconnect()
        self._cursor = self._connection.cursor()

    def reconnect(self):
        self._connection.reconnect()
        self._cursor = None

    def __getattr__(self, name):
        return getattr(self._cursor, name)


class ConnectionPool:
    """
    Represent a pool of connections to a DB-API 2.0 compliant database.

    @ivar connectionFactory: factory for connections, default to L{Connection}.
    @type connectionFactory: any callable.

    @ivar transactionFactory: factory for transactions, default to
        L{Transaction}.
    @type transactionFactory: any callable

    @ivar shutdownID: L{None} or a handle on the shutdown event trigger which
        will be used to stop the connection pool workers when the reactor
        stops.

    @ivar _reactor: The reactor which will be used to schedule startup and
        shutdown events.
    @type _reactor: L{IReactorCore} provider
    """

    CP_ARGS = "min max name noisy openfun reconnect good_sql".split()

    noisy = False  # If true, generate informational log messages
    min = 3  # Minimum number of connections in pool
    max = 5  # Maximum number of connections in pool
    name = None  # Name to assign to thread pool for debugging
    openfun = None  # A function to call on new connections
    reconnect = False  # Reconnect when connections fail
    good_sql = "select 1"  # A query which should always succeed

    running = False  # True when the pool is operating
    connectionFactory = Connection
    transactionFactory = Transaction

    # Initialize this to None so it's available in close() even if start()
    # never runs.
    shutdownID = None

    def __init__(self, dbapiName, *connargs, **connkw):
        """
        Create a new L{ConnectionPool}.

        Any positional or keyword arguments other than those documented here
        are passed to the DB-API object when connecting. Use these arguments to
        pass database names, usernames, passwords, etc.

        @param dbapiName: an import string to use to obtain a DB-API compatible
            module (e.g. C{'pyPgSQL.PgSQL'})

        @keyword cp_min: the minimum number of connections in pool (default 3)

        @keyword cp_max: the maximum number of connections in pool (default 5)

        @keyword cp_noisy: generate informational log messages during operation
            (default C{False})

        @keyword cp_openfun: a callback invoked after every C{connect()} on the
            underlying DB-API object. The callback is passed a new DB-API
            connection object. This callback can setup per-connection state
            such as charset, timezone, etc.

        @keyword cp_reconnect: detect connections which have failed and reconnect
            (default C{False}). Failed connections may result in
            L{ConnectionLost} exceptions, which indicate the query may need to
            be re-sent.

        @keyword cp_good_sql: an sql query which should always succeed and change
            no state (default C{'select 1'})

        @keyword cp_reactor: use this reactor instead of the global reactor
            (added in Twisted 10.2).
        @type cp_reactor: L{IReactorCore} provider
        """
        self.dbapiName = dbapiName
        self.dbapi = reflect.namedModule(dbapiName)

        if getattr(self.dbapi, "apilevel", None) != "2.0":
            log.msg("DB API module not DB API 2.0 compliant.")

        if getattr(self.dbapi, "threadsafety", 0) < 1:
            log.msg("DB API module not sufficiently thread-safe.")

        reactor = connkw.pop("cp_reactor", None)
        if reactor is None:
            from twisted.internet import reactor
        self._reactor = reactor

        self.connargs = connargs
        self.connkw = connkw

        for arg in self.CP_ARGS:
            cpArg = f"cp_{arg}"
            if cpArg in connkw:
                setattr(self, arg, connkw[cpArg])
                del connkw[cpArg]

        self.min = min(self.min, self.max)
        self.max = max(self.min, self.max)

        # All connections, hashed on thread id
        self.connections = {}

        # These are optional so import them here
        from twisted.python import threadable, threadpool

        self.threadID = threadable.getThreadID
        self.threadpool = threadpool.ThreadPool(self.min, self.max)
        self.startID = self._reactor.callWhenRunning(self._start)

    def _start(self):
        self.startID = None
        return self.start()

    def start(self):
        """
        Start the connection pool.

        If you are using the reactor normally, this function does *not*
        need to be called.
        """
        if not self.running:
            self.threadpool.start()
            self.shutdownID = self._reactor.addSystemEventTrigger(
                "during", "shutdown", self.finalClose
            )
            self.running = True

    def runWithConnection(self, func, *args, **kw):
        """
        Execute a function with a database connection and return the result.

        @param func: A callable object of one argument which will be executed
            in a thread with a connection from the pool. It will be passed as
            its first argument a L{Connection} instance (whose interface is
            mostly identical to that of a connection object for your DB-API
            module of choice), and its results will be returned as a
            L{Deferred}. If the method raises an exception the transaction will
            be rolled back. Otherwise, the transaction will be committed.
            B{Note} that this function is B{not} run in the main thread: it
            must be threadsafe.

        @param args: positional arguments to be passed to func

        @param kw: keyword arguments to be passed to func

        @return: a L{Deferred} which will fire the return value of
            C{func(Transaction(...), *args, **kw)}, or a
            L{twisted.python.failure.Failure}.
        """
        return threads.deferToThreadPool(
            self._reactor, self.threadpool, self._runWithConnection, func, *args, **kw
        )

    def _runWithConnection(self, func, *args, **kw):
        conn = self.connectionFactory(self)
        try:
            result = func(conn, *args, **kw)
            conn.commit()
            return result
        except BaseException:
            try:
                conn.rollback()
            except BaseException:
                log.err(None, "Rollback failed")
            raise

    def runInteraction(self, interaction, *args, **kw):
        """
        Interact with the database and return the result.

        The 'interaction' is a callable object which will be executed in a
        thread using a pooled connection. It will be passed an L{Transaction}
        object as an argument (whose interface is identical to that of the
        database cursor for your DB-API module of choice), and its results will
        be returned as a L{Deferred}. If running the method raises an
        exception, the transaction will be rolled back. If the method returns a
        value, the transaction will be committed.

        NOTE that the function you pass is *not* run in the main thread: you
        may have to worry about thread-safety in the function you pass to this
        if it tries to use non-local objects.

        @param interaction: a callable object whose first argument is an
            L{adbapi.Transaction}.

        @param args: additional positional arguments to be passed to
            interaction

        @param kw: keyword arguments to be passed to interaction

        @return: a Deferred which will fire the return value of
            C{interaction(Transaction(...), *args, **kw)}, or a
            L{twisted.python.failure.Failure}.
        """
        return threads.deferToThreadPool(
            self._reactor,
            self.threadpool,
            self._runInteraction,
            interaction,
            *args,
            **kw,
        )

    def runQuery(self, *args, **kw):
        """
        Execute an SQL query and return the result.

        A DB-API cursor which will be invoked with C{cursor.execute(*args,
        **kw)}. The exact nature of the arguments will depend on the specific
        flavor of DB-API being used, but the first argument in C{*args} be an
        SQL statement. The result of a subsequent C{cursor.fetchall()} will be
        fired to the L{Deferred} which is returned. If either the 'execute' or
        'fetchall' methods raise an exception, the transaction will be rolled
        back and a L{twisted.python.failure.Failure} returned.

        The C{*args} and C{**kw} arguments will be passed to the DB-API
        cursor's 'execute' method.

        @return: a L{Deferred} which will fire the return value of a DB-API
            cursor's 'fetchall' method, or a L{twisted.python.failure.Failure}.
        """
        return self.runInteraction(self._runQuery, *args, **kw)

    def runOperation(self, *args, **kw):
        """
        Execute an SQL query and return L{None}.

        A DB-API cursor which will be invoked with C{cursor.execute(*args,
        **kw)}. The exact nature of the arguments will depend on the specific
        flavor of DB-API being used, but the first argument in C{*args} will be
        an SQL statement. This method will not attempt to fetch any results
        from the query and is thus suitable for C{INSERT}, C{DELETE}, and other
        SQL statements which do not return values. If the 'execute' method
        raises an exception, the transaction will be rolled back and a
        L{Failure} returned.

        The C{*args} and C{*kw} arguments will be passed to the DB-API cursor's
        'execute' method.

        @return: a L{Deferred} which will fire with L{None} or a
            L{twisted.python.failure.Failure}.
        """
        return self.runInteraction(self._runOperation, *args, **kw)

    def close(self):
        """
        Close all pool connections and shutdown the pool.
        """
        if self.shutdownID:
            self._reactor.removeSystemEventTrigger(self.shutdownID)
            self.shutdownID = None
        if self.startID:
            self._reactor.removeSystemEventTrigger(self.startID)
            self.startID = None
        self.finalClose()

    def finalClose(self):
        """
        This should only be called by the shutdown trigger.
        """
        self.shutdownID = None
        self.threadpool.stop()
        self.running = False
        for conn in self.connections.values():
            self._close(conn)
        self.connections.clear()

    def connect(self):
        """
        Return a database connection when one becomes available.

        This method blocks and should be run in a thread from the internal
        threadpool. Don't call this method directly from non-threaded code.
        Using this method outside the external threadpool may exceed the
        maximum number of connections in the pool.

        @return: a database connection from the pool.
        """

        tid = self.threadID()
        conn = self.connections.get(tid)
        if conn is None:
            if self.noisy:
                log.msg(f"adbapi connecting: {self.dbapiName}")
            conn = self.dbapi.connect(*self.connargs, **self.connkw)
            if self.openfun is not None:
                self.openfun(conn)
            self.connections[tid] = conn
        return conn

    def disconnect(self, conn):
        """
        Disconnect a database connection associated with this pool.

        Note: This function should only be used by the same thread which called
        L{ConnectionPool.connect}. As with C{connect}, this function is not
        used in normal non-threaded Twisted code.
        """
        tid = self.threadID()
        if conn is not self.connections.get(tid):
            raise Exception("wrong connection for thread")
        if conn is not None:
            self._close(conn)
            del self.connections[tid]

    def _close(self, conn):
        if self.noisy:
            log.msg(f"adbapi closing: {self.dbapiName}")
        try:
            conn.close()
        except BaseException:
            log.err(None, "Connection close failed")

    def _runInteraction(self, interaction, *args, **kw):
        conn = self.connectionFactory(self)
        trans = self.transactionFactory(self, conn)
        try:
            result = interaction(trans, *args, **kw)
            trans.close()
            conn.commit()
            return result
        except BaseException:
            try:
                conn.rollback()
            except BaseException:
                log.err(None, "Rollback failed")
            raise

    def _runQuery(self, trans, *args, **kw):
        trans.execute(*args, **kw)
        return trans.fetchall()

    def _runOperation(self, trans, *args, **kw):
        trans.execute(*args, **kw)

    def __getstate__(self):
        return {
            "dbapiName": self.dbapiName,
            "min": self.min,
            "max": self.max,
            "noisy": self.noisy,
            "reconnect": self.reconnect,
            "good_sql": self.good_sql,
            "connargs": self.connargs,
            "connkw": self.connkw,
        }

    def __setstate__(self, state):
        self.__dict__ = state
        self.__init__(self.dbapiName, *self.connargs, **self.connkw)


__all__ = ["Transaction", "ConnectionPool"]