pgtools/pgmanager.py
author Radek Brich <radek.brich@devl.cz>
Wed, 10 Aug 2011 18:34:54 +0200
changeset 6 4ab077c93b2d
parent 4 80634cb1c65b
permissions -rw-r--r--
Add table diff tool.
Ignore whitespace changes - Everywhere: Within whitespace: At end of lines:
0
eaae9539e910 Postgres tools.
Radek Brich <radek.brich@devl.cz>
parents:
diff changeset
     1
# -*- coding: utf-8 -*-
eaae9539e910 Postgres tools.
Radek Brich <radek.brich@devl.cz>
parents:
diff changeset
     2
#
eaae9539e910 Postgres tools.
Radek Brich <radek.brich@devl.cz>
parents:
diff changeset
     3
# PgManager - manage database connections
eaae9539e910 Postgres tools.
Radek Brich <radek.brich@devl.cz>
parents:
diff changeset
     4
#
eaae9539e910 Postgres tools.
Radek Brich <radek.brich@devl.cz>
parents:
diff changeset
     5
# Requires: Python 2.6, psycopg2
eaae9539e910 Postgres tools.
Radek Brich <radek.brich@devl.cz>
parents:
diff changeset
     6
#
eaae9539e910 Postgres tools.
Radek Brich <radek.brich@devl.cz>
parents:
diff changeset
     7
# Copyright (c) 2010, 2011  Radek Brich <radek.brich@devl.cz>
eaae9539e910 Postgres tools.
Radek Brich <radek.brich@devl.cz>
parents:
diff changeset
     8
#
eaae9539e910 Postgres tools.
Radek Brich <radek.brich@devl.cz>
parents:
diff changeset
     9
# Permission is hereby granted, free of charge, to any person obtaining a copy
eaae9539e910 Postgres tools.
Radek Brich <radek.brich@devl.cz>
parents:
diff changeset
    10
# of this software and associated documentation files (the "Software"), to deal
eaae9539e910 Postgres tools.
Radek Brich <radek.brich@devl.cz>
parents:
diff changeset
    11
# in the Software without restriction, including without limitation the rights
eaae9539e910 Postgres tools.
Radek Brich <radek.brich@devl.cz>
parents:
diff changeset
    12
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
eaae9539e910 Postgres tools.
Radek Brich <radek.brich@devl.cz>
parents:
diff changeset
    13
# copies of the Software, and to permit persons to whom the Software is
eaae9539e910 Postgres tools.
Radek Brich <radek.brich@devl.cz>
parents:
diff changeset
    14
# furnished to do so, subject to the following conditions:
eaae9539e910 Postgres tools.
Radek Brich <radek.brich@devl.cz>
parents:
diff changeset
    15
#
eaae9539e910 Postgres tools.
Radek Brich <radek.brich@devl.cz>
parents:
diff changeset
    16
# The above copyright notice and this permission notice shall be included in
eaae9539e910 Postgres tools.
Radek Brich <radek.brich@devl.cz>
parents:
diff changeset
    17
# all copies or substantial portions of the Software.
eaae9539e910 Postgres tools.
Radek Brich <radek.brich@devl.cz>
parents:
diff changeset
    18
#
eaae9539e910 Postgres tools.
Radek Brich <radek.brich@devl.cz>
parents:
diff changeset
    19
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
eaae9539e910 Postgres tools.
Radek Brich <radek.brich@devl.cz>
parents:
diff changeset
    20
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
eaae9539e910 Postgres tools.
Radek Brich <radek.brich@devl.cz>
parents:
diff changeset
    21
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
eaae9539e910 Postgres tools.
Radek Brich <radek.brich@devl.cz>
parents:
diff changeset
    22
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
eaae9539e910 Postgres tools.
Radek Brich <radek.brich@devl.cz>
parents:
diff changeset
    23
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
eaae9539e910 Postgres tools.
Radek Brich <radek.brich@devl.cz>
parents:
diff changeset
    24
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
eaae9539e910 Postgres tools.
Radek Brich <radek.brich@devl.cz>
parents:
diff changeset
    25
# THE SOFTWARE.
eaae9539e910 Postgres tools.
Radek Brich <radek.brich@devl.cz>
parents:
diff changeset
    26
eaae9539e910 Postgres tools.
Radek Brich <radek.brich@devl.cz>
parents:
diff changeset
    27
