Update contextmanagers.py

Changes after testing
This commit is contained in:
FelipeV34 2021-08-19 01:00:35 -05:00
parent a5a18ea75f
commit 40534b1e55

View File

@ -1,29 +1,29 @@
"""psycopg context managers for connection and cursor objects.""" """psycopg context managers for connection and cursor objects."""
import psycopg2
from contextlib import contextmanager from contextlib import contextmanager
from extensions import connect, cursor
@contextmanager @contextmanager
def connection(**kwargs): def connect(**kwargs):
"""Context manager to yield and close a db connection.""" """Context manager to yield and close a db connection."""
connection = connect(**kwargs) conn = psycopg2.connect(**kwargs)
try: try:
yield connection yield conn
finally: finally:
connection.close() conn.close()
@contextmanager @contextmanager
def cursor(connection): def cursor(conn):
""" """
Context manager to yield and close a cursor given an existing db Context manager to yield and close a cursor given an existing db
connection. connection.
""" """
cursor = connection.cursor() cursor = conn.cursor()
try: try:
yield cursor yield cursor
@ -33,17 +33,17 @@ def cursor(connection):
@contextmanager @contextmanager
def connection_cursor(**kwargs): def connect_cursor(**kwargs):
""" """
Context manager to yield a cursor and close it and its connection for Context manager to yield a cursor and close it and its connection for
purposes or performing a single db call. purposes or performing a single db call.
""" """
connection = connect(**kwargs) conn = psycopg2.connect(**kwargs)
cursor = connection.cursor() cur = conn.cursor()
try: try:
yield cursor yield cur
finally: finally:
cursor.close() cur.close()
connection.close() conn.close()