We Keep Coding

How to create a postgres connection that uses context manager?

I was trying to create a wrapper class that used enter/exit to allow for a context manager to close a postgres db connection.
class ClosingConnection:
def __init__(self, schema_name: str) -> None:
"""
:param schema_name: the db schema (i.e. the tenant)
"""
super().__init__()
self.schema_name = schema_name
def __enter__(self):
try:
self.conn = psycopg2.connect(
host=os.environ["DB_HOST"],
port=os.environ["DB_PORT"],
database=os.environ["DB_NAME"],
user=os.environ["DB_USERNAME"],
password=password,
options=f"-c search_path={self.schema_name}",
cursor_factory=psycopg2.extras.DictCursor,
)
return self.conn
except psycopg2.OperationalError:
pass
def __exit__(self, exc_type, exc_val, exc_tb):
self.conn.close()
Usage
def get_db(schema_name: str):
return ClosingConnection(schema_name)
with get_db("test") as db:
cursor = db.cursor()
cursor.execute("some sql")
rv = cursor.fetchall()
cursor.close()
This works fine to retrieve data, but if I do an upsert, only the first upsert is performed. The remaining ones rv returns correctly, but the actual database is never updated.
In contrast if I get rid of the closing class, this works fine
Usage
def get_db(schema_name: str):
try:
self.conn = psycopg2.connect(
host=os.environ["DB_HOST"],
port=os.environ["DB_PORT"],
database=os.environ["DB_NAME"],
user=os.environ["DB_USERNAME"],
password=password,
options=f"-c search_path={self.schema_name}",
cursor_factory=psycopg2.extras.DictCursor,
)
return self.conn
except psycopg2.OperationalError:
pass
conn = get_db("test")
try:
cursor = db.cursor()
cursor.execute("some sql")
rv = cursor.fetchall()
cursor.close()
finally:
conn.close()
For curiosity here is the upsert command I'm using (I've also tried the older with upsert version)
INSERT INTO settings_account (setting_name, setting_value)
VALUES (%(setting_name)s, %(setting_value)s)
ON CONFLICT (setting_name)
DO
UPDATE SET setting_value=EXCLUDED.setting_value
RETURNING *
Why is the wrapper class causing the db writes to fail? Is the idea of wrapper class flawed?

Thanks the comment above, i was able to get this working with:
class ClosingConnection:
def __init__(self, schema_name: str) -> None:
"""
:param schema_name: the db schema (i.e. the tenant)
"""
super().__init__()
self.schema_name = schema_name
def __enter__(self):
try:
self.conn = psycopg2.connect(
host=os.environ["DB_HOST"],
port=os.environ["DB_PORT"],
database=os.environ["DB_NAME"],
user=os.environ["DB_USERNAME"],
password=password,
options=f"-c search_path={self.schema_name}",
cursor_factory=psycopg2.extras.DictCursor,
)
return self.conn
except psycopg2.OperationalError:
pass
def __exit__(self, exc_type, exc_val, exc_tb):
self.conn.commit()
self.conn.close()
def get_db(schema_name: str):
return ClosingConnection(schema_name)
with get_db("test") as db:
with db.cursor() as cursor:
cursor.execute("some sql")
rv = cursor.fetchall()

psycopg2 - Unkeyed Connection error in multithreading