"""Postgres database connection manager
eaae9539e910 Postgres tools.
Radek Brich <radek.brich@devl.cz>
parents:
diff changeset
    28
eaae9539e910 Postgres tools.
Radek Brich <radek.brich@devl.cz>
parents:
diff changeset
    29
PgManager wraps psycopg2 connect function, adding following features:
eaae9539e910 Postgres tools.
Radek Brich <radek.brich@devl.cz>
parents:
diff changeset
    30
eaae9539e910 Postgres tools.
Radek Brich <radek.brich@devl.cz>
parents:
diff changeset
    31
 * Manage database connection parameters - link connection parameters
eaae9539e910 Postgres tools.
Radek Brich <radek.brich@devl.cz>
parents:
diff changeset
    32
   to an unique identifier, retrieve connection object by this identifier
eaae9539e910 Postgres tools.
Radek Brich <radek.brich@devl.cz>
parents:
diff changeset
    33
eaae9539e910 Postgres tools.
Radek Brich <radek.brich@devl.cz>
parents:
diff changeset
    34
 * Connection pooling - connections with same identifier are pooled and reused
eaae9539e910 Postgres tools.
Radek Brich <radek.brich@devl.cz>
parents:
diff changeset
    35
eaae9539e910 Postgres tools.
Radek Brich <radek.brich@devl.cz>
parents:
diff changeset
    36
 * Easy query using the with statement - retrieve cursor directly by connection
eaae9539e910 Postgres tools.
Radek Brich <radek.brich@devl.cz>
parents:
diff changeset
    37
   identifier, don't worry about connections
eaae9539e910 Postgres tools.
Radek Brich <radek.brich@devl.cz>
parents:
diff changeset
    38
eaae9539e910 Postgres tools.
Radek Brich <radek.brich@devl.cz>
parents:
diff changeset
    39
 * Dict rows - cursor has additional methods like fetchall_dict(), which
eaae9539e910 Postgres tools.
Radek Brich <radek.brich@devl.cz>
parents:
diff changeset
    40
   returns dict row instead of ordinary list-like row
eaae9539e910 Postgres tools.
Radek Brich <radek.brich@devl.cz>
parents:
diff changeset
    41
eaae9539e910 Postgres tools.
Radek Brich <radek.brich@devl.cz>
parents:
diff changeset
    42
Example:
eaae9539e910 Postgres tools.
Radek Brich <radek.brich@devl.cz>
parents:
diff changeset
    43
eaae9539e910 Postgres tools.
Radek Brich <radek.brich@devl.cz>
parents:
diff changeset
    44
import pgmanager
eaae9539e910 Postgres tools.
Radek Brich <radek.brich@devl.cz>
parents:
diff changeset
    45
eaae9539e910 Postgres tools.
Radek Brich <radek.brich@devl.cz>
parents:
diff changeset
    46
pgm = pgmanager.get_instance()
eaae9539e910 Postgres tools.
Radek Brich <radek.brich@devl.cz>
parents:
diff changeset
    47
pgm.create_conn(hostaddr='127.0.0.1', dbname='postgres')
eaae9539e910 Postgres tools.
Radek Brich <radek.brich@devl.cz>
parents:
diff changeset
    48
eaae9539e910 Postgres tools.
Radek Brich <radek.brich@devl.cz>
parents:
diff changeset
    49
with pgm.cursor() as curs:
eaae9539e910 Postgres tools.
Radek Brich <radek.brich@devl.cz>
parents:
diff changeset
    50
    curs.execute('SELECT now() AS now')
eaae9539e910 Postgres tools.
Radek Brich <radek.brich@devl.cz>
parents:
diff changeset
    51
    row = curs.fetchone_dict()
eaae9539e910 Postgres tools.
Radek Brich <radek.brich@devl.cz>
parents:
diff changeset
    52
    print row.now
eaae9539e910 Postgres tools.
Radek Brich <radek.brich@devl.cz>
parents:
diff changeset
    53
eaae9539e910 Postgres tools.
Radek Brich <radek.brich@devl.cz>
parents:
diff changeset
    54
First, we have obtained PgManager instance. This is like calling
eaae9539e910 Postgres tools.
Radek Brich <radek.brich@devl.cz>
parents:
diff changeset
    55
PgManager(), although in our example the instance is global. That means
eaae9539e910 Postgres tools.
Radek Brich <radek.brich@devl.cz>
parents:
diff changeset
    56
getting the instance in another module brings us all the defined connections
eaae9539e910 Postgres tools.
Radek Brich <radek.brich@devl.cz>
parents:
diff changeset
    57
etc.
eaae9539e910 Postgres tools.
Radek Brich <radek.brich@devl.cz>
parents:
diff changeset
    58
eaae9539e910 Postgres tools.
Radek Brich <radek.brich@devl.cz>
parents:
diff changeset
    59
On second line we created connection named 'default' (this name can be left out).
eaae9539e910 Postgres tools.
Radek Brich <radek.brich@devl.cz>
parents:
diff changeset
    60
The with statement obtains connection (actually connects to database when needed),
eaae9539e910 Postgres tools.
Radek Brich <radek.brich@devl.cz>
parents:
diff changeset
    61
then returns cursor for this connection. On exit, the connection is returned
eaae9539e910 Postgres tools.
Radek Brich <radek.brich@devl.cz>
parents:
diff changeset
    62
to the pool or closed (depending on number of connections on pool and setting
eaae9539e910 Postgres tools.
Radek Brich <radek.brich@devl.cz>
parents:
diff changeset
    63
of keep_open parameter).
eaae9539e910 Postgres tools.
Radek Brich <radek.brich@devl.cz>
parents:
diff changeset
    64
eaae9539e910 Postgres tools.
Radek Brich <radek.brich@devl.cz>
parents:
diff changeset
    65
The row returned by fetchone_dict() is special dict object, which can be accessed
eaae9539e910 Postgres tools.
Radek Brich <radek.brich@devl.cz>
parents:
diff changeset
    66
using item or attribute access, that is row['now'] or row.now.
eaae9539e910 Postgres tools.
Radek Brich <radek.brich@devl.cz>
parents:
diff changeset
    67
"""
eaae9539e910 Postgres tools.
Radek Brich <radek.brich@devl.cz>
parents:
diff changeset
    68
eaae9539e910 Postgres tools.
Radek Brich <radek.brich@devl.cz>
parents:
diff changeset
    69
from contextlib import contextmanager
eaae9539e910 Postgres tools.
Radek Brich <radek.brich@devl.cz>
parents:
diff changeset
    70
import logging
eaae9539e910 Postgres tools.
Radek Brich <radek.brich@devl.cz>
parents:
diff changeset
    71
import threading
eaae9539e910 Postgres tools.
Radek Brich <radek.brich@devl.cz>
parents:
diff changeset
    72
import select
eaae9539e910 Postgres tools.
Radek Brich <radek.brich@devl.cz>
parents:
diff changeset
    73
eaae9539e910 Postgres tools.
Radek Brich <radek.brich@devl.cz>
parents:
diff changeset
    74
import psycopg2
eaae9539e910 Postgres tools.
Radek Brich <radek.brich@devl.cz>
parents:
diff changeset
    75
import psycopg2.extensions
eaae9539e910 Postgres tools.
Radek Brich <radek.brich@devl.cz>
parents:
diff changeset
    76
eaae9539e910 Postgres tools.
Radek Brich <radek.brich@devl.cz>
parents:
diff changeset
    77
from psycopg2 import DatabaseError, IntegrityError
eaae9539e910 Postgres tools.
Radek Brich <radek.brich@devl.cz>
parents:
diff changeset
    78
eaae9539e910 Postgres tools.
Radek Brich <radek.brich@devl.cz>
parents:
diff changeset
    79
