~launchpad-pqm/launchpad/devel

# Copyright 2009 Canonical Ltd.  This software is licensed under the
# GNU Affero General Public License version 3 (see the file LICENSE).

"""A self maintaining mock database for tests.

The first time a `MockDbConnection` is used, it functions as a proxy
to the real database connection. Queries and results are recorded into
a script.

For subsequent runs, if the same queries are issued in the same order
then results are returned from the script and the real database is
not used. If the script is detected as being invalid, it is removed and
a `RetryTest` exception raised for the test runner to deal with.
"""

__metaclass__ = type
__all__ = [
        'MockDbConnection', 'script_filename',
        'ScriptPlayer', 'ScriptRecorder',
        ]

import cPickle as pickle
import gzip
import os.path
import urllib

import psycopg2

# from zope.testing.testrunner import RetryTest

from lp.services.config import config


SCRIPT_DIR = os.path.join(config.root, 'mockdbscripts~')


def script_filename(key):
    """Calculate and return the script filename to use."""
    key = urllib.quote(key, safe='')
    return os.path.join(SCRIPT_DIR, key) + '.pickle.gz'


class ScriptEntry:
    """An entry in our test's log of database calls."""

    # The connection number used for this command. All connections used
    # by a test store their commands in a single list to preserve global
    # ordering, and we use connection_number to differentiate them.
    connection_number = None

    # If the command raised an exception, it is stored here.
    exception = None

    def __init__(self, connection):
        if connection is not None:
            self.connection_number = connection.connection_number


class ConnectScriptEntry(ScriptEntry):
    """An entry created instantiating a Connection."""
    args = None # Arguments passed to the connect() method.
    kw = None # Keyword arguments passed to the connect() method.
    
    def __init__(self, connection, *args, **kw):
        super(ConnectScriptEntry, self).__init__(connection)
        self.args = args
        self.kw = kw


class ExecuteScriptEntry(ScriptEntry):
    """An entry created via Cursor.execute()."""
    query = None # Query passed to Cursor.execute().
    params = None # Parameters passed to Cursor.execute().
    results = None # Cursor.fetchall() results as a list.
    description = None # Cursor.description as per DB-API.
    rowcount = None # Cursor.rowcount after Cursor.fetchall() as per DB-API.

    def __init__(self, connection, query, params):
        super(ExecuteScriptEntry, self).__init__(connection)
        self.query = query
        self.params = params


class CloseScriptEntry(ScriptEntry):
    """An entry created via Connection.close()."""


class CommitScriptEntry(ScriptEntry):
    """An entry created via Connection.commit()."""


class RollbackScriptEntry(ScriptEntry):
    """An entry created via Connection.rollback()."""


class SetIsolationLevelScriptEntry(ScriptEntry):
    """An entry created via Connection.set_isolation_level()."""
    level = None # The requested isolation level
    def __init__(self, connection, level):
        super(SetIsolationLevelScriptEntry, self).__init__(connection)
        self.level = level