I am implementing connection pooling with Postgres for a multithreading application, Here I am creating two pools and storing in a singleton instance to be able to access the same instance from the entire application with the class RepositoryConnPool
Then we have the second class that is for context management, this I made to be able to execute multiple query's with the same connection in a with context, I have two pools of connections because when I execute a query i want to be able to execute a select query to replica database and modification of data to the master database. I am getting the following error while executing in a multithreading environment. This is not happening while I run without multithreading. Here is the error log :
ERROR ;2020-06-03 11:05:42,826; ThreadPoolExecutor-0_4; modules.table_modify_module; table_modify_module:insert_host;60: trying to put unkeyed connection
ERROR ;2020-06-03 11:05:42,826; ThreadPoolExecutor-0_21; table_modify_module; table_modify_module:insert_host;60: trying to put unkeyed connection
ERROR ;2020-06-03 11:05:42,885; ThreadPoolExecutor-0_12; table_modify_module; table_modify_module:get_inst_detail;367: trying to put unkeyed connection
class RepositoryConnPool:
__instance = None
#This method overcharge is to make use of singleton pattern
def __new__(cls):
if cls.__instance is None:
cls.__instance = object.__new__(cls)
return cls.__instance
def __init__(self):
try:
#I get connection details from vault
inventory_master = vm.get_secret_by_base_path("xxxxxx")
inventory_repl = vm.get_secret_by_base_path("xxxxx")
self.repo_connection_pool_master = self.setup_conn_pool(inventory_master)
self.repo_connection_pool_repl = self.setup_conn_pool(inventory_repl)
except Exception as e:
logger.error("GET repository address failed: ", exc_info=e)
def setup_conn_pool(self,inventory_data):
repo_dbausername = 'dba_admin'
repo_dbapassword = inventory_data['accounts']['dba_admin']['password']
repo_host = inventory_data['fqdn']
repo_port = inventory_data['port']
try:
repo_connection_pool = psycopg2.pool.ThreadedConnectionPool(10, 30, user=repo_dbausername,
password=repo_dbapassword,
host=repo_host,
port=repo_port,
database="xxxxx")
except Exception as e:
logger.error(e)
return repo_connection_pool
#This method will fetch a connection from pool
def get_master_connection(self):
return self.repo_connection_pool_master.getconn()
def get_repl_connection(self):
return self.repo_connection_pool_repl.getconn()
class RepositoryConnection:
#Gets connection from RepositoryConnPool
def __init__(self):
self.repo_conn_pool = RepositoryConnPool()
self._repo_conn_master = self.repo_conn_pool.get_master_connection()
self._repo_conn_repl = self.repo_conn_pool.get_repl_connection()
def execute_query(self, query, commit=False, call_master = False, retry_count = 0):
cursor = None
try:
if query.strip().upper().startswith("SELECT") and call_master == False:
repo_query = query
cursor = self._repo_conn_repl.cursor()
cursor.execute(repo_query)
if commit:
self._repo_conn_repl.commit()
else:
try:
result = cursor.fetchall()
return result
except:
logging.info(f"Query executed:{query} and no result was returned")
else:
repo_query = query
cursor = self._repo_conn_master.cursor()
cursor.execute(repo_query)
if commit:
self._repo_conn_master.commit()
else:
try:
result = cursor.fetchall()
return result
except:
logging.info(f"Query executed:{query} and no result was returned")
#This needs to change to handle connection error connection in different way
except Exception as e:
if(retry_count < 3):
logger.warning(f"There was an error during execution of query, retry will be done {query} ")
return self.execute_query(query,commit,call_master = True, retry_count= retry_count+1 )
else:
logger.fatal("All 3 connection retry failed")
finally:
if cursor:
cursor.close()
#Checks if connection is up
def __enter__(self):
try:
if self._repo_conn_master and self._repo_conn_repl:
return self
else:
raise Exception()
except Exception as e:
logger.error(e)
#Returns connection to pool
def __exit__(self, exc_type, exc_val, exc_tb):
if self._repo_conn_master:
self.repo_conn_pool.repo_connection_pool_master.putconn(self._repo_conn_master)
if self._repo_conn_repl:
self.repo_conn_pool.repo_connection_pool_repl.putconn(self._repo_conn_repl)

Python - Manage cursor connection outside context manager