eaae9539e910 Postgres tools.
Radek Brich <radek.brich@devl.cz>
parents:
diff changeset
    80
class PgManagerError(Exception):
eaae9539e910 Postgres tools.
Radek Brich <radek.brich@devl.cz>
parents:
diff changeset
    81
eaae9539e910 Postgres tools.
Radek Brich <radek.brich@devl.cz>
parents:
diff changeset
    82
    pass
eaae9539e910 Postgres tools.
Radek Brich <radek.brich@devl.cz>
parents:
diff changeset
    83
eaae9539e910 Postgres tools.
Radek Brich <radek.brich@devl.cz>
parents:
diff changeset
    84
eaae9539e910 Postgres tools.
Radek Brich <radek.brich@devl.cz>
parents:
diff changeset
    85
class ConnectionInfo:
eaae9539e910 Postgres tools.
Radek Brich <radek.brich@devl.cz>
parents:
diff changeset
    86
eaae9539e910 Postgres tools.
Radek Brich <radek.brich@devl.cz>
parents:
diff changeset
    87
    def __init__(self, dsn, isolation_level=None, init_statement=None, keep_open=1):
eaae9539e910 Postgres tools.
Radek Brich <radek.brich@devl.cz>
parents:
diff changeset
    88
        self.dsn = dsn
eaae9539e910 Postgres tools.
Radek Brich <radek.brich@devl.cz>
parents:
diff changeset
    89
        self.isolation_level = isolation_level
eaae9539e910 Postgres tools.
Radek Brich <radek.brich@devl.cz>
parents:
diff changeset
    90
        self.init_statement = init_statement
eaae9539e910 Postgres tools.
Radek Brich <radek.brich@devl.cz>
parents:
diff changeset
    91
        self.keep_open = keep_open
eaae9539e910 Postgres tools.
Radek Brich <radek.brich@devl.cz>
parents:
diff changeset
    92
eaae9539e910 Postgres tools.
Radek Brich <radek.brich@devl.cz>
parents:
diff changeset
    93
eaae9539e910 Postgres tools.
Radek Brich <radek.brich@devl.cz>
parents:
diff changeset
    94
class RowDict(dict):
eaae9539e910 Postgres tools.
Radek Brich <radek.brich@devl.cz>
parents:
diff changeset
    95
eaae9539e910 Postgres tools.
Radek Brich <radek.brich@devl.cz>
parents:
diff changeset
    96
    def __getattr__(self, key):
eaae9539e910 Postgres tools.
Radek Brich <radek.brich@devl.cz>
parents:
diff changeset
    97
        return self[key]
eaae9539e910 Postgres tools.
Radek Brich <radek.brich@devl.cz>
parents:
diff changeset
    98
eaae9539e910 Postgres tools.
Radek Brich <radek.brich@devl.cz>
parents:
diff changeset
    99
eaae9539e910 Postgres tools.
Radek Brich <radek.brich@devl.cz>
parents:
diff changeset
   100
class Cursor(psycopg2.extensions.cursor):
eaae9539e910 Postgres tools.
Radek Brich <radek.brich@devl.cz>
parents:
diff changeset
   101
eaae9539e910 Postgres tools.
Radek Brich <radek.brich@devl.cz>
parents:
diff changeset
   102
    def execute(self, query, args=None):
eaae9539e910 Postgres tools.
Radek Brich <radek.brich@devl.cz>
parents:
diff changeset
   103
        try:
eaae9539e910 Postgres tools.
Radek Brich <radek.brich@devl.cz>
parents:
diff changeset
   104
            return super(Cursor, self).execute(query, args)
eaae9539e910 Postgres tools.
Radek Brich <radek.brich@devl.cz>
parents:
diff changeset
   105
        finally:
1
ddce8990b976 Fix pgmanager logging in Python3.
Radek Brich <radek.brich@devl.cz>
parents: 0
diff changeset
   106
            log.debug(self.query.decode('utf8'))
0
eaae9539e910 Postgres tools.
Radek Brich <radek.brich@devl.cz>
parents:
diff changeset
   107
eaae9539e910 Postgres tools.
Radek Brich <radek.brich@devl.cz>
parents:
diff changeset
   108
    def callproc(self, procname, args=None):
eaae9539e910 Postgres tools.
Radek Brich <radek.brich@devl.cz>
parents:
diff changeset
   109
        try:
eaae9539e910 Postgres tools.
Radek Brich <radek.brich@devl.cz>
parents:
diff changeset
   110
            return super(Cursor, self).callproc(procname, args)
eaae9539e910 Postgres tools.
Radek Brich <radek.brich@devl.cz>
parents:
diff changeset
   111
        finally:
1
ddce8990b976 Fix pgmanager logging in Python3.
Radek Brich <radek.brich@devl.cz>
parents: 0
diff changeset
   112
            log.debug(self.query.decode('utf8'))
0
eaae9539e910 Postgres tools.
Radek Brich <radek.brich@devl.cz>
parents:
diff changeset
   113
eaae9539e910 Postgres tools.
Radek Brich <radek.brich@devl.cz>
parents:
diff changeset
   114
    def row_dict(self, row, lstrip=None):
eaae9539e910 Postgres tools.
Radek Brich <radek.brich@devl.cz>
parents:
diff changeset
   115
        adjustname = lambda a: a
eaae9539e910 Postgres tools.
Radek Brich <radek.brich@devl.cz>
parents:
diff changeset
   116
        if lstrip:
eaae9539e910 Postgres tools.
Radek Brich <radek.brich@devl.cz>
parents:
diff changeset
   117
            adjustname = lambda a: a.lstrip(lstrip)
eaae9539e910 Postgres tools.
Radek Brich <radek.brich@devl.cz>
parents:
diff changeset
   118
        return RowDict(zip([adjustname(desc[0]) for desc in self.description], row))
eaae9539e910 Postgres tools.
Radek Brich <radek.brich@devl.cz>
parents:
diff changeset
   119
eaae9539e910 Postgres tools.
Radek Brich <radek.brich@devl.cz>
parents:
diff changeset
   120
    def fetchone_dict(self, lstrip=None):
