Synnefo

# Copyright 2012 GRNET S.A. All rights reserved.

#

# Redistribution and use in source and binary forms, with or

# without modification, are permitted provided that the following

# conditions are met:

#

#   1. Redistributions of source code must retain the above

#      copyright notice, this list of conditions and the following

#      disclaimer.

#

#   2. Redistributions in binary form must reproduce the above

#      copyright notice, this list of conditions and the following

#      disclaimer in the documentation and/or other materials

#      provided with the distribution.

#

# THIS SOFTWARE IS PROVIDED BY GRNET S.A. ``AS IS'' AND ANY EXPRESS

# OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED

# WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR

# PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL GRNET S.A OR

# CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,

# SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT

# LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF

# USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED

# AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT

# LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN

# ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE

# POSSIBILITY OF SUCH DAMAGE.

#

# The views and conclusions contained in the software and

# documentation are those of the authors and should not be

# interpreted as representing official policies, either expressed

# or implied, of GRNET S.A.

#

import psycopg2

from synnefo.lib.pool import ObjectPool

import logging

logger = logging.getLogger(__name__)

# Quick-n-dirty logging setup to stderr

#LOG_FORMAT = "%(asctime)-15s %(levelname)-6s %(message)s"

#handler = logging.StreamHandler()

#handler = logging.FileHandler("/tmp/aa1")

#handler.setFormatter(logging.Formatter(LOG_FORMAT))

#logger.addHandler(handler)

#logger.setLevel(logging.DEBUG)

#logger.fatal("kalhmera1")

#import sys

#print >>sys.stderr, "logger is %d" % id(logger)

_pool_kwargs = None

_pool = None

# How many times to retry on getting a live connection

# from the pool, before giving up.

RETRY_LIMIT = 1000

class PooledConnection(object):

    """Thin wrapper around a psycopg2 connection for a pooled connection.

    It takes care to put itself back into the pool upon connection closure.

    """

    def __init__(self, pool, conn):

        #print >>sys.stderr,"LOGGER is %d" % id(logger)

        logger.fatal("INIT POOLED CONN: pool = %s, conn = %s", pool, conn)

        self._pool = pool

        self._conn = conn

    def close(self):

        logger.debug("CLOSE POOLED CONN: self._pool = %s", self._pool)

        if not self._pool:

            return

        pool = self._pool

        logger.debug("PUT POOLED CONN: About to return %s to the pool", self)

        pool.pool_put(self)

        logger.debug("FINISHED PUT POOLED CONN: Returned %s to the pool", self)

    def __getattr__(self, attr):

        """Proxy every other call to the real connection"""

        return getattr(self._conn, attr)

    def __setattr__(self, attr, val):

        if attr not in ("_pool", "_conn"):

            setattr(self._conn, attr, val)

        object.__setattr__(self, attr, val)

    #def __del__(self):

    #    pass

    #    print >>sys.stderr, "DELETED"

class Psycopg2ConnectionPool(ObjectPool):

    """A synnefo.lib.pool.ObjectPool of psycopg2 connections.

    Every connection knows how to return itself to the pool

    when it gets close()d.

    """

    def __init__(self, **kw):

        ObjectPool.__init__(self, size=kw["synnefo_poolsize"])

        kw.pop("synnefo_poolsize")

        self._connection_args = kw

    def _pool_create(self):

        logger.info("CREATE: About to get a new connection from psycopg2")

        conn = psycopg2._original_connect(**self._connection_args)

        logger.info("CREATED: Got connection %s from psycopg2", conn)

        return PooledConnection(self, conn)

    def _pool_cleanup(self, pooledconn):

        logger.debug("CLEANING, conn = %d", id(pooledconn))

        try:

            # Reset this connection before putting it back

            # into the pool

            cursor = pooledconn.cursor()

            cursor.execute("ABORT; RESET ALL")

        except psycopg2.Error:

            # Since we're not going to be putting the psycopg2 connection

            # back into the pool, close it uncoditionally.

            logger.error("DEAD connection, conn = %d, %s",

                         id(pooledconn), pooledconn)

            try:

                pooledconn._conn.close()

            except:

                pass

            return True

        return False