class ScriptRecorder:
    key = None # A unique key identifying this test.
    script_filename = None # Path to our script file.
    log = None
    connections = None

    def __init__(self, key):
        self.key = key
        self.script_filename = script_filename(key)
        self.log = []
        self.connections = []

    def connect(self, connect_func, *args, **kw):
        """Open a connection to the database, returning a `MockDbConnection`.
        """
        try:
            connection = connect_func(*args, **kw)
            exception = None
        except (psycopg2.Warning, psycopg2.Error), connect_exception:
            connection = None
            exception = connect_exception

        connection = MockDbConnection(self, connection, *args, **kw)

        self.connections.append(connection)
        if connection is not None:
            connection.connection_number = self.connections.index(connection)
        entry = ConnectScriptEntry(connection, *args, **kw)
        self.log.append(entry)
        if exception:
            entry.exception = exception
            #pylint: disable-msg=W0706
            raise exception
        return connection

    def cursor(self, connection):
        """Return a MockDbCursor."""
        real_cursor = connection.real_connection.cursor()
        return MockDbCursor(connection, real_cursor)

    def execute(self, cursor, query, params=None):
        """Handle Cursor.execute()."""
        con = cursor.connection
        entry = ExecuteScriptEntry(con, query, params)

        real_cursor = cursor.real_cursor
        try:
            real_cursor.execute(query, params)
        except (psycopg2.Warning, psycopg2.Error), exception:
            entry.exception = exception
            self.log.append(entry)
            raise

        entry.rowcount = real_cursor.rowcount
        try:
            entry.results = list(real_cursor.fetchall())
            entry.description = real_cursor.description
            # Might have changed now fetchall() has been done.
            entry.rowcount = real_cursor.rowcount
        except psycopg2.ProgrammingError, exc:
            if exc.args[0] == 'no results to fetch':
                # No results, such as an UPDATE query.
                entry.results = None
            else:
                raise

        self.log.append(entry)
        return entry

    def close(self, connection):
        """Handle Connection.close()."""
        entry = CloseScriptEntry(connection)
        try:
            if connection.real_connection is not None:
                connection.real_connection.close()
        except (psycopg2.Warning, psycopg2.Error), exception:
            entry.exception = exception
            self.log.append(entry)
            raise
        else:
            self.log.append(entry)

    def commit(self, connection):
        """Handle Connection.commit()."""
        entry = CommitScriptEntry(connection)
        try:
            connection.real_connection.commit()
        except (psycopg2.Warning, psycopg2.Error), exception:
            entry.exception = exception
            self.log.append(entry)
            raise
        else:
            self.log.append(entry)

    def rollback(self, connection):
        """Handle Connection.rollback()."""
        entry = RollbackScriptEntry(connection)
        try:
            connection.real_connection.rollback()
        except (psycopg2.Warning, psycopg2.Error), exception:
            entry.exception = exception
            self.log.append(entry)
            raise
        else:
            self.log.append(entry)

    def set_isolation_level(self, connection, level):
        """Handle Connection.set_isolation_level()."""
        entry = SetIsolationLevelScriptEntry(connection, level)
        try:
            connection.real_connection.set_isolation_level(level)
        except (psycopg2.Warning, psycopg2.Error), exception:
            entry.exception = exception
            self.log.append(entry)
            raise
        else:
            self.log.append(entry)

    def store(self):
        """Store the script for future use by a ScriptPlayer."""
        # Create script directory if necessary.
        if not os.path.isdir(SCRIPT_DIR):
            os.makedirs(SCRIPT_DIR, mode=0700)

        # Save log to a pickle. The pickle contains the key for this test,
        # followed by a list of ScriptEntry-derived objects.
        obj_to_store = [self.key] + self.log
        pickle.dump(
                obj_to_store, gzip.open(self.script_filename, 'wb'),
                pickle.HIGHEST_PROTOCOL
                )

        # Trash all the connected connections. This isn't strictly necessary
        # but protects us from silly mistakes.
        while self.connections:
            con = self.connections.pop()
            if con is not None and not con._closed:
                con.close()


def noop_if_invalid(func):
    """Decorator that causes the decorated method to be a noop if the
    script this method belongs too is invalid.

    This allows teardown to complete when a ReplayScript has
    raised a RetryTest exception. Normally during teardown DB operations
    are made when doing things like aborting the transaction.
    """
    def dont_retry_func(self, *args, **kw):
        if self.invalid:
            return None
        else:
            return func(self, *args, **kw)
    return dont_retry_func