I'm new at Python and I'm trying to build a start project to get into this language.
I created a SQLite3 DB and managed to make transactions with it.
Everything works fine.
I wanted to get deeper in Python so I've been searching and discovered Decorators and Context Manager and I was trying to implement these concepts on my Query Execution's functions. However, I'm stuck with a problem.
I've created a class that handles the open and close connection tasks.
DB_ContextManager.py class:
class DB_ContextManager():
def __init__(self, db_connection):
self.db_connection = db_connection
def __enter__(self):
self.conn = sqlite3.connect(self.db_connection)
return self.conn
def __exit__(self, exc_type, exc_val, exc_tb): # obligatory params
self.conn.close()
And also created ConnectionDB.py which is responsible for executing queries.
from Database.DB_ContextManager import DB_ContextManager as DB_CM
# Handles SELECT queries
def ExecuteSelectQuery(self, pQuery):
try:
with DB_CM(db_connection_string) as conn:
cur = conn.cursor()
cur.execute(pQuery)
result = cur.fetchall()
return result
except Exception as e:
LH.Handler(log_folder, 'ConnectionDB', 'Queries', 'ExecuteSelectQuery', e)
raise DE.ConnectionDB_Exception()
# Handles INSERTs, UPDATEs, DELETEs queries
def ExecuteNonQuery(self, pQuery):
try:
with DB_CM(db_connection_string) as conn:
cur = conn.cursor()
cur.execute(pQuery)
except Exception as e:
LH.Handler(log_folder, 'ConnectionDB', 'Queries', 'ExecuteSelectNonQuery', e)
raise DE.ConnectionDB_Exception()
As you can see
with DB_CM(db_connection_string) as conn:
cur = conn.cursor()
cur.execute(pQuery)
is repeated in each function
To avoid this situation, I'd like to create a Decorator function that encapsulates this piece of code.
My problem is that the cursor 'dies' inside the ContextManager and, for example, ExecuteSelectQuery needs the cursor to fetch the return data after the query was executed.
I know it's a small project and thinking so long term in future may not be necessary. But, remember, it's a start project and I'm learning to apply new concepts.
SOLUTION
As #blhsing suggested, I return the connection object instead of the cursor in the ContextManager.
Also I handles commit() and rollback() in it.
So, summarizing:
ConnectionDB.py
def ExecuteSelectQuery(self, pQuery):
with DB_CM(db_connection_string, pQuery) as cur:
result = cur.fetchall()
return result
def ExecuteSelectNonQuery(self, pQuery):
with DB_CM(db_connection_string, pQuery) as cur:
pass
and ConnectionDB.py
class DB_ContextManager():
def __init__(self, db_connection, pQuery):
self.db_connection = db_connection
self.query = pQuery
def __enter__(self):
try:
self.conn = sqlite3.connect(self.db_connection)
cur = self.conn.cursor()
cur.execute(self.query)
self.conn.commit()
return cur
except Exception as e:
LH.Handler(log_folder, 'DB_ContextManager', 'DB_ContextManager', '__enter__', e)
self.conn.rollback()
raise DE.ConnectionDB_Exception()
def __exit__(self, exc_type, exc_val, exc_tb): # obligatory params
self.conn.close()

You can make the context manager return the cursor instead of the connection object:
class DB_CM():
def __init__(self, db_connection):
self.db_connection = db_connection
def __enter__(self):
self.conn = sqlite3.connect(self.db_connection)
cur = self.conn.cursor()
cur.execute(pQuery)
return cur
def __exit__(self, exc_type, exc_val, exc_tb): # obligatory params
self.conn.close()
so that the try block of ExecuteSelectQuery can be revised as:
with DB_CM(db_connection_string) as cur:
result = cur.fetchall()
return result
and the try block of ExecuteNonQuery can be simply:
with DB_CM(db_connection_string):
pass

Why using pymysql to operate the MySQL with Class does not work, but contextlib works well?

