Using unittest in python with docker and psycopg2 - python

BRIEF DESCRIPTION OF THE PROBLEM: Psycopg2 won't connect to a test DB within docker from unittest, but connects fine from console.
ERROR MESSAGE:
psycopg2.OperationalError: server closed the connection unexpectedly
This probably means the server terminated abnormally before or while processing the request.
DETAILS:
I'm trying to set up a testing database in docker, that will be created and populated before testing and then removed after.
Here's the fuction to set up database:
def set_up_empty_test_db():
client = docker.from_env()
try:
testdb = client.containers.get("TestDB")
testdb.stop()
testdb.remove()
testdb = client.containers.run(
"postgres",
ports={5432: 5433},
detach=True,
name="TestDB",
environment=["POSTGRES_PASSWORD=yourPassword"],
)
except NotFound:
testdb = client.containers.run(
"postgres",
ports={5432: 5433},
detach=True,
name="TestDB",
environment=["POSTGRES_PASSWORD=yourPassword"],
)
while testdb.status != "running":
testdb = client.containers.get("TestDB")
return
If I launch this function from console it works without an error and I can see TestDB container running. I can successfully initiate connection:
conn = psycopg2.connect("host='localhost' user='postgres' password='yourPassword' port='5433'")
But it doesn't work when unit testing. Here's the testing code:
class TestCreateCity(unittest.TestCase):
def setUp(self):
set_up_empty_test_db()
con = psycopg2.connect("host='localhost' user='postgres' password='yourPassword' port='5433'")
self.assertIsNone(con.isolation_level)
cur = con.cursor()
sql_file = open(os.path.join(ROOT_DIR + "/ddl/creates/schema_y.sql"), "r")
cur.execute(sql_file.readline())
con.commit()
con.close()
self.session = Session(creator=con)
def test_create_city(self):
cs.create_city("Test_CITY", "US")
q = self.session.query(City).filter(City.city_name == "Test_CITY").one()
self.assertIs(q)
self.assertEqual(q.city_name, "Test_CITY")
self.assertEqual(q.country_code, "US")
It breaks when trying to initiate connection. Please advise.

I know this is an old question, but I needed to do the same thing today. You try and connect to the postgres server too quickly after starting it - that's why it works in the console.
All you need to do is replace:
set_up_empty_test_db()
con = psycopg2.connect("host='localhost' user='postgres' password='yourPassword' port='5433'")
with:
set_up_empty_test_db()
con = None
while con == None:
try:
con = psycopg2.connect("host='localhost' user='postgres' password='yourPassword' port='5433'")
except psycopg2.OperationalError:
time.sleep(0.5);
Hope this helps someone else. Cheers!

Related

accessing postgresql database in python using functions

Let me start off by saying I am extremely new to Python and Postgresql so I feel like I'm in way over my head. My end goal is to get connected to the dvdrental database in postgresql and be able to access/manipulate the data. So far I have:
created a .config folder and a database.ini is within there with my login credentials.
in my src i have a config.py folder and use config parser, see below:
def config(filename='.config/database.ini', section='postgresql'):
# create a parser
parser = ConfigParser()
# read config file
parser.read(filename)
# get section, default to postgresql
db = {}
if parser.has_section(section):
params = parser.items(section)
for param in params:
db[param[0]] = param[1]
else:
raise Exception('Section {0} not found in the {1} file'.format(section, filename))
return db
then also in my src I have a tasks.py file that has a basic connect function, see below:
import pandas as pd
from clients.config import config
import psycopg
def connect():
""" Connect to the PostgreSQL database server """
conn = None
try:
# read connection parameters
params = config()
# connect to the PostgreSQL server
print('Connecting to the PostgreSQL database...')
conn = psycopg.connect(**params)
# create a cursor
cur = conn.cursor()
# execute a statement
print('PostgreSQL database version:')
cur.execute('SELECT version()')
# display the PostgreSQL database server version
db_version = cur.fetchone()
print(db_version)
# close the communication with the PostgreSQL
cur.close()
except (Exception, psycopg.DatabaseError) as error:
print(error)
finally:
if conn is not None:
conn.close()
print('Database connection closed.')
if __name__ == '__main__':
connect()
Now this runs and prints out the Postgresql database version which is all well & great but I'm struggling to figure out how to change the code so that it's more generalized and maybe just creates a cursor?
I need the connect function to basically just connect to the dvdrental database and create a cursor so that I can then use my connection to select from the database in other needed "tasks" -- for example I'd like to be able to create another function like the below:
def select_from_table(cursor, table_name, schema):
cursor.execute(f"SET search_path TO {schema}, public;")
results= cursor.execute(f"SELECT * FROM {table_name};").fetchall()
return results
but I'm struggling with how to just create a connection to the dvdrental database & a cursor so that I'm able to actually fetch data and create pandas tables with it and whatnot.
so it would be like
task 1 is connecting to the database
task 2 is interacting with the database (selecting tables and whatnot)
task 3 is converting the result from 2 into a pandas df
thanks so much for any help!! This is for a project in a class I am taking and I am extremely overwhelmed and have been googling-researching non-stop and seemingly end up nowhere fast.
The fact that you established the connection is honestly the hardest step. I know it can be overwhelming but you're on the right track.
Just copy these three lines from connect into the select_from_table method
params = config()
conn = psycopg.connect(**params)
cursor = conn.cursor()
It will look like this (also added conn.close() at the end):
def select_from_table(cursor, table_name, schema):
params = config()
conn = psycopg.connect(**params)
cursor = conn.cursor()
cursor.execute(f"SET search_path TO {schema}, public;")
results= cursor.execute(f"SELECT * FROM {table_name};").fetchall()
conn.close()
return results