class ScriptPlayer:
    """Replay database queries from a script."""

    key = None # Unique key identifying this test.
    script_filename = None # File storing our statement/result script.
    log = None # List of ScriptEntry objects loaded from _script_filename.
    connections = None # List of connections using this script.

    invalid = False # If True, the script is invalid and we are tearing down.

    def __init__(self, key):
        self.key = key
        self.script_filename = script_filename(key)
        self.log = pickle.load(gzip.open(self.script_filename, 'rb'))
        try:
            stored_key = self.log.pop(0)
            assert stored_key == key, "Script loaded for wrong key."
        except IndexError:
            self.handleInvalidScript(
                    "Connection key not stored in script."
                    )

        self.connections = []

    def getNextEntry(self, connection, expected_entry_class):
        """Pull the next entry from the script.

        Invokes handleInvalidScript on error, including some entry validation.
        """
        try:
            entry = self.log.pop(0)
        except IndexError:
            self.handleInvalidScript('Ran out of commands.')

        # This guards against file format changes as well as file corruption.
        if not isinstance(entry, ScriptEntry):
            self.handleInvalidScript('Unexpected object type in script')

        if connection.connection_number != entry.connection_number:
            self.handleInvalidScript(
                    'Expected query to connection %s '
                    'but got query to connection %s'
                    % (entry.connection_number, connection.connection_number)
                    )

        if not isinstance(entry, expected_entry_class):
            self.handleInvalidScript(
                    'Expected %s but got %s'
                    % (expected_entry_class, entry.__class__)
                    )

        return entry

    @noop_if_invalid
    def connect(self, connect_func, *args, **kw):
        """Return a `MockDbConnection`.
       
        Does not actually connect to the database - we are in replay mode.
        """
        connection = MockDbConnection(self, None, *args, **kw)
        self.connections.append(connection)
        connection.connection_number = self.connections.index(connection)
        entry = self.getNextEntry(connection, ConnectScriptEntry)
        if (entry.args, entry.kw) != (args, kw):
            self.handleInvalidScript("Connection parameters have changed.")
        if entry.exception is not None:
            raise entry.exception
        return connection

    def cursor(self, connection):
        """Return a MockDbCursor."""
        return MockDbCursor(connection, real_cursor=None)

    @noop_if_invalid
    def execute(self, cursor, query, params=None):
        """Handle Cursor.execute()."""
        connection = cursor.connection
        entry = self.getNextEntry(connection, ExecuteScriptEntry)

        if query != entry.query:
            self.handleInvalidScript(
                    'Unexpected command. Expected %s. Got %s.'
                    % (entry.query, query)
                    )

        if params != entry.params:
            self.handleInvalidScript(
                    'Unexpected parameters. Expected %r. Got %r.'
                    % (entry.params, params)
                    )

        if entry.exception is not None:
            raise entry.exception

        return entry

    @noop_if_invalid
    def close(self, connection):
        """Handle Connection.close()."""
        entry = self.getNextEntry(connection, CloseScriptEntry)
        if entry.exception is not None:
            raise entry.exception

    def commit(self, connection):
        """Handle Connection.commit()."""
        entry = self.getNextEntry(connection, CommitScriptEntry)
        if entry.exception is not None:
            raise entry.exception

    @noop_if_invalid
    def rollback(self, connection):
        """Handle Connection.rollback()."""
        entry = self.getNextEntry(connection, RollbackScriptEntry)
        if entry.exception is not None:
            raise entry.exception

    @noop_if_invalid
    def set_isolation_level(self, connection, level):
        """Handle Connection.set_isolation_level()."""
        entry = self.getNextEntry(connection, SetIsolationLevelScriptEntry)
        if entry.level != level:
            self.handleInvalidScript("Different isolation level requested.")
        if entry.exception is not None:
            raise entry.exception

    def handleInvalidScript(self, reason):
        """Remove the script from disk and raise a RetryTest exception."""
        if os.path.exists(self.script_filename):
            os.unlink(self.script_filename)
        self.invalid = True
        raise RetryTest(reason) # Leaving this as a name error: this should
        # not be called unless we reinstate the retry behavior in zope.testing.


class MockDbConnection:
    """Connection to our Mock database."""

    real_connection = None
    connection_number = None
    script = None

    def __init__(self, script, real_connection, *args, **kw):
        """Initialize the `MockDbConnection`.

        `MockDbConnection` intances are generally only created in the
        *Script.connect() methods as the attempt needs to be recorded
        (even if it fails).

        If we have a real_connection, we are proxying and recording results.
        If real_connection is None, we are replaying results from the script.

        *args and **kw are the arguments passed to open the real connection
        and are used by the script to confirm the db connection details have
        not been changed; a `RetryTest` exception may be raised in replay mode.
        """
        self.script = script
        self.real_connection = real_connection

    def cursor(self):
        """As per DB-API."""
        return self.script.cursor(self)

    _closed = False

    def _checkClosed(self):
        """Guard that raises an exception if the connection is closed."""
        if self._closed is True:
            raise psycopg2.InterfaceError('Connection closed.')

    def close(self):
        """As per DB-API."""
        # DB-API says closing a closed connection should raise an exception
        # ("exception will be raised if any operation is attempted
        # wht the [closed] connection"), but psycopg1 doesn't do this.
        # It would be nice if our wrapper could be more strict than psycopg1,
        # but unfortunately the sqlos/sqlobject combination relies on this
        # behavior. So we have to emulate it.
        if self._closed:
            return
        #self._checkClosed()
        self.script.close(self)
        self._closed = True

    def commit(self):
        """As per DB-API."""
        self._checkClosed()
        self.script.commit(self)

    def rollback(self):
        """As per DB-API."""
        self._checkClosed()
        self.script.rollback(self)

    def set_isolation_level(self, level):
        """As per psycopg1 extension."""
        self._checkClosed()
        self.script.set_isolation_level(self, level)

    # Exceptions exposed on connection, as per optional DB-API extension.
    ## Disabled, as psycopg1 does not implement this extension.
    ## Warning = psycopg.Warning
    ## Error = psycopg.Error
    ## InterfaceError = psycopg.InterfaceError
    ## DatabaseError = psycopg.DatabaseError
    ## DataError = psycopg.DataError
    ## OperationalError = psycopg.OperationalError
    ## IntegrityError = psycopg.IntegrityError
    ## InternalError = psycopg.InternalError
    ## ProgrammingError = psycopg.ProgrammingError
    ## NotSupportedError = psycopg.NotSupportedError


