How to modularize a code containing 'with' in Python - python

I have this code that I am using to get information from a mysql database
def query_result_connect(_query):
with SSHTunnelForwarder((ssh_host, ssh_port),
ssh_password=ssh_password,
ssh_username=ssh_user,
remote_bind_address=('127.0.0.1', 3306)) as server:
connection = mdb.connect(user=sql_username,
passwd=sql_password,
db=sql_main_database,
host='127.0.0.1',
port=server.local_bind_port)
cursor = connection.cursor()
cursor.execute(_query)
connection.commit()
try:
y = pd.read_sql(_query, connection)
return y
except TypeError as e:
x = cursor.fetchall()
return x
I would like to create a function that includes the following part.
with SSHTunnelForwarder((ssh_host, ssh_port),
ssh_password=ssh_password,
ssh_username=ssh_user,
remote_bind_address=('127.0.0.1', 3306)) as server:
connection = mdb.connect(user=sql_username,
passwd=sql_password,
db=sql_main_database,
host='127.0.0.1',
port=server.local_bind_port)
and execute it in the query_result_connect() function. The problem is that I don't know how to include more code within the 'with' statement. The code should look something like this:
# Maybe introduce some arguments
def db_connection():
with SSHTunnelForwarder((ssh_host, ssh_port),
ssh_password=ssh_password,
ssh_username=ssh_user,
remote_bind_address=('127.0.0.1', 3306)) as server:
connection = mdb.connect(user=sql_username,
passwd=sql_password,
db=sql_main_database,
host='127.0.0.1',
port=server.local_bind_port)
# Maybe return something
def query_result_connect(_query):
# call the db_connection() function somehow.
# Write the following code in a way that is within the 'with' statement of the db_connection() function.
cursor = connection.cursor()
cursor.execute(_query)
connection.commit()
try:
y = pd.read_sql(_query, connection)
return y
except TypeError as e:
x = cursor.fetchall()
return x
Thank you

What's about to make "do_connection" to be a context manager itself?
#contextmanager
def do_connection():
# prepare connection
# yield connection
# close connection (__exit__). Perhaps you even want to call "commit" here.
Then, you will use it like this:
with do_connection() as connection:
cursor = connection.cursor()
...
It is a common approach to use context managers for creating DB connections.

You could make you own Connection class, that works like a conext manager.
__enter__ sets up ssh tunnel and db connection.
__exit__, tries to close the cursor, db connection and the ssh tunnel.
from sshtunnel import SSHTunnelForwarder
import psycopg2, traceback
class MyDatabaseConnection:
def __init__(self):
self.ssh_host = '...'
self.ssh_port = 22
self.ssh_user = '...'
self.ssh_password = '...'
self.local_db_port = 59059
def _connect_db(self, dsn):
try:
self.con = psycopg2.connect(dsn)
self.cur = self.con.cursor()
except:
traceback.print_exc()
def _create_tunnel(self):
try:
self.tunnel = SSHTunnelForwarder(
(self.ssh_host, self.ssh_port),
ssh_password=self.ssh_password,
ssh_username=self.ssh_user,
remote_bind_address=('localhost', 5959),
local_bind_address=('localhost', self.local_db_port)
)
self.tunnel.start()
if self.tunnel.local_bind_port == self.local_db_port:
return True
except:
traceback.print_exc()
def __enter__(self):
if self._create_tunnel():
self._connect_db(
"dbname=mf port=%s host='localhost' user=mf_usr" %
self.local_db_port
)
return self
def __exit__(self, *args):
for c in ('cur', 'con', 'tunnel'):
try:
obj = getattr(self, c)
obj.close()
obj = None
del obj
except:
pass
with MyDatabaseConnection() as db:
print(db)
db.cur.execute('Select count(*) from platforms')
print(db.cur.fetchone())
Out:
<__main__.MyDatabaseConnection object at 0x1017cb6d0>
(8,)
Note:
I am connecting to Postgres, but that should work using mysql as well. Probably you need to adjust to match your own needs.