Postgresql - How to open an Autocommit Connection to handle a Vacuum Full

I am trying to achieve the same thing as in earlier question psycopg2: How to execute vacuum postgresql query in python script; however, the recommendation to open an autocommit connection includes a link which is broken.
The below code runs without error BUT the table is not vacuumed.
How does this need to be written to call the Vacuum Full correctly?
#!/usr/bin/python
import psycopg2
from config import config
def connect():
""" Connect to the PostgreSQL database server """
conn = None
try:
# read connection parameters
params = config()
# connect to the PostgreSQL server
conn = psycopg2.connect(**params)
conn.autocommit=1
# create a cursor
cur = conn.cursor()
# execute Vacuum Full
cur.execute('Vacuum Full netsuite_display')
# close the communication with the PostgreSQL
cur.close()
except (Exception, psycopg2.DatabaseError) as error:
print(error)
finally:
if conn is not None:
conn.close()
print('Database connection closed.')
if __name__ == '__main__':
connect()

When to close MySQL connection discord.py

I want to make a discord bot that saves data to a MySQL Database (currently localhost) but the problem is I don't know when to close the database connection.
Currently, when the user enters the command it always creates a new connection to the database but as you can imagine it's very slow to always connect then execute the query then close the connection, and finally after then returning the data.
Here is an example:
def open_case(case_id):
search_query = f"SELECT `CASE_ID`, `USER_ID`, `LINK_REASON`, `LINK_SCREENSHOT` FROM `Report` WHERE CASE_ID ='{case_id}'"`
mydb = mysql.connector.connect(
host = "localhost",
database = "report",
password = "root",
username = "root"
)
cursor = mydb.cursor()
try:
cursor.execute(search_query)
result = cursor.fetchall()
mydb.close()
return result
except:
return print("Error case not found")
mydb.close()
But I'm afraid if I connect to the DB at the beginning the bot crashers or so and then I've never closed the connection to the database.
Is there a way to make it better?
every connection has an idle timeout, that would detect a unused open connection and close it.
But your approach is very good and clean, as it closes the connection
A much simpler approach is to add finqallyas all try catch and ecept will run through it
Also use prepared statements, when using variabels
def open_case(case_id):
search_query = "SELECT `CASE_ID`, `USER_ID`, `LINK_REASON`, `LINK_SCREENSHOT` FROM `Report` WHERE CASE_ID =%s"
mydb = mysql.connector.connect(
host = "localhost",
database = "report",
password = "root",
username = "root"
)
cursor = mydb.cursor()
try:
cursor.execute(search_query,(case_id,))
result = cursor.fetchall()
return result
except:
result = print("Error case not found")
finally:
mydb.close()
open_case(1)

Can't SSH to Google Cloud Compute Engine with sshtunnel

I prepared two VM instance with Compute Engine on GCP.
ServerA: Data processing and read/write to SQL(mysql) on ServerB.
ServerB: SQL Server (f1-micro* This is not Cloud SQL, but normal VM instance.)
Trying to access SSH from A to B in order to read/write DB on ServerB with the code below.
error code
error: ERROR | Problem setting SSH Forwarder up: Couldn't open tunnel localhost:3306 <> localhost:3306 might be in use or destination not reachable
sshtunnel.HandlerSSHTunnelForwarderError: An error occurred while opening tunnels.
#SSH connection
with SSHTunnelForwarder(
('PublicIP of ServerA', 22),
ssh_pkey=SSH_PKEY_PATH,
ssh_username=SSH_USER,
remote_bind_address=('localhost', 3306),
local_bind_address=('localhost', 3306)
) as ssh:
try:
#DB connection
connection = mysql.connector.connect(
host='localhost',
port = 3306,
user=MYSQL_USER,
passwd=MYSQL_PASS,
db=MYSQL_DB,
charset='utf8'
)
# print(connection.is_connected())
# Get Cur
cur = connection.cursor()
sql = "use dbname"
cur.execute(sql)
for i in range(len(sqlList)):
print("DB Access:" + str(sqlList[i]))
sql = str(sqlList[i])
# sql = 'create table test (id int, content varchar(32))'
cur.execute(sql)
sqlOUTPUT = cur.fetchall()
# rows = cur.fetchall()
# for row in rows:
# print(row)
except mysql.connector.Error as err:
print("Something went wrong: {}".format(err))
connection.rollback()
raise err
finally:
#Cur close
cur.close()
# Commit
connection.commit()
#DB Connection close
connection.close()
return sqlOUTPUT
But after "local_bind_address=(localhost, MYSQL_PORT)", an error occurs despite it goes through with the same code and same private key on the shell of B or on VSCode local environment.
I don't understand why it goes through with same code using shell and VSCode although it doesn't work on GCE.
Any help?
You might be able to debug this further and discard any issues with sshtunnel if you try to create the tunnel outside of the script from the client VM, with:
$ gcloud compute ssh server-a --zone=your-zone --ssh-flag='-NL 3306:127.0.0.1:3306' &
Then attempt a connection with:
$ mysql -h 127.0.0.1

Connect to MySQL on GCE using Google Colaboratory

I am trying to access MySQL on GCE VM instance through Google Colaboratory though, unfortunately, it does not work properly but does work in the local VSCode environment.
While executing TestExec.py, it shows SSH Connected so it seems that ssh connection is successfully done though, however, it seems to be stuck on MySQL connection.
Output on Google Colab:
Please help with the solutions/tips?
sqlList = []
sqlList.append("select * from table name;")
HOST = 'ComputeEngine PublicIP'
PORT = 22
USER = 'username'
DBUSER = 'username for db'
KEY_FILE = 'private key file path'
DBNAME = 'dbname'
DBPORT = 3306
SSH_BASTION_ADDRESS = HOST
SSH_PORT = PORT
SSH_USER = USER
SSH_PKEY_PATH = KEY_FILE
MYSQL_HOST = HOST
MYSQL_PORT = 3306
MYSQL_USER = DBUSER
MYSQL_PASS = 'MySQL Login PW'
MYSQL_DB = DBNAME
with SSHTunnelForwarder(
(SSH_BASTION_ADDRESS, SSH_PORT),
ssh_pkey=SSH_PKEY_PATH,
ssh_username=SSH_USER,
# ssh_password=PASSPHRASE,
remote_bind_address=('localhost', MYSQL_PORT),
local_bind_address=('localhost', MYSQL_PORT)
) as ssh:
print("SSH Connected")
print(ssh.local_bind_port)
try:
connection = mysql.connector.connect(
host='localhost',
port = ssh.local_bind_port,
user=MYSQL_USER,
passwd=MYSQL_PASS,
db=MYSQL_DB,
charset='utf8'
)
print(connection.is_connected())
print("DB Connected")
cur = connection.cursor()
sql = "use dbname"
cur.execute(sql)
# rows = cur.fetchall()
# for row in rows:
# print(row)
for i in range(len(sqlList)):
print(sqlList[i])
sql = str(sqlList[i])
# sql = 'create table test (id int, content varchar(32))'
cur.execute(sql)
rows = cur.fetchall()
for row in rows:
print(row)
except mysql.connector.Error as err:
print("Something went wrong: {}".format(err))
connection.rollback()
raise err
finally:
cur.close()
connection.commit()
connection.close()
It seemed the error occurred because of the connection. I have changed the method of connection.
import pymysql.cursors
connection = pymysql.connect(
host='localhost',
user='user',
password='password!',
db='dbname',
charset='utf8mb4',
cursorclass=pymysql.cursors.DictCursor
)

Categories

Resources