I use pymysql (python3.5)to operate the MySQL(ver>5.5), and I want to use it simply, so I write a Class:
class Mysql(object):
def __init__(self):
self.conn = connect(host=HOST, port=PORT, user=USER, passwd=PASSWORD, db=DB, charset=CHARSET)
"""DictCursor：A cursor which returns results as a dictionary"""
self.cursor = self.conn.cursor(cursor=pymysql.cursors.DictCursor)
def __enter__(self):
return self.cursor
def __exit__(self, exc_type, exc_val, exc_tb):
self.cursor.close()
self.conn.close()
Then, I wrote this:
def do_add(_name, _age):
with Mysql() as cursor:
param = (_name, _age)
sql = """insert into mydb(name, age) values (%s, %s)"""
cursor.execute(sql, param)
That does not work, I can not insert data into MySQL.
But, when I use the context Manager, it works very well:
#contextlib.contextmanager
def mysql():
conn = connect(host=HOST, port=PORT, user=USER, passwd=PASSWORD, db=DB, charset=CHARSET)
cursor = conn.cursor(cursor=pymysql.cursors.DictCursor)
try:
yield cursor
conn.commit()
except:
conn.rollback()
finally:
cursor.close()
conn.close()
I want to know why I can not use Class?

Python Postgres psycopg2 ThreadedConnectionPool exhausted

I have looked into several 'too many clients' related topic here but still can't solve my problem, so I have to ask this again, for me specific case.
Basically, I set up my local Postgres server and need to do tens of thousands of queries, so I used the Python psycopg2package. Here are my codes:
import psycopg2
import pandas as pd
import numpy as np
from flashtext import KeywordProcessor
from psycopg2.pool import ThreadedConnectionPool
from concurrent.futures import ThreadPoolExecutor
df = pd.DataFrame({'S':['California', 'Ohio', 'Texas'], 'T':['Dispatcher', 'Zookeeper', 'Mechanics']})
# df = pd.concat([df]*10000) # repeat df 10000 times
DSN = "postgresql://User:password#localhost/db"
tcp = ThreadedConnectionPool(1, 800, DSN)
def do_one_query(inputS, inputT):
conn = tcp.getconn()
c = conn.cursor()
q = r"SELECT * from eridata where "State" = 'California' and "Title" = 'Dispatcher' limit 1;"
c.execute(q)
all_results = c.fetchall()
for row in all_results:
return row
tcp.putconn(conn, close=True)
cnt=0
for idx, row in df.iterrows():
cnt+=1
with ThreadPoolExecutor(max_workers=1) as pool:
ret = pool.submit(do_one_query, row["S"], row["T"])
print ret.result()
print cnt
The code runs well with a small df. If I repeat df by 10000 times, I got error message saying connection pool exhausted
. I though the connections I used have been closed by this line:
tcp.putconn(conn, close=True)
But I guess actually they are not closed? How can I get around this issue?

I've struggled to find really detailed information on how the ThreadedConnectionPool works. https://bbengfort.github.io/observations/2017/12/06/psycopg2-transactions.html ain't bad, but it turns out that its claim that getconn blocks until a connection becomes available is incorrect. Checking the code, all ThreadedConnectionPool adds is a lock around the AbstractConnectionPool methods to prevent race conditions. If more than maxconn connections are attempted used at any point, the connection pool exhausted PoolError will be raised.
If you want something a bit simpler than the accepted answer, further wrapping the methods in a Semaphore providing the blocking until a connection becomes available should do the trick:
from psycopg2.pool import ThreadedConnectionPool as _ThreadedConnectionPool
from threading import Semaphore
class ThreadedConnectionPool(_ThreadedConnectionPool):
def __init__(self, minconn, maxconn, *args, **kwargs):
self._semaphore = Semaphore(maxconn)
super().__init__(minconn, maxconn, *args, **kwargs)
def getconn(self, *args, **kwargs):
self._semaphore.acquire()
try:
return super().getconn(*args, **kwargs)
except:
self._semaphore.release()
raise
def putconn(self, *args, **kwargs):
try:
super().putconn(*args, **kwargs)
finally:
self._semaphore.release()
# closeall is inherited as is. This means the Semaphore does
# not get reset, but neither do the core structures for
# maintaining the pool in the original ThreadedConnectionPool
# so a closed pool is not intended to be reused once closed.
Note that ConnectionPools, both standard and threaded, only come with the three putconn, getconn and closeall methods, and nothing fancy like context management. So the above should cover all existing functionality.

