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()
Related
I want to create a class in Python that establishes a connection to SnowFlake. I have a user.txt file that specifies an account, warehouse, database, schema, and my user. Here is the code I have so far:
import pandas as pd
import snowflake.connector
import os
from getpass import getpass
import sfcommon.sfdb as sf
class Database:
def __init__(self):
self.environ = os.environ['SFPASSWORD'] = getpass(prompt='Password: ', stream=None)
self._conn = sf.get_connect()
self._cursor = self._conn.cursor()
def __enter__(self):
return self
def __exit__(self, exc_type, exc_val, exc_tb):
self.close()
#property
def connection(self):
return self._conn
#property
def cursor(self):
return self._cursor
def commit(self):
self.connection.commit()
def close(self, commit=True):
if commit:
self.commit()
self.connection.close()
def execute(self, sql, params=None):
self.cursor.execute(sql, params or ())
def fetchall(self):
return self.cursor.fetchall()
def fetchone(self):
return self.cursor.fetchone()
def query(self, sql, params=None):
self.cursor.execute(sql, params or ())
return self.fetchall()
sql = "select * from test_database"
After running this class, I would hope to run this block of code to retrieve all the rows in my database in a pandas dataframe:
with Database() as test:
resultSet = pd.read_sql(sql, conn)
I have managed to create a connection to Snowflake but have been trying to get it into a class for easier readability.
Also, I have this block of code that I'm not sure how to integrate into my Database class
%reload_ext sql_magic
%config SQL.conn_name = 'conn'
You can use the below code for dataframe .
---------------------------------------------------------------------------------
import snowflake.connector
import pandas as pd
# creates a connection
def openConn():
ctx = snowflake.connector.connect(
account='',
user='',
password='',
database='',
schema='public',
warehouse='',
role='',
)
return ctx
def main():
conn = openConn()
cs = conn.cursor()
cs.execute("""<query>""")
# Put it all to a data frame
sql_data = pd.DataFrame(cs.fetchall())
print(sql_data)
if __name__ == "__main__":
main()
------------------------------------------------------------------------------
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
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?
I have a database class like:
rlock = RLock()
class PostgresDB:
host = Settings.DB_HOST
name = Settings.DB_NAME
username = Settings.DB_USERNAME
pwd = Settings.DB_PWD
def __init__(self):
self.con = None
def _create_connection(self):
if self.con is None or self.con.closed:
self.con = psycopg2.connect("host={h} dbname={n} user={u} password={p}".format(
h=self.host, n=self.name, u=self.username, p=self.pwd,
cursor_factory=RealDictCursor))
#contextlib.contextmanager
def reading_cursor(self, sql, params=None):
rlock.acquire()
self._create_connection()
try:
cur = self.con.cursor()
cur.execute(sql, params)
yield cur
except DatabaseError as e:
traceback.print_exc()
yield Exception()
finally:
self.con.close()
rlock.release()
#contextlib.contextmanager
def writing_cursor(self, sql, params=None):
rlock.acquire()
self._create_connection()
try:
cur = self.con.cursor()
cur.execute(sql, params)
yield cur
self.con.commit()
except DatabaseError as e:
traceback.print_exc()
yield Exception()
finally:
self.con.close()
rlock.release()
Note that I am explicitly asking for RealDictCursor.
However, when any request is done, the result is a list of tuples, not dicts. I am not sure what I am missing?
Ok this is incredibly stupid... If you look closely you will see I passed cursor_factory to .format() instead of passing it to .connect() ...
i was planning to change my project to multiprocesses so i can use more resources,here's my database module code
import pymysql
import threading
class tdb:
def __init__(self):
self.totalEffected = 0
pass
def start(self):
self.conn = pymysql.connect(host='xxxx', port=3306, user='root', passwd='xxxx', db='xxxx', charset='utf8')
def select(self,sql,args=None):
cur = self.conn.cursor()
cur.execute(sql,args)
result = cur.fetchall()
cur.close()
return result
def execute(self,sql,args=None):
cur = self.conn.cursor()
result = cur.execute(sql,args)
cur.close()
self.totalEffected+=result
return result
# def __commit(self,callback):
def __commitCallback(self,result):
print('commit result:',result)
self.conn.close()
def errorc(self,*args):
print('error')
def end(self):
# init()
# p.apply_async(self.conn.commit, callback=self.__commitCallback,error_callback=self.errorc)
if self.totalEffected!=0:
thread = threading.Thread(target=self.t)
thread.start()
else:
self.conn.close()
# p.apply(self.conn.commit)
# self.conn.close()
# print('result:' ,result.get())
def t(self):
self.conn.commit()
self.conn.close()
the only operation that really need to handle is conn.commit(), i use thread to do it ,so i can immediately return. i once use Pool.apply_async(),but it didn't callback, so i want to know how to make the other process call me , so i don't have to spend my time waiting recieve.