eaae9539e910 Postgres tools.
Radek Brich <radek.brich@devl.cz>
parents:
diff changeset
   121
        row = super(Cursor, self).fetchone()
eaae9539e910 Postgres tools.
Radek Brich <radek.brich@devl.cz>
parents:
diff changeset
   122
        if not row:
eaae9539e910 Postgres tools.
Radek Brich <radek.brich@devl.cz>
parents:
diff changeset
   123
            return row
eaae9539e910 Postgres tools.
Radek Brich <radek.brich@devl.cz>
parents:
diff changeset
   124
        return self.row_dict(row, lstrip)
eaae9539e910 Postgres tools.
Radek Brich <radek.brich@devl.cz>
parents:
diff changeset
   125
eaae9539e910 Postgres tools.
Radek Brich <radek.brich@devl.cz>
parents:
diff changeset
   126
    def fetchall_dict(self, lstrip=None):
eaae9539e910 Postgres tools.
Radek Brich <radek.brich@devl.cz>
parents:
diff changeset
   127
        rows = super(Cursor, self).fetchall()
eaae9539e910 Postgres tools.
Radek Brich <radek.brich@devl.cz>
parents:
diff changeset
   128
        return [self.row_dict(row, lstrip) for row in rows]
eaae9539e910 Postgres tools.
Radek Brich <radek.brich@devl.cz>
parents:
diff changeset
   129
eaae9539e910 Postgres tools.
Radek Brich <radek.brich@devl.cz>
parents:
diff changeset
   130
eaae9539e910 Postgres tools.
Radek Brich <radek.brich@devl.cz>
parents:
diff changeset
   131
class Connection(psycopg2.extensions.connection):
eaae9539e910 Postgres tools.
Radek Brich <radek.brich@devl.cz>
parents:
diff changeset
   132
eaae9539e910 Postgres tools.
Radek Brich <radek.brich@devl.cz>
parents:
diff changeset
   133
    def cursor(self, name=None):
eaae9539e910 Postgres tools.
Radek Brich <radek.brich@devl.cz>
parents:
diff changeset
   134
        if name is None:
eaae9539e910 Postgres tools.
Radek Brich <radek.brich@devl.cz>
parents:
diff changeset
   135
            return super(Connection, self).cursor(cursor_factory=Cursor)
eaae9539e910 Postgres tools.
Radek Brich <radek.brich@devl.cz>
parents:
diff changeset
   136
        else:
eaae9539e910 Postgres tools.
Radek Brich <radek.brich@devl.cz>
parents:
diff changeset
   137
            return super(Connection, self).cursor(name, cursor_factory=Cursor)
eaae9539e910 Postgres tools.
Radek Brich <radek.brich@devl.cz>
parents:
diff changeset
   138
eaae9539e910 Postgres tools.
Radek Brich <radek.brich@devl.cz>
parents:
diff changeset
   139
eaae9539e910 Postgres tools.
Radek Brich <radek.brich@devl.cz>
parents:
diff changeset
   140
class PgManager:
eaae9539e910 Postgres tools.
Radek Brich <radek.brich@devl.cz>
parents:
diff changeset
   141
eaae9539e910 Postgres tools.
Radek Brich <radek.brich@devl.cz>
parents:
diff changeset
   142
    def __init__(self):
eaae9539e910 Postgres tools.
Radek Brich <radek.brich@devl.cz>
parents:
diff changeset
   143
        self.conn_known = {}  # available connections
eaae9539e910 Postgres tools.
Radek Brich <radek.brich@devl.cz>
parents:
diff changeset
   144
        self.conn_pool = {}
eaae9539e910 Postgres tools.
Radek Brich <radek.brich@devl.cz>
parents:
diff changeset
   145
        self.lock = threading.Lock()
eaae9539e910 Postgres tools.
Radek Brich <radek.brich@devl.cz>
parents:
diff changeset
   146
eaae9539e910 Postgres tools.
Radek Brich <radek.brich@devl.cz>
parents:
diff changeset
   147
    def __del__(self):
eaae9539e910 Postgres tools.
Radek Brich <radek.brich@devl.cz>
parents:
diff changeset
   148
        for conn in tuple(self.conn_known.keys()):
eaae9539e910 Postgres tools.
Radek Brich <radek.brich@devl.cz>
parents:
diff changeset
   149
            self.destroy_conn(conn)
eaae9539e910 Postgres tools.
Radek Brich <radek.brich@devl.cz>
parents:
diff changeset
   150
eaae9539e910 Postgres tools.
Radek Brich <radek.brich@devl.cz>
parents:
diff changeset
   151
    def create_conn(self, name='default', isolation_level=None, dsn=None, **kw):
eaae9539e910 Postgres tools.
Radek Brich <radek.brich@devl.cz>
parents:
diff changeset
   152
        '''Create named connection.'''
eaae9539e910 Postgres tools.
Radek Brich <radek.brich@devl.cz>
parents:
diff changeset
   153
        if name in self.conn_known:
2
efee419b7a2d Fix pgmanager exceptions.
Radek Brich <radek.brich@devl.cz>
parents: 0
diff changeset
   154
            raise PgManagerError('Connection name "%s" already registered.' % name)
0
eaae9539e910 Postgres tools.
Radek Brich <radek.brich@devl.cz>
parents:
diff changeset
   155
eaae9539e910 Postgres tools.
Radek Brich <radek.brich@devl.cz>
parents:
diff changeset
   156
        if dsn is None:
eaae9539e910 Postgres tools.
Radek Brich <radek.brich@devl.cz>
parents:
diff changeset
   157
            dsn = ' '.join([x[0]+'='+str(x[1]) for x in kw.items()])
eaae9539e910 Postgres tools.
Radek Brich <radek.brich@devl.cz>
parents:
diff changeset
   158
eaae9539e910 Postgres tools.
Radek Brich <radek.brich@devl.cz>
parents:
diff changeset
   159
        isolation_level = self._normalize_isolation_level(isolation_level)