You need to use a queue on top of your pool.
Something like the following should work:
import gevent, sys, random, psycopg2, logging
from contextlib import contextmanager
from gevent.queue import Queue
from gevent.socket import wait_read, wait_write
from psycopg2.pool import ThreadedConnectionPool
from psycopg2 import extensions, OperationalError
import sys
logger = logging.getLogger(__name__)
poolsize = 100 #number of max connections
pdsn = '' # put your dsn here
if sys.version_info[0] >= 3:
integer_types = (int,)
else:
import __builtin__
integer_types = (int, __builtin__.long)
class ConnectorError(Exception):
""" This is a base class for all CONNECTOR related exceptions """
pass
#simplified calls etc. db.fetchall(SQL, arg1, arg2...)
def cursor(): return Pcursor()
def fetchone(PSQL, *args): return Pcursor().fetchone(PSQL, *args)
def fetchall(PSQL, *args): return Pcursor().fetchall(PSQL, *args)
def execute(PSQL, *args): return Pcursor().execute(PSQL, *args)
#singleton connection pool, gets reset if a connection is bad or drops
_pgpool = None
def pgpool():
global _pgpool
if not _pgpool:
try:
_pgpool = PostgresConnectionPool(maxsize=poolsize)
except psycopg2.OperationalError as exc:
_pgpool = None
return _pgpool
class Pcursor(object):
def __init__(self, **kwargs):
#in case of a lost connection lets sit and wait till it's online
global _pgpool
if not _pgpool:
while not _pgpool:
try:
pgpool()
except:
logger.debug('Attempting Connection To Postgres...')
gevent.sleep(1)
def fetchone(self, PSQL, *args):
with _pgpool.cursor() as cursor:
try:
cursor.execute(PSQL, args)
except TypeError:
cursor.execute(PSQL, args[0])
except Exception as exc:
print(sys._getframe().f_back.f_code)
print(sys._getframe().f_back.f_code.co_name)
logger.warning(str(exc))
logger.debug(cursor.query)
return cursor.fetchone()
def fetchall(self, PSQL, *args):
with _pgpool.cursor() as cursor:
try:
cursor.execute(PSQL, args)
except TypeError:
cursor.execute(PSQL, args[0])
except Exception as exc:
print(sys._getframe().f_back.f_code)
print(sys._getframe().f_back.f_code.co_name)
logger.warning(str(exc))
logger.debug(cursor.query)
return cursor.fetchall()
def execute(self, PSQL, *args):
with _pgpool.cursor() as cursor:
try:
cursor.execute(PSQL, args)
except TypeError:
cursor.execute(PSQL, args[0])
except Exception as exc:
print(sys._getframe().f_back.f_code)
print(sys._getframe().f_back.f_code.co_name)
logger.warning(str(exc))
finally:
logger.debug(cursor.query)
return cursor.query
def fetchmany(self, PSQL, *args):
with _pgpool.cursor() as cursor:
try:
cursor.execute(PSQL, args)
except TypeError:
cursor.execute(PSQL, args[0])
while 1:
items = cursor.fetchmany()
if not items:
break
for item in items:
yield item
class AbstractDatabaseConnectionPool(object):
def __init__(self, maxsize=poolsize):
if not isinstance(maxsize, integer_types):
raise TypeError('Expected integer, got %r' % (maxsize, ))
self.maxsize = maxsize
self.pool = Queue()
self.size = 0
def create_connection(self):
#overridden by PostgresConnectionPool
raise NotImplementedError()
def get(self):
pool = self.pool
if self.size >= self.maxsize or pool.qsize():
return pool.get()
self.size += 1
try:
new_item = self.create_connection()
except:
self.size -= 1
raise
return new_item
def put(self, item):
self.pool.put(item)
def closeall(self):
while not self.pool.empty():
conn = self.pool.get_nowait()
try:
conn.close()
except Exception:
pass
#contextmanager
def connection(self, isolation_level=None):
conn = self.get()
try:
if isolation_level is not None:
if conn.isolation_level == isolation_level:
isolation_level = None
else:
conn.set_isolation_level(isolation_level)
yield conn
except:
if conn.closed:
conn = None
self.closeall()
raise
else:
if conn.closed:
raise OperationalError("Cannot commit because connection was closed: %r" % (conn, ))
finally:
if conn is not None and not conn.closed:
if isolation_level is not None:
conn.set_isolation_level(isolation_level)
self.put(conn)
#contextmanager
def cursor(self, *args, **kwargs):
isolation_level = kwargs.pop('isolation_level', None)
with self.connection(isolation_level) as conn:
try:
yield conn.cursor(*args, **kwargs)
except:
global _pgpool
_pgpool = None
del(self)
class PostgresConnectionPool(AbstractDatabaseConnectionPool):
def __init__(self,**kwargs):
try:
self.pconnect = ThreadedConnectionPool(1, poolsize, dsn=pdsn)
except:
global _pgpool
_pgpool = None
raise ConnectorError('Database Connection Failed')
maxsize = kwargs.pop('maxsize', None)
self.kwargs = kwargs
AbstractDatabaseConnectionPool.__init__(self, maxsize)
def create_connection(self):
self.conn = self.pconnect.getconn()
self.conn.autocommit = True
return self.conn
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 ConnectorError("Bad result from poll: %r" % state)
extensions.set_wait_callback(gevent_wait_callback)
Then you can call your connection via this:
import db
db.Pcursor().execute(PSQL, arg1, arg2, arg3)
Basically I borrowed the gevent example of async postgres and modified it to support threadpooling via pyscopg2.
https://github.com/gevent/gevent/blob/master/examples/psycopg2_pool.py
I added what psycogreen does inside the module, so all you need to do is import and call the class. Each call to the class stacks a new query on the queue, but only uses the pool at a certain size. This way you don't run out of connections. This is essentially similar to what PGBouncer does, which I think would also eliminate your problem.
https://pgbouncer.github.io/