Related

Return database errors in Python

I have a Python Azure function. In the code, I have a Connection timeout error exception from pyodbc connection.
def connection():
username = ''
server = ''
password = ''
database = ''
connection= pyodbc.connect('DRIVER={ODBC Driver 17 for SQL Server};SERVER='+server+';DATABASE='+database+';ENCRYPT=yes;UID='+username+';PWD='+ password)
return connection
When the connection times out the code fails with a connection timeout error in the logs. But I need the function to return the error value instead of showing the error message in the logs. Is there any way to achieve this in Python? Currently, I am using a try-except statement.
try:
conn = connection()
# .....
except pyodbc.OperationalError as e:
raise e
Edit:
Helper code:
def connection():
username = ''
server = ''
password = ''
database = ''
connection= pyodbc.connect('DRIVER={ODBC Driver 17 for SQL Server};SERVER='+server+';DATABASE='+database+';ENCRYPT=yes;UID='+username+';PWD='+ password)
return conctn
def read_table(table, cnxn):
conctn.timeout = 3
df = pd.read_sql('SELECT * FROM ' + table, cnxn)
if df.empty:
raise customerror
return df
Main code:
try:
conn = connection()
df1 = read_table("table_name", conn)
except pyodbc.OperationalError as e:
return e
except customerror as e:
return e
else:
.....
When the connection times out, the main code fails with the below error.
SELECT * FROM table_name': ('HYT00', '[HYT00] [Microsoft][ODBC Driver 17 for SQL Server]Query timeout expired (0) (SQLExecDirectW)
When the query times out I do not want the code to fail with the above error message instead I want to exit/terminate the code with a custom time out error message.
You're describing, vaguely, something like...
def connection():
username = ''
server = ''
password = ''
database = ''
try:
return (
pyodbc.connect('DRIVER={ODBC Driver 17 for SQL Server};SERVER='+server+';DATABASE='+database+';ENCRYPT=yes;UID='+username+';PWD='+ password),
None
)
except pyodbc.OperationalError as e:
return (
None,
e
)
Then, an operational error will never cause a crash, and your calling code can be something like...
conn, e = connection()
if e:
# do something with the error
else:
# do something with the connection

AWS: Unable to insert values in Redshift table