eaae9539e910 Postgres tools.
Radek Brich <radek.brich@devl.cz>
parents:
diff changeset
   160
        ci = ConnectionInfo(dsn, isolation_level)
eaae9539e910 Postgres tools.
Radek Brich <radek.brich@devl.cz>
parents:
diff changeset
   161
eaae9539e910 Postgres tools.
Radek Brich <radek.brich@devl.cz>
parents:
diff changeset
   162
        self.conn_known[name] = ci
eaae9539e910 Postgres tools.
Radek Brich <radek.brich@devl.cz>
parents:
diff changeset
   163
        self.conn_pool[name] = []
eaae9539e910 Postgres tools.
Radek Brich <radek.brich@devl.cz>
parents:
diff changeset
   164
eaae9539e910 Postgres tools.
Radek Brich <radek.brich@devl.cz>
parents:
diff changeset
   165
    def close_conn(self, name='default'):
eaae9539e910 Postgres tools.
Radek Brich <radek.brich@devl.cz>
parents:
diff changeset
   166
        '''Close all connections of given name.
eaae9539e910 Postgres tools.
Radek Brich <radek.brich@devl.cz>
parents:
diff changeset
   167
        
eaae9539e910 Postgres tools.
Radek Brich <radek.brich@devl.cz>
parents:
diff changeset
   168
        Connection credentials are still saved.
eaae9539e910 Postgres tools.
Radek Brich <radek.brich@devl.cz>
parents:
diff changeset
   169
        
eaae9539e910 Postgres tools.
Radek Brich <radek.brich@devl.cz>
parents:
diff changeset
   170
        '''
eaae9539e910 Postgres tools.
Radek Brich <radek.brich@devl.cz>
parents:
diff changeset
   171
        while len(self.conn_pool[name]):
eaae9539e910 Postgres tools.
Radek Brich <radek.brich@devl.cz>
parents:
diff changeset
   172
            conn = self.conn_pool[name].pop()
eaae9539e910 Postgres tools.
Radek Brich <radek.brich@devl.cz>
parents:
diff changeset
   173
            conn.close()
eaae9539e910 Postgres tools.
Radek Brich <radek.brich@devl.cz>
parents:
diff changeset
   174
eaae9539e910 Postgres tools.
Radek Brich <radek.brich@devl.cz>
parents:
diff changeset
   175
    def destroy_conn(self, name='default'):
eaae9539e910 Postgres tools.
Radek Brich <radek.brich@devl.cz>
parents:
diff changeset
   176
        '''Destroy connection.
eaae9539e910 Postgres tools.
Radek Brich <radek.brich@devl.cz>
parents:
diff changeset
   177
        
eaae9539e910 Postgres tools.
Radek Brich <radek.brich@devl.cz>
parents:
diff changeset
   178
        Counterpart of create_conn.
eaae9539e910 Postgres tools.
Radek Brich <radek.brich@devl.cz>
parents:
diff changeset
   179
        
eaae9539e910 Postgres tools.
Radek Brich <radek.brich@devl.cz>
parents:
diff changeset
   180
        '''
eaae9539e910 Postgres tools.
Radek Brich <radek.brich@devl.cz>
parents:
diff changeset
   181
        if not name in self.conn_known:
2
efee419b7a2d Fix pgmanager exceptions.
Radek Brich <radek.brich@devl.cz>
parents: 0
diff changeset
   182
            raise PgManagerError('Connection name "%s" not registered.' % name)
0
eaae9539e910 Postgres tools.
Radek Brich <radek.brich@devl.cz>
parents:
diff changeset
   183
eaae9539e910 Postgres tools.
Radek Brich <radek.brich@devl.cz>
parents:
diff changeset
   184
        self.close_conn(name)
eaae9539e910 Postgres tools.
Radek Brich <radek.brich@devl.cz>
parents:
diff changeset
   185
eaae9539e910 Postgres tools.
Radek Brich <radek.brich@devl.cz>
parents:
diff changeset
   186
        del self.conn_known[name]
eaae9539e910 Postgres tools.
Radek Brich <radek.brich@devl.cz>
parents:
diff changeset
   187
        del self.conn_pool[name]
eaae9539e910 Postgres tools.
Radek Brich <radek.brich@devl.cz>
parents:
diff changeset
   188
eaae9539e910 Postgres tools.
Radek Brich <radek.brich@devl.cz>
parents:
diff changeset
   189
    def get_conn(self, name='default'):
eaae9539e910 Postgres tools.
Radek Brich <radek.brich@devl.cz>
parents:
diff changeset
   190
        '''Get connection of name 'name' from pool.'''
eaae9539e910 Postgres tools.
Radek Brich <radek.brich@devl.cz>
parents:
diff changeset
   191
        self.lock.acquire()
eaae9539e910 Postgres tools.
Radek Brich <radek.brich@devl.cz>
parents:
diff changeset
   192
        try:
eaae9539e910 Postgres tools.
Radek Brich <radek.brich@devl.cz>
parents:
diff changeset
   193
            if not name in self.conn_known:
2
efee419b7a2d Fix pgmanager exceptions.
Radek Brich <radek.brich@devl.cz>
parents: 0
diff changeset
   194
                raise PgManagerError("Connection name '%s' not registered." % name)
0
eaae9539e910 Postgres tools.
Radek Brich <radek.brich@devl.cz>
parents:
diff changeset
   195
eaae9539e910 Postgres tools.
Radek Brich <radek.brich@devl.cz>
parents:
diff changeset
   196
            conn = None
eaae9539e910 Postgres tools.
Radek Brich <radek.brich@devl.cz>
parents:
diff changeset
   197
            while len(self.conn_pool[name]) and conn is None:
eaae9539e910 Postgres tools.
Radek Brich <radek.brich@devl.cz>
parents:
diff changeset
   198
                conn = self.conn_pool[name].pop()
eaae9539e910 Postgres tools.
Radek Brich <radek.brich@devl.cz>
parents:
diff changeset
   199
                if conn.closed:
eaae9539e910 Postgres tools.
Radek Brich <radek.brich@devl.cz>
parents:
diff changeset
   200
                    conn = None