Your problem here is, that you actually do not return the connection to the pool, but close it forever with
tcp.putconn(conn, close=True)
See the documentation here http://initd.org/psycopg/docs/pool.html
If close is True, discard the connection from the pool.
So, if you put 800 connections into your pool, after 801 loops you will get the "exhausted error" because your connection pool size is zero.

I think the reason why you get PoolError("exhausted connections") maybe you return before close connection when all_results is not None. so, connection pool exhausted
def do_one_query(inputS, inputT):
...
for row in all_results:
return row <---- return row before putconn when all_results is not None,
tcp.putconn(conn, close=True)
for idx, row in df.iterrows():
cnt+=1
with ThreadPoolExecutor(max_workers=1) as pool:
ret = pool.submit(do_one_query, row["S"], row["T"])
print ret.result()
print cnt
I make a ugly implementation with when exhausted or connection lost, try reconnect to get new conn, like below
class PostgresConnectionPool:
def __init__(self, minconn, maxconn, *args, **kwargs):
self.pool = ThreadedConnectionPool(minconn=minconn, maxconn=maxconn, *args, **kwargs)
def get_conn(self):
try:
# check if connection lost or pool exhausted
con = self.pool.getconn()
cur = con.cursor()
cur.execute("select 1;")
except (OperationalError, PoolError) as oe:
print(f"get pg connection with err:{oe}, reconnect")
# reconnect
key = str(uuid.uuid4())
con = self.pool._connect(key)
return con