class MockDbCursor:
    """A fake DB-API cursor as produced by MockDbConnection.cursor.
    
    The real work is done by the associated ScriptRecorder or ScriptPlayer
    using the common interface, making this class independent on what
    mode the mock database is running in.
    """
    # The ExecuteScriptEntry for the in progress query, if there is one.
    # It stores any unfetched results and cursor metadata such as the
    # resultset description.
    _script_entry = None

    arraysize = 100 # As per DB-API.
    connection = None # As per DB-API optional extension.
    real_cursor = None # The real cursor if there is one.

    def __init__(self, connection, real_cursor=None):
        self.connection = connection
        self.real_cursor = real_cursor

    @property
    def description(self):
        """As per DB-API, pulled from the script entry."""
        if self._script_entry is None:
            return None
        return self._script_entry.description

    @property
    def rowcount(self):
        """Return the rowcount only if all the results have been consumed.

        As per DB-API, pulled from the script entry.
        """
        if not isinstance(self._script_entry, ExecuteScriptEntry):
            return -1

        results = self._script_entry.results

        if results is None: # DELETE or UPDATE set rowcount.
            return self._script_entry.rowcount
        
        if results is None or self._fetch_position < len(results):
            return -1
        return self._script_entry.rowcount

    _closed = False

    def close(self):
        """As per DB-API."""
        self._checkClosed()
        self._closed = True
        if self.real_cursor is not None:
            self.real_cursor.close()
            self.real_cursor = None
            self.connection = None

    def _checkClosed(self):
        """Raise an exception if the cursor or connection is closed."""
        if self._closed is True:
            raise psycopg2.Error('Cursor closed.')
        self.connection._checkClosed()

    # Index in our results that the next fetch will return. We don't consume
    # the results list: if we are recording we need to keep the results 
    # so we can serialize at the end of the test.
    _fetch_position = 0

    def execute(self, query, parameters=None):
        """As per DB-API."""
        self._checkClosed()
        self._script_entry = self.connection.script.execute(
                self, query, parameters
                )
        self._fetch_position = 0

    def executemany(self, query, seq_of_parameters=None):
        """As per DB-API."""
        self._checkClosed()
        raise NotImplementedError('executemany')

    def fetchone(self):
        """As per DB-API."""
        self._checkClosed()
        if self._script_entry is None:
            raise psycopg2.Error("No query issued yet")
        if self._script_entry.results is None:
            raise psycopg2.Error("Query returned no results")
        try:
            row = self._script_entry.results[self._fetch_position]
            self._fetch_position += 1
            return row
        except IndexError:
            return None

    def fetchmany(self, size=None):
        """As per DB-API."""
        self._checkClosed()
        raise NotImplementedError('fetchmany')

    def fetchall(self):
        """As per DB-API."""
        self._checkClosed()
        if self._script_entry is None:
            raise psycopg2.Error('No query issued yet')
        if self._script_entry.results is None:
            raise psycopg2.Error('Query returned no results')
        results = self._script_entry.results[self._fetch_position:]
        self._fetch_position = len(results)
        return results

    def nextset(self):
        """As per DB-API."""
        self._checkClosed()
        raise NotImplementedError('nextset')

    def setinputsizes(self, sizes):
        """As per DB-API."""
        self._checkClosed()
        return # No-op.

    def setoutputsize(self, size, column=None):
        """As per DB-API."""
        self._checkClosed()
        return # No-op.

    ## psycopg1 does not support this extension.
    ##
    ## def next(self):
    ##     """As per iterator spec and DB-API optional extension."""
    ##     row = self.fetchone()
    ##     if row is None:
    ##         raise StopIteration
    ##     else:
    ##         return row

    ## def __iter__(self):
    ##     """As per iterator spec and DB-API optional extension."""
    ##     return self