eaae9539e910 Postgres tools.
Radek Brich <radek.brich@devl.cz>
parents:
diff changeset
   201
eaae9539e910 Postgres tools.
Radek Brich <radek.brich@devl.cz>
parents:
diff changeset
   202
            if conn is None:
eaae9539e910 Postgres tools.
Radek Brich <radek.brich@devl.cz>
parents:
diff changeset
   203
                ci = self.conn_known[name]
eaae9539e910 Postgres tools.
Radek Brich <radek.brich@devl.cz>
parents:
diff changeset
   204
                conn = psycopg2.connect(ci.dsn, connection_factory=Connection)
eaae9539e910 Postgres tools.
Radek Brich <radek.brich@devl.cz>
parents:
diff changeset
   205
                if not ci.isolation_level is None:
eaae9539e910 Postgres tools.
Radek Brich <radek.brich@devl.cz>
parents:
diff changeset
   206
                    conn.set_isolation_level(ci.isolation_level)
eaae9539e910 Postgres tools.
Radek Brich <radek.brich@devl.cz>
parents:
diff changeset
   207
                if ci.init_statement:
eaae9539e910 Postgres tools.
Radek Brich <radek.brich@devl.cz>
parents:
diff changeset
   208
                    curs = conn.cursor()
eaae9539e910 Postgres tools.
Radek Brich <radek.brich@devl.cz>
parents:
diff changeset
   209
                    curs.execute(ci.init_statement)
eaae9539e910 Postgres tools.
Radek Brich <radek.brich@devl.cz>
parents:
diff changeset
   210
                    curs.close()
eaae9539e910 Postgres tools.
Radek Brich <radek.brich@devl.cz>
parents:
diff changeset
   211
        finally:
eaae9539e910 Postgres tools.
Radek Brich <radek.brich@devl.cz>
parents:
diff changeset
   212
            self.lock.release()
eaae9539e910 Postgres tools.
Radek Brich <radek.brich@devl.cz>
parents:
diff changeset
   213
        return conn
eaae9539e910 Postgres tools.
Radek Brich <radek.brich@devl.cz>
parents:
diff changeset
   214
eaae9539e910 Postgres tools.
Radek Brich <radek.brich@devl.cz>
parents:
diff changeset
   215
    def put_conn(self, conn, name='default'):
eaae9539e910 Postgres tools.
Radek Brich <radek.brich@devl.cz>
parents:
diff changeset
   216
        '''Put connection back to pool.
eaae9539e910 Postgres tools.
Radek Brich <radek.brich@devl.cz>
parents:
diff changeset
   217
        
eaae9539e910 Postgres tools.
Radek Brich <radek.brich@devl.cz>
parents:
diff changeset
   218
        Name must be same as used for get_conn,
eaae9539e910 Postgres tools.
Radek Brich <radek.brich@devl.cz>
parents:
diff changeset
   219
        otherwise things become broken.
eaae9539e910 Postgres tools.
Radek Brich <radek.brich@devl.cz>
parents:
diff changeset
   220
        
eaae9539e910 Postgres tools.
Radek Brich <radek.brich@devl.cz>
parents:
diff changeset
   221
        '''
eaae9539e910 Postgres tools.
Radek Brich <radek.brich@devl.cz>
parents:
diff changeset
   222
        self.lock.acquire()
eaae9539e910 Postgres tools.
Radek Brich <radek.brich@devl.cz>
parents:
diff changeset
   223
        try:
eaae9539e910 Postgres tools.
Radek Brich <radek.brich@devl.cz>
parents:
diff changeset
   224
            if not name in self.conn_known:
2
efee419b7a2d Fix pgmanager exceptions.
Radek Brich <radek.brich@devl.cz>
parents: 0
diff changeset
   225
                raise PgManagerError("Connection name '%s' not registered." % name)
0
eaae9539e910 Postgres tools.
Radek Brich <radek.brich@devl.cz>
parents:
diff changeset
   226
eaae9539e910 Postgres tools.
Radek Brich <radek.brich@devl.cz>
parents:
diff changeset
   227
            if len(self.conn_pool[name]) >= self.conn_known[name].keep_open:
eaae9539e910 Postgres tools.
Radek Brich <radek.brich@devl.cz>
parents:
diff changeset
   228
                conn.close()
eaae9539e910 Postgres tools.
Radek Brich <radek.brich@devl.cz>
parents:
diff changeset
   229
                return
eaae9539e910 Postgres tools.
Radek Brich <radek.brich@devl.cz>
parents:
diff changeset
   230
eaae9539e910 Postgres tools.
Radek Brich <radek.brich@devl.cz>
parents:
diff changeset
   231
            if conn.get_transaction_status() == psycopg2.extensions.TRANSACTION_STATUS_UNKNOWN:
eaae9539e910 Postgres tools.
Radek Brich <radek.brich@devl.cz>
parents:
diff changeset
   232
                conn.close()
eaae9539e910 Postgres tools.
Radek Brich <radek.brich@devl.cz>
parents:
diff changeset
   233
                return
eaae9539e910 Postgres tools.
Radek Brich <radek.brich@devl.cz>
parents:
diff changeset
   234
eaae9539e910 Postgres tools.
Radek Brich <radek.brich@devl.cz>
parents:
diff changeset
   235
            # connection returned to the pool must not be in transaction
eaae9539e910 Postgres tools.
Radek Brich <radek.brich@devl.cz>
parents:
diff changeset
   236
            if conn.get_transaction_status() != psycopg2.extensions.TRANSACTION_STATUS_IDLE:
eaae9539e910 Postgres tools.
Radek Brich <radek.brich@devl.cz>
parents:
diff changeset
   237
                conn.rollback()
eaae9539e910 Postgres tools.
Radek Brich <radek.brich@devl.cz>
parents:
diff changeset
   238
eaae9539e910 Postgres tools.
Radek Brich <radek.brich@devl.cz>
parents:
diff changeset
   239
            self.conn_pool[name].append(conn)
eaae9539e910 Postgres tools.
Radek Brich <radek.brich@devl.cz>
parents:
diff changeset
   240
        finally:
eaae9539e910 Postgres tools.
Radek Brich <radek.brich@devl.cz>
parents:
diff changeset
   241
            self.lock.release()
eaae9539e910 Postgres tools.
Radek Brich <radek.brich@devl.cz>
parents:
diff changeset
   242
eaae9539e910 Postgres tools.
Radek Brich <radek.brich@devl.cz>
parents:
diff changeset
   243
    @contextmanager
eaae9539e910 Postgres tools.
Radek Brich <radek.brich@devl.cz>
parents:
diff changeset
   244
    def cursor(self, name='default'):
eaae9539e910 Postgres tools.
Radek Brich <radek.brich@devl.cz>
parents:
diff changeset
   245
        '''Cursor context.
eaae9539e910 Postgres tools.
Radek Brich <radek.brich@devl.cz>
parents:
diff changeset
   246
        
eaae9539e910 Postgres tools.
Radek Brich <radek.brich@devl.cz>
parents:
diff changeset
   247
        Uses any connection of name 'name' from pool
eaae9539e910 Postgres tools.
Radek Brich <radek.brich@devl.cz>
parents:
diff changeset
   248
        and returns cursor for that connection.
eaae9539e910 Postgres tools.
Radek Brich <radek.brich@devl.cz>
parents:
diff changeset
   249
        
eaae9539e910 Postgres tools.
Radek Brich <radek.brich@devl.cz>
parents:
diff changeset
   250
        '''
eaae9539e910 Postgres tools.
Radek Brich <radek.brich@devl.cz>
parents:
diff changeset
   251
        conn = self.get_conn(name)
eaae9539e910 Postgres tools.
Radek Brich <radek.brich@devl.cz>
parents:
diff changeset
   252
eaae9539e910 Postgres tools.
Radek Brich <radek.brich@devl.cz>
parents:
diff changeset
   253
        try:
eaae9539e910 Postgres tools.
Radek Brich <radek.brich@devl.cz>
parents:
diff changeset
   254
            curs = conn.cursor()
eaae9539e910 Postgres tools.
Radek Brich <radek.brich@devl.cz>
parents:
diff changeset
   255
            yield curs
eaae9539e910 Postgres tools.
Radek Brich <radek.brich@devl.cz>
parents:
diff changeset
   256
        finally:
eaae9539e910 Postgres tools.
Radek Brich <radek.brich@devl.cz>
parents:
diff changeset
   257
            curs.close()
eaae9539e910 Postgres tools.
Radek Brich <radek.brich@devl.cz>
parents:
diff changeset
   258
            self.put_conn(conn, name)
eaae9539e910 Postgres tools.
Radek Brich <radek.brich@devl.cz>
parents:
diff changeset
   259
eaae9539e910 Postgres tools.
Radek Brich <radek.brich@devl.cz>
parents:
diff changeset
   260
    def wait_for_notify(self, name='default', timeout=5):
eaae9539e910 Postgres tools.
Radek Brich <radek.brich@devl.cz>
parents:
diff changeset
   261
        '''Wait for asynchronous notifies, return the last one.
eaae9539e910 Postgres tools.
Radek Brich <radek.brich@devl.cz>
parents:
diff changeset
   262
        
eaae9539e910 Postgres tools.
Radek Brich <radek.brich@devl.cz>
parents:
diff changeset
   263
        Returns None on timeout.
eaae9539e910 Postgres tools.
Radek Brich <radek.brich@devl.cz>
parents:
diff changeset
   264
        
eaae9539e910 Postgres tools.
Radek Brich <radek.brich@devl.cz>
parents:
diff changeset
   265
        '''
eaae9539e910 Postgres tools.
Radek Brich <radek.brich@devl.cz>
parents:
diff changeset
   266
        conn = self.get_conn(name)
eaae9539e910 Postgres tools.
Radek Brich <radek.brich@devl.cz>
parents:
diff changeset
   267
        
eaae9539e910 Postgres tools.
Radek Brich <radek.brich@devl.cz>
parents:
diff changeset
   268
        try:
eaae9539e910 Postgres tools.
Radek Brich <radek.brich@devl.cz>
parents:
diff changeset
   269
            # any residual notify?
eaae9539e910 Postgres tools.
Radek Brich <radek.brich@devl.cz>
parents:
diff changeset
   270
            # then return it, that should not break anything
eaae9539e910 Postgres tools.
Radek Brich <radek.brich@devl.cz>
parents:
diff changeset
   271
            if conn.notifies:
eaae9539e910 Postgres tools.
Radek Brich <radek.brich@devl.cz>
parents:
diff changeset
   272
                return conn.notifies.pop()
eaae9539e910 Postgres tools.
Radek Brich <radek.brich@devl.cz>
parents:
diff changeset
   273
eaae9539e910 Postgres tools.
Radek Brich <radek.brich@devl.cz>
parents:
diff changeset
   274
            if select.select([conn], [], [], timeout) == ([], [], []):
eaae9539e910 Postgres tools.
Radek Brich <radek.brich@devl.cz>
parents:
diff changeset
   275
                # timeout
eaae9539e910 Postgres tools.
Radek Brich <radek.brich@devl.cz>
parents:
diff changeset
   276
                return None
eaae9539e910 Postgres tools.
Radek Brich <radek.brich@devl.cz>
parents:
diff changeset
   277
            else:
eaae9539e910 Postgres tools.
Radek Brich <radek.brich@devl.cz>
parents:
diff changeset
   278
                conn.poll()
eaae9539e910 Postgres tools.
Radek Brich <radek.brich@devl.cz>
parents:
diff changeset
   279
eaae9539e910 Postgres tools.
Radek Brich <radek.brich@devl.cz>
parents:
diff changeset
   280
                # return just the last notify (we do not care for older ones)
eaae9539e910 Postgres tools.
Radek Brich <radek.brich@devl.cz>
parents:
diff changeset
   281
                if conn.notifies:
eaae9539e910 Postgres tools.
Radek Brich <radek.brich@devl.cz>
parents:
diff changeset
   282
                    return conn.notifies.pop()