I am trying to insert values into a table within my redshift cluster, it is connected as I can read the table but I can't insert values on it. When I use SELECT statements it works fine but when I try to insert values from lambda function, it is aborted with no error or log info about why was it aborted.
The query part is like this:
conn = psycopg2.connect(dbname = 'dev',
host =
'redshift-cluster-summaries.c0xcgwtgz65l.us-east-2.redshift.amazonaws.com',
port = '5439',
user = '****',
password = '****%')
cur = conn.cursor()
cur.execute("INSERT INTO public.summaries(topic,summary)
values('data', 'data_summary');")
#print(cur.fetchone())
cur.close()
conn.close()
As I said, there is no log information about why was it aborted, neither it is giving me any kind of error. Actually, when I just use a Select statement, it works.
Is there anyone who can guide me through what could be going on?
You forgot to do conn.commit()
conn = psycopg2.connect(dbname = 'dev',
host = 'redshift-cluster-summaries.c0xcgwtgz65l.us-east-2.redshift.amazonaws.com',
port = '5439',
user = '****',
password = '****%')
cur = conn.cursor()
cur.execute("INSERT INTO public.summaries(topic,summary) values('data', 'data_summary');")
cur.close()
conn.commit()
conn.close()
a bit improved way to run this
from contextlib import contextmanager
#contextmanager
def cursor():
with psycopg2.connect(dbname = 'dev',
host = 'redshift-cluster-summaries.c0xcgwtgz65l.us-east-2.redshift.amazonaws.com',
port = '5439',
user = '****',
password = '****%') as conn:
try:
yield conn.cursor()
finally:
try:
conn.commit()
except psycopg2.InterfaceError:
pass
def run_insert(query):
with cursor() as cur:
cur.execute(query)
cur.close()
run_insert("INSERT INTO public.summaries(topic,summary) values('data', 'data_summary');")

2013, 'Lost connection to MySQL server during query'

This is my MySQLConnector code.
import logging
import pymysql
logger = logging.getLogger(__name__)
class MySQLConnector:
def __init__(self, config):
self.connection = None
self._connect(config=config)
def _connect(self, config):
config["charset"] = "utf8mb4"
config["cursorclass"] = pymysql.cursors.DictCursor
self.connection = pymysql.connect(**config)
def read(self, query, params=None):
result = None
try:
with self.connection.cursor() as cursor:
cursor.execute(query, params)
result = cursor.fetchall()
self.connection.commit()
except Exception as e:
logger.error(e)
return result
And I use that like this.
connector = MySQLConnector(config=config)
while True:
query = "SELECT * FROM my_table"
print(connector.read(query=query))
time.sleep(30)
This works well but after few hours later, this process raise error (2013, 'Lost connection to MySQL server during query'), and I can't find the reason.

Use multiple functions in MySQLdb (python)

this code below runs just fine, but i want to separate this code into function (this is my first time using MySQLdb) ,
import MySQLdb
# Open database connection
db = MySQLdb.connect("localhost","user_name","pass","db_name" )
cursor = db.cursor()
sql = "SELECT activity_log.datetime FROM activity_log"
cursor.execute(sql)
date_data = cursor.fetchall()
for content in date_data:
print content
# disconnect from server
db.close()
Question: how could I create one database connection and use it to multiple functions, this is what i've wrote so far (doesn't work):
import MySQLdb
class DB():
def __init__(self):
db = MySQLdb.connect("locahost", "user_name", "pass", "db_name")
self.cur = db.cursor()
def time_statistic(self):
sql = "SELECT activity_log.datetime FROM activity_log"
self.cur.execute(sql)
self.date_data = self.cursor.fetchone()
for content in self.date_data:
print content
def test1(self):
pass
if __name__ == '__main__':
db = DB.connect("db_name" )
db.time_statistic(self)
db.test1(self)
db.close()
You need to pass the arguments to init to make sure that the class connects to the right DB.
import MySQLdb
class DB():
def __init__(self, server, user, password, db_name):
db = MySQLdb.connect(server, user, password, db_name)
self.cur = db.cursor()
def time_statistic(self):
sql = "SELECT activity_log.datetime FROM activity_log"
self.cur.execute(sql)
self.date_data = self.cursor.fetchone()
for content in self.date_data:
print content
def test1(self):
pass
if __name__ == '__main__':
db = DB(<server>, <user>, <password>, <db_name>)
db.time_statistic()
db.test1()
Replace the arguments in <> with actual values you need to connect to the db. You may also want to add some error handling in the above code.

Python 2.7 connection string to PostgreSQL (OOP method)

I'm new to python and I'm trying to make this work. I'm using Python 2.7 and PostgreSQL 9.3:
#! F:\Python2.7.6\python
import psycopg2
class Database:
host = "192.168.56.101"
user = "testuser"
passwd = "passwd"
db = "test"
def __init__(self):
self.connection = psycopg2.connect( host = self.host,
user = self.user,
password = self.passwd,
dbname = self.db )
self.cursor = self.connection.cursor
def query(self, q):
cursor = self.cursor
cursor.execute(q)
return cursor.fetchall()
def __del__(self):
self.connection.close()
if __name__ == "__main__":
db = Database()
q = "DELETE FROM testschema.test"
db.query(q)
However I am getting an error "AttributeError: 'builtin_function_or_method' object has no attribute 'execute'". I figure I should put something like self.execute = something in the Database class, but I can't figure it out what exactly I need to put there. Any suggestions?
You are missing the parenthesis at the end
self.cursor = self.connection.cursor()
or
cursor = self.cursor()
But not both

Categories