This extention makes it possible to hold a global configuration of all needed database, and access this databases using a special context aware named connection pools. The context can help to automate transaction commits and rollbacks and return special connection cursors in the context. See module and class docs for more information.
Python, 532 lines
Download
Copy to clipboard1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249 250 251 252 253 254 255 256 257 258 259 260 261 262 263 264 265 266 267 268 269 270 271 272 273 274 275 276 277 278 279 280 281 282 283 284 285 286 287 288 289 290 291 292 293 294 295 296 297 298 299 300 301 302 303 304 305 306 307 308 309 310 311 312 313 314 315 316 317 318 319 320 321 322 323 324 325 326 327 328 329 330 331 332 333 334 335 336 337 338 339 340 341 342 343 344 345 346 347 348 349 350 351 352 353 354 355 356 357 358 359 360 361 362 363 364 365 366 367 368 369 370 371 372 373 374 375 376 377 378 379 380 381 382 383 384 385 386 387 388 389 390 391 392 393 394 395 396 397 398 399 400 401 402 403 404 405 406 407 408 409 410 411 412 413 414 415 416 417 418 419 420 421 422 423 424 425 426 427 428 429 430 431 432 433 434 435 436 437 438 439 440 441 442 443 444 445 446 447 448 449 450 451 452 453 454 455 456 457 458 459 460 461 462 463 464 465 466 467 468 469 470 471 472 473 474 475 476 477 478 479 480 481 482 483 484 485 486 487 488 489 490 491 492 493 494 495 496 497 498 499 500 501 502 503 504 505 506 507 508 509 510 511 512 513 514 515 516 517 518 519 520 521 522 523 524 525 526 527 528 529 530 531 532 | # -*- coding: utf-8 -*-
# file: helpers/database.py
"""
This module provides a possibility to process a configuration file
and load database connection information from that configuration.
To load configuration and open the configured connections
the following code can be executed::
from helpers import database
database.fileConfig(CONF_FILE_NAME)
Here CONF_FILE_NAME is the name of the configuration file.
See fileConfig(fileName) for more details.
All registered connection pools are closed automatically on exiting the script
(using 'atexit' events)
A public databaseConfig variable is an instance of a private _DatabaseConfig class,
and holds the connection pools.
Possible examples of the usage are::
from helpers import database
# find the database.conf file and use it as a database configuration
database.fileConfig('database.conf');
database.get_connection_pool_names() # get all defined pool names
pool_name = 'videoflip'
database.has_connection_pool(pool_name) # check if the pool with pool_name is defined
videoflip_pool = database.get_connection_pool(pool_name) # get videoflip pool
with videoflip_pool.context('testing cursor', ContextType.TRANSACTIONAL ) as cur:
cur.execute('SELECT 1')
r = cur.fetchone()
print r
Decorator 'connectionPoolAware(pool_name, argument_name=None)' can be used to mark methods
with the defined parameters, to pass the named connection pool to the method every time it is called.
But this approach seems to be not so interesting, as one always can get a needed pool
directry with database.get_connection_pool('pool_name')
See documentation for class 'PersistentConnectionPoolWithContext' that is returned by all the methods
returning connection pool of this module.
"""
import logging
from contextlib import contextmanager
import psycopg2
from psycopg2.pool import PersistentConnectionPool
from psycopg2.extras import DictCursor
import ConfigParser
import threading
import os.path
import atexit
import collections
__all__ = [ 'databaseConfig',
'fileConfig',
'DEFAULT_DATABASE_CONFIGURATION_FILE_NAME',
'ContextType',
'PersistentConnectionPoolWithContext',
'get_connection_pool',
'has_connection_pool',
'get_connection_pool_names',
'InvalidMethodException',
'InvalidDatabaseConfigurationException' ]
class Error(Exception):
pass
class InvalidMethodException(Error):
pass
class InvalidDatabaseConfigurationException(Error):
pass
########################################################################################
#
# Find file on the python path (http://code.activestate.com/recipes/52224/)
# This code is actually to be located in a utility library, just moved it here
# to make the module code independent
#
########################################################################################
def _find(path, matchFunc):
apath = os.path.abspath(path)
if matchFunc(apath):
return apath
for dirname in sys.path:
candidate = os.path.join(dirname, path)
if matchFunc(candidate):
return candidate
raise IOError("Can't find file %s" % path)
def find_file(path):
"""Find a file in the python system path
"""
return _find(path,matchFunc=os.path.isfile)
def find_dir(path):
"""Find a directory in the pathon system path
"""
return _find(path,matchFunc=os.path.isdir)
########################################################################################
#
# EasySafeConfigParser class
#
########################################################################################
class EasySafeConfigParser(SafeConfigParser):
'This implementation addes get_default method, that does not raise NoOptionError and returning default value in case the option does not exist'
def get_default(self, section, option, default=None):
try :
return SafeConfigParser.get(self, section, option )
except NoOptionError :
return default
########################################################################################
databaseConfig = None
_logger = logging.getLogger('helpers.database')
_database_config_lock = threading.RLock()
DEFAULT_DATABASE_CONFIGURATION_FILE_NAME='database.conf'
def fileConfig(file_name=DEFAULT_DATABASE_CONFIGURATION_FILE_NAME):
"""
Configure database pools from the given configuration file.
Database configuration section names should start with prefix 'database_'
a name after that prefix is used as name for the registered connection pool.
Configuration file should be named database.conf (default name) and located
in the current directory or if not found there in the PYTHON_PATH directories.
This makes it possible to have one global configuration file, and then
add a local one for debugging purposes. Note, that config data is not being
merged and the file, first found is being used as a config file.
Typical configuration file can look like that::
[database_testdb1]
dbname=testdb1
host=localhost
port=5453
user=test
password=test
encoding=latin1
statement_timeout=0
[database_testdb2]
dbname=testdb2
host=localhost
port=5454
user=test
password=test
After the database access is configured, you can access database pools like::
from myvideo.helpers import database
database.fileConfig()
testdb2_pool = database.get_connection_pool('testdb2')
"""
# find configuration file
file_name = find_file( file_name )
# create configuration
global databaseConfig #IGNORE:W0603
databaseConfig = _DatabaseConfig(os.path.abspath(file_name))
def get_connection_pool(pool_name, config_file_name=DEFAULT_DATABASE_CONFIGURATION_FILE_NAME):
"""Returns a connection pool with a given name
@param pool_name: a connection pool name to look up
"""
if databaseConfig is None :
fileConfig(config_file_name)
return databaseConfig.get_connection_pool(pool_name)
def get_connection_pool_set(pool_name, config_file_name=DEFAULT_DATABASE_CONFIGURATION_FILE_NAME):
"""Returns a set of connection pools with a given name
@param pool_name: a connection pool name to look up
"""
if databaseConfig is None :
fileConfig(config_file_name)
return databaseConfig.get_connection_pool_set(pool_name)
def has_connection_pool(pool_name, config_file_name=DEFAULT_DATABASE_CONFIGURATION_FILE_NAME):
"""True if a connection pool with a given name is registered
@param pool_name: a connection pool name to look up
"""
if databaseConfig is None :
fileConfig(config_file_name)
return databaseConfig.has_connection_pool(pool_name)
def get_connection_pool_names(config_file_name=DEFAULT_DATABASE_CONFIGURATION_FILE_NAME):
"""Returns all registered connection pool names"""
if databaseConfig is None :
fileConfig(config_file_name)
return databaseConfig.get_connection_pool_names()
class _DatabaseConfig(object):
"""Private class holding database configuration and registered connection pools"""
__slots__ = ('logger', 'cp', 'fileName', '_pool_map', '_pool_config_section_map')
# define an atexit function to be called at exit for the created pools
def _atexit_close_pool(self, connection_pool):
if connection_pool:
if not connection_pool.closed:
self.logger.debug('Closing all connections of the connection pool %s', connection_pool.pool_name)
connection_pool.closeall()
def __init__(self, fileName):
# Initialize logging
self.logger = logging.getLogger(self.__class__.__name__)
self.logger.debug('Creating DatabaseConfig from file %s', fileName )
self.fileName = fileName
# load configuration
self.cp = EasySafeConfigParser()
self.cp.read(fileName)
with _database_config_lock:
self._pool_config_section_map = {}
self._pool_map = {}
# initialize database connection pools
for section_name in self.cp.sections() :
s = section_name.split('_', 1)
if s[0] != 'database' : continue
pool_name = s[1]
if not pool_name : continue
self._pool_config_section_map[pool_name] = section_name
pool_tuple = tuple( self.database_config_section_to_connection_pool_set(section_name, pool_name) )
self._pool_map[pool_name] = pool_tuple
if not self._pool_map :
self.logger.error('Could not find any database configuration sections in the given configuration file %s', fileName )
def reload(self):
"""
Read database configuration.
Read configuration file and check if some parameters had been changed there.
If changed -- recreate pool and close the old one
Also should recheck if the database connection are alive
"""
# TODO: implement refreshing of the configuration
pass
def database_config_section_to_connection_pool_set(self, section_name, pool_name):
"""Read the given section name from the private config parser and create a corresponding connection pool.
This pool does not try to connect to the database immediately (with minconn=0)
"""
if not pool_name:
pool_name = section_name
if not self.cp.has_section(section_name) :
raise ConfigParser.NoSectionError(section_name)
hosts = self.cp.get_default(section_name, 'host', 'localhost').split(',')
port = int( self.cp.get_default(section_name, 'port', 5432) )
dbname = self.cp.get_default(section_name, 'dbname', 'postgres')
common_dsn = ' '.join( '{0[0]}={0[1]}'.format(item) for item in self.cp.items(section_name) if item[0].lower() not in ('host','dsn','encoding','statement_timeout') )
host_dsn_list = [ ( host, 'host={0} {1}'.format(host.strip(), common_dsn) ) for host in hosts if host.strip() ]
if not host_dsn_list :
raise ConfigParser.NoOptionError('host', section_name)
try :
encoding=self.cp.get(section_name, 'encoding')
except ConfigParser.NoOptionError :
encoding=None
try :
statement_timeout=self.cp.get(section_name, 'statement_timeout')
statement_timeout=int(statement_timeout)
except ConfigParser.NoOptionError :
statement_timeout=None
except ValueError :
statement_timeout=None
self.logger.warning("Database configuration section {0} contains option 'statement_timeout' that should be numeric, but the specified value is '{1}', skipping this option.".format(section_name, statement_timeout) )
for i, ( host, dsn ) in enumerate( host_dsn_list ):
pool_name_i = '{0}_{1}'.format(pool_name, i) if len(host_dsn_list) > 1 else pool_name
connection_pool = PersistentConnectionPoolWithContext(pool_name_i,
minconn=0,
maxconn=5,
host=host,
port=port,
dbname=dbname,
encoding=encoding,
statement_timeout=statement_timeout,
dsn=dsn)
self.logger.debug('Created a connection pool to {0} database with dsn: {1}'.format(pool_name_i, dsn))
atexit.register(self._atexit_close_pool, connection_pool)
yield connection_pool
def get_connection_pool(self, pool_name):
"""Returns a connection pool with a given name
@param pool_name: a connection pool name to look up
"""
with _database_config_lock :
pool_set = self._pool_map[pool_name]
if isinstance(pool_set, collections.Sequence ):
if len(pool_set) != 1:
raise InvalidMethodException('Trying to get a connection pool from defined set of pools')
return pool_set[0]
else:
return pool_set
def get_connection_pool_set(self, pool_name):
"""Returns a set of connection pools with a given name
@param pool_name: a connection pool name to look up
"""
with _database_config_lock :
pool_set = self._pool_map[pool_name]
if isinstance(pool_set, collections.Sequence ):
return pool_set
else :
return tuple( pool_set )
def has_connection_pool(self, pool_name):
"""True if a connection pool with a given name is registered
@param pool_name: a connection pool name to look up
"""
with _database_config_lock :
return pool_name in self._pool_map
def get_connection_pool_names(self):
"""Returns all registered connection pool names"""
with _database_config_lock :
return self._pool_map.keys()
def connectionPoolAware(pool_name, argument_name=None):
"""This decorator redefines the value of the decorated function argument named <pool_name>
and sets it to the connection pool named pool_name defined in the DatabaseConfig class.
Note: this decorator seems to have no practical use, as it is much easier to fetch
a needed connection pool by name using database.get_connection_pool('pool_name')
directly inside the method.
Typical way of decorating the method would be::
@connectionPoolAware(pool_name='slave', argument_name='slavePool')
def test_database_method(slavePool=None)
with slavePool.context('testing connection pool aware method') as cur
cur.execute('SELECT VERSION()')
# and then call the defined method like that
test_database_method()
"""
def decorator(f):
"""Function wrapper, that returns a callable that wraps decorated function 'f' (passed as a parameter)"""
def executor(*argv, **kwargs):
"""Actual wrapper, that adds some new parameters to the decorated function 'f'"""
if databaseConfig :
if databaseConfig.has_connection_pool(pool_name) :
# calculate the argument name to be assigned
n = argument_name if argument_name is not None else pool_name
p = databaseConfig.get_connection_pool( pool_name )
if p :
kwargs[ n ] = p
else :
_logger.error( 'Pool named %s is not initialized', pool_name )
else :
_logger.error( 'Pool named %s is not defined!', pool_name )
else :
_logger.error( 'Database configuration should be initialized with fileConfig(fileName) call.')
f(*argv, **kwargs) #IGNORE:W0142
return executor
return decorator
class ContextType(object):
"""Context type to be used by the PersistentConnectionPoolWithContext.context() method.
The values define bits, for the control of the context, to be returned:
SIMPLE -- does nothing
TRANSACTIONAL -- will commit the results if no error happened and rollback in case of database error
DICT_CURSOR -- returns DictCursor, that can be accessed using column names together with the column indexes.
Cannot be used with NAMED_CURSOR.
NAMED_CURSOR -- return a named cursor, that uses CURSOR internally and is good for fetching large data sets.
Cannot be used with DICT_CURSOR.
CLOSE_CONNECTION -- forces a close of the connection, before returning it to the pool,
actually leading to throwing of that closed connection away.
This can be used in case we always need some fresh connections from the pool.
"""
__slots__ = ()
SIMPLE = 0
TRANSACTIONAL = 1
DICT_CURSOR = 2
NAMED_CURSOR = 4
CLOSE_CONNECTION = 8
class PersistentConnectionPoolWithContext(PersistentConnectionPool):
"""This class extends psycopy2 PersistentConnectionPool class and adds context() method,
that returns a database cursor aware context manager that automatically fetches a connection from the pool,
and, if needed, commits the changes on exiting from the context. In case some exception happens in the context,
the connection is rolled back and exception is thrown further.
"""
def __init__(self, pool_name, minconn, maxconn, host, port, dbname, encoding, statement_timeout, *args, **kwargs):
self.logger = logging.getLogger(self.__class__.__name__)
self.pool_name = pool_name
self.address = ( host, port, )
self.dbname = dbname
self.encoding = encoding
self.statement_timeout = statement_timeout
PersistentConnectionPool.__init__(self, minconn, maxconn, *args, **kwargs) #IGNORE:W0142
def _connect(self, key=None):
conn = PersistentConnectionPool._connect(self, key)
if self.encoding and conn.encoding and conn.encoding.lower() != self.encoding.lower() :
conn.set_client_encoding(self.encoding)
debug_string = 'Pool {0} created new connection with {1} encoding'.format( self.pool_name, conn.encoding )
if self.statement_timeout is not None :
cur = conn.cursor()
try:
# one can improve this by getting an old statement timeout value and then restoring it on
# reterning connection back to the pool
cur.execute('set STATEMENT_TIMEOUT to %s;', ( self.statement_timeout, ))
except psycopg2.Error as e :
self.logger.error('Could not set connection statement_timeout to {0}: {1}'.format( self.statement_timeout, e ) )
conn.rollback()
finally:
cur.close()
del cur
conn.commit()
debug_string += ' and statement_timeout set to {0} ms'.format(self.statement_timeout)
self.logger.debug(debug_string + '.')
del debug_string
return conn
@contextmanager
def context(self, description, contextType=ContextType.DICT_CURSOR, cursorName=None):
"""Return a context of the cursor to be used with the pool.
The following example code can be used to demonstrate the way this context is to be used::
with pool.context('experimenting with pool context', ContextType.TRANSACTIONAL | ContextType.CLOSE_CONNECTION ) as cur :
cur.execute('SELECT 1')
cur.fetchall()
# as we are out of the context, cursor is closed, transaction is committed and
# connection is closed (as requested by the parameters)
@param description: a description of the cursor that we are getting within this context,
this string should answer the question 'what is the cursor needed for?'
@param contextType: a bitmap defined by combining ContextType flags. Default is ContextType.DICT_CURSOR.
@param cursorName: a name for the named cursor. Should be defined when using ContextType.NAMED_CURSOR.
"""
if ( contextType & ContextType.NAMED_CURSOR ) > 0 :
if cursorName is None :
cursorName = 'undefined'
d = 'for ' + str(description) if description is not None else ''
conn = self.getconn()
if ( contextType & ContextType.DICT_CURSOR ) > 0 :
cur = conn.cursor(cursor_factory=DictCursor)
self.logger.debug("Got dict cursor %s", d)
elif ( contextType & ContextType.NAMED_CURSOR ) > 0 :
cur = conn.cursor(cursorName)
self.logger.debug("Got named cursor(%s) %s", cursorName, d)
else :
cur = conn.cursor()
self.logger.debug("Got generic cursor %s", d)
try:
yield cur
except:
if ( contextType & ContextType.TRANSACTIONAL ) > 0 :
conn.rollback()
self.logger.error("Rolling back transaction of cursor %s", d)
raise
else:
if ( contextType & ContextType.TRANSACTIONAL ) > 0 :
conn.commit()
self.logger.debug("Committed transaction of cursor %s", d)
finally:
cur.close()
if ( contextType & ContextType.CLOSE_CONNECTION ) > 0 :
self.putconn(conn, close=True)
self.logger.debug("Closed cursor %s and connection", d)
else :
self.putconn(conn)
self.logger.debug("Closed cursor %s", d)
def _test():
logging.getLogger().setLevel(logging.DEBUG)
fileConfig()
testdb1_pool = get_connection_pool('testdb1')
testdb2_pool = get_connection_pool('testdb2')
with videoflip_pool.context('getting test context to testdb1') as curs:
curs.execute('SELECT %s::text as foo', ( 'à èìòù', ) )
r = curs.fetchone()
print curs.statusmessage, r
with master_pool.context('getting test context to testdb2') as curs:
curs.execute('SELECT %s::text as foo', ( 'à èìòù', ) )
r = curs.fetchone()
print curs.statusmessage, r
if __name__ == '__main__':
_test()
|
Global configuration mechanics was inspired by the logging module global configuration apprach.