def _init_pool(kw):

    global _pool

    global _pool_kwargs

    _pool_kwargs = kw

    _pool = Psycopg2ConnectionPool(**kw)

def _get_pool(kw):

    if not _pool:

        logger.debug("POOLINIT: Initializing DB connection pool")

        _init_pool(kw)

    if _pool_kwargs != kw:

        raise NotImplementedError(("Requested pooled psycopg2 connection with "

                                   "args %s != %s." % (kw, _pool_kwargs)))

    return _pool

def _verify_connection(pooledconn):

    try:

        cursor = pooledconn.cursor()

        cursor.execute("SELECT 1")

    except psycopg2.Error:

        # The connection has died.

        pooledconn.close()

        return False

    return True

def _pooled_connect(**kw):

    poolsize = kw.get("synnefo_poolsize", 0)

    if not poolsize:

        kw.pop("synnefo_poolsize", None)

        return psycopg2._original_connect(**kw)

    pool = _get_pool(kw)

    logger.debug("GET: Pool: set: %d, semaphore: %d",

                 len(pool._set), pool._semaphore._Semaphore__value)

    r = None

    n = 0

    while not r:

        r = pool.pool_get()

        if not _verify_connection(r):

            logger.error("DEADCONNECTION: Got dead connection %d from pool",

                         id(r))

            r = None

            n += 1

            if n > RETRY_LIMIT:

                raise RuntimeError(("Could not get live connection from pool"

                                    "after %d retries." % RETRY_LIMIT))

    logger.debug("GOT: Got connection %d from pool", id(r))

    return r

def monkey_patch_psycopg2():

    """Monkey-patch psycopg2's connect(), to retrieve connections from a pool.

    To enable pooling, you need to pass a synnefo_poolsize argument

    inside psycopg2's connection options.

    """

    if hasattr(psycopg2, '_original_connect'):

        return

    psycopg2._original_connect = psycopg2.connect

    psycopg2.connect = _pooled_connect

# Copyright (C) 2010 Daniele Varrazzo <daniele.varrazzo@gmail.com>

# and licensed under the MIT license:

#

# Permission is hereby granted, free of charge, to any person obtaining a copy

# of this software and associated documentation files (the "Software"), to deal

# in the Software without restriction, including without limitation the rights

# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell

# copies of the Software, and to permit persons to whom the Software is

# furnished to do so, subject to the following conditions:

#

# The above copyright notice and this permission notice shall be included in

# all copies or substantial portions of the Software.

#

# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR

# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,

# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE

# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER

# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,

# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN

# THE SOFTWARE.

"""A wait callback to allow psycopg2 cooperation with gevent.

Use `make_psycopg_green()` to enable gevent support in Psycopg.

"""

import psycopg2

from psycopg2 import extensions

from gevent.socket import wait_read, wait_write

def make_psycopg_green():

    """Configure Psycopg to be used with gevent in non-blocking way."""

    if not hasattr(extensions, 'set_wait_callback'):

        raise ImportError(

            "support for coroutines not available in this Psycopg version (%s)"

            % psycopg2.__version__)

    extensions.set_wait_callback(gevent_wait_callback)

def gevent_wait_callback(conn, timeout=None):

    """A wait callback useful to allow gevent to work with Psycopg."""

    while 1:

        state = conn.poll()

        if state == extensions.POLL_OK:

            break

        elif state == extensions.POLL_READ:

            wait_read(conn.fileno(), timeout=timeout)

        elif state == extensions.POLL_WRITE:

            wait_write(conn.fileno(), timeout=timeout)

        else:

            raise psycopg2.OperationalError("Bad result from poll: %r" % state)