eaae9539e910 Postgres tools.
Radek Brich <radek.brich@devl.cz>
parents:
diff changeset
   283
                return None
eaae9539e910 Postgres tools.
Radek Brich <radek.brich@devl.cz>
parents:
diff changeset
   284
        finally:
eaae9539e910 Postgres tools.
Radek Brich <radek.brich@devl.cz>
parents:
diff changeset
   285
            # clean notifies
eaae9539e910 Postgres tools.
Radek Brich <radek.brich@devl.cz>
parents:
diff changeset
   286
            while conn.notifies:
eaae9539e910 Postgres tools.
Radek Brich <radek.brich@devl.cz>
parents:
diff changeset
   287
                conn.notifies.pop()
eaae9539e910 Postgres tools.
Radek Brich <radek.brich@devl.cz>
parents:
diff changeset
   288
            self.put_conn(conn, name)
eaae9539e910 Postgres tools.
Radek Brich <radek.brich@devl.cz>
parents:
diff changeset
   289
eaae9539e910 Postgres tools.
Radek Brich <radek.brich@devl.cz>
parents:
diff changeset
   290
    def _normalize_isolation_level(self, level):
eaae9539e910 Postgres tools.
Radek Brich <radek.brich@devl.cz>
parents:
diff changeset
   291
        if type(level) == str:
eaae9539e910 Postgres tools.
Radek Brich <radek.brich@devl.cz>
parents:
diff changeset
   292
            if level.lower() == 'autocommit':
eaae9539e910 Postgres tools.
Radek Brich <radek.brich@devl.cz>
parents:
diff changeset
   293
                return psycopg2.extensions.ISOLATION_LEVEL_AUTOCOMMIT
eaae9539e910 Postgres tools.
Radek Brich <radek.brich@devl.cz>
parents:
diff changeset
   294
            if level.lower() == 'read_committed':
eaae9539e910 Postgres tools.
Radek Brich <radek.brich@devl.cz>
parents:
diff changeset
   295
                return psycopg2.extensions.ISOLATION_LEVEL_READ_COMMITTED
eaae9539e910 Postgres tools.
Radek Brich <radek.brich@devl.cz>
parents:
diff changeset
   296
            if level.lower() == 'serializable':
eaae9539e910 Postgres tools.
Radek Brich <radek.brich@devl.cz>
parents:
diff changeset
   297
                return psycopg2.extensions.ISOLATION_LEVEL_SERIALIZABLE
eaae9539e910 Postgres tools.
Radek Brich <radek.brich@devl.cz>
parents:
diff changeset
   298
            raise PgManagerError('Unknown isolation level name: "%s"', level)
eaae9539e910 Postgres tools.
Radek Brich <radek.brich@devl.cz>
parents:
diff changeset
   299
        return level
eaae9539e910 Postgres tools.
Radek Brich <radek.brich@devl.cz>
parents:
diff changeset
   300
eaae9539e910 Postgres tools.
Radek Brich <radek.brich@devl.cz>
parents:
diff changeset
   301
eaae9539e910 Postgres tools.
Radek Brich <radek.brich@devl.cz>
parents:
diff changeset
   302
try:
eaae9539e910 Postgres tools.
Radek Brich <radek.brich@devl.cz>
parents:
diff changeset
   303
    NullHandler = logging.NullHandler
eaae9539e910 Postgres tools.
Radek Brich <radek.brich@devl.cz>
parents:
diff changeset
   304
except AttributeError:
eaae9539e910 Postgres tools.
Radek Brich <radek.brich@devl.cz>
parents:
diff changeset
   305
    class NullHandler(logging.Handler):
eaae9539e910 Postgres tools.
Radek Brich <radek.brich@devl.cz>
parents:
diff changeset
   306
        def emit(self, record):
eaae9539e910 Postgres tools.
Radek Brich <radek.brich@devl.cz>
parents:
diff changeset
   307
            pass
eaae9539e910 Postgres tools.
Radek Brich <radek.brich@devl.cz>
parents:
diff changeset
   308
eaae9539e910 Postgres tools.
Radek Brich <radek.brich@devl.cz>
parents:
diff changeset
   309
eaae9539e910 Postgres tools.
Radek Brich <radek.brich@devl.cz>
parents:
diff changeset
   310
log = logging.getLogger("pgmanager")
eaae9539e910 Postgres tools.
Radek Brich <radek.brich@devl.cz>
parents:
diff changeset
   311
log.addHandler(NullHandler())
eaae9539e910 Postgres tools.
Radek Brich <radek.brich@devl.cz>
parents:
diff changeset
   312
eaae9539e910 Postgres tools.
Radek Brich <radek.brich@devl.cz>
parents:
diff changeset
   313
eaae9539e910 Postgres tools.
Radek Brich <radek.brich@devl.cz>
parents:
diff changeset
   314
instance = None
eaae9539e910 Postgres tools.
Radek Brich <radek.brich@devl.cz>
parents:
diff changeset
   315
eaae9539e910 Postgres tools.
Radek Brich <radek.brich@devl.cz>
parents:
diff changeset
   316
eaae9539e910 Postgres tools.
Radek Brich <radek.brich@devl.cz>
parents:
diff changeset
   317
def get_instance():
eaae9539e910 Postgres tools.
Radek Brich <radek.brich@devl.cz>
parents:
diff changeset
   318
    global instance
eaae9539e910 Postgres tools.
Radek Brich <radek.brich@devl.cz>
parents:
diff changeset
   319
    if instance is None:
eaae9539e910 Postgres tools.
Radek Brich <radek.brich@devl.cz>
parents:
diff changeset
   320
        instance = PgManager()
eaae9539e910 Postgres tools.
Radek Brich <radek.brich@devl.cz>
parents:
diff changeset
   321
    return instance
eaae9539e910 Postgres tools.
Radek Brich <radek.brich@devl.cz>
parents:
diff changeset
   322
eaae9539e910 Postgres tools.
Radek Brich <radek.brich@devl.cz>
parents:
diff changeset
   323