How to create a function that connects, executes and disconnects

user_13019246 0
Viewed 312

Today I have spent lots of time learning abit slowly with Postgres and I have been creating a code that does some stuff with the database such as insert, select etc etc.

I have realized that most of my code is copy paste when it comes to Connect & Disconnet and I know some people do not like it also depends on what I am doing so before people gets mad at me, do not take this as a bad code but more that I would like to imrpove of course <3

What I have done so far is:

import psycopg2
import psycopg2.extras
from loguru import logger

DATABASE_CONNECTION = {
    "host": "TEST",
    "database": "TEST",
    "user": "TEST",
    "password": "TEST"
}

def register_datas(store, data):
    """
    Register a data to database
    :param store:
    :param data:
    :return:
    """

    ps_connection = psycopg2.connect(**DATABASE_CONNECTION)
    ps_cursor = ps_connection.cursor()
    ps_connection.autocommit = True

    sql_update_query = "INSERT INTO public.store_items (store, name) VALUES (%s, %s);"

    try:

        data_tuple = (store, data["name"])

        ps_cursor.execute(sql_update_query, data_tuple)

        has_registered = ps_cursor.rowcount

        ps_cursor.close()
        ps_connection.close()

        return bool(has_registered)

    except (Exception, psycopg2.DatabaseError) as error:
        logger.exception("Error: %s" % error)
        ps_connection.rollback()
        ps_cursor.close()
        ps_connection.close()
        return False

def get_all_keywords(keywords):
    """
    Get all keywords
    :param positive_or_negative:
    :return:
    """

    ps_connection = psycopg2.connect(**DATABASE_CONNECTION)
    ps_cursor = ps_connection.cursor(cursor_factory=psycopg2.extras.DictCursor)

    sql_update_query = "SELECT keyword FROM public.keywords WHERE filter_type = %s;"

    try:

        data_tuple = (keywords,)
        ps_cursor.execute(sql_update_query, data_tuple)

        all_keywords = [keyword["keyword"] for keyword in ps_cursor]

        ps_cursor.close()
        ps_connection.close()

        return all_keywords

    except (Exception, psycopg2.DatabaseError) as error:
        logger.exception("Error: %s" % error)
        ps_connection.rollback()
        ps_cursor.close()
        ps_connection.close()
        return []


def check_if_store_exists(store):
    """
    Check if the store exists in database
    :param store:
    :return:
    """

    ps_connection = psycopg2.connect(**DATABASE_CONNECTION)
    ps_cursor = ps_connection.cursor()

    sql_update_query = "SELECT store FROM public.store_config WHERE store = %s;"

    try:

        data_tuple = (store,)
        ps_cursor.execute(sql_update_query, data_tuple)

        exists = bool(ps_cursor.fetchone())

        ps_cursor.close()
        ps_connection.close()

        return exists

    except (Exception, psycopg2.DatabaseError) as error:
        logger.exception("Error: %s" % error)
        ps_connection.rollback()
        ps_cursor.close()
        ps_connection.close()
        return []

and I do see that I have same code where I do:

ps_connection = psycopg2.connect(**DATABASE_CONNECTION)
ps_cursor = ps_connection.cursor()
...
...
...
...
ps_cursor.close()
ps_connection.close()
return data

and to shorter this code, I wonder if its possible to do a function where I call the function that connects -> lets me do the query/execution -> close the connection and then return the data I want to return?

postgresql python
1 Answers

Context manager

There is a pattern built-in Python. It is named context manager. It has two purposes:

  • Do operations before and after some logic in with block.
  • Catch errors inside with block and allow to handle it in a custom way.

To create a context manager, you can go in either of two ways. One (I like it more) is to define a class satisfying the following protocol:

  • __enter__(self)
  • __exit__(self, exc_type, exc_value, traceback)

Any class that satisfies the protocol can be used in with statement and works as a context manager. Accordingly to Python duck-typing principles, the interpreter "knows" how to use the class in with statement.

Example:

class QuickConnection:  
    def __init__(self):
        self.ps_connection = psycopg2.connect(**DATABASE_CONNECTION)
        self.ps_cursor = ps_connection.cursor(cursor_factory=psycopg2.extras.DictCursor)

    def __enter__(self):
        return self.ps_cursor

    def __exit__(self, err_type, err_value, traceback):
        if err_type and err_value:
            self.ps_connection.rollback()
        self.ps_cursor.close()
        self.ps_connection.close()
        return False

Return value of __exit__ method does matter. If True returned then all the errors happened in the with block suppressed. If False returned then the errors raised at the end of the __exit__ execution. The return values of __exit__ is better to keep explicit since the feature itself is not that obvious.

Or use contextlib.contextmanager decorator

from contextlib import contextmanager

@contextmanager
def quick_connection():
    ps_connection = psycopg2.connect(**DATABASE_CONNECTION)
    ps_cursor = ps_connection.cursor(cursor_factory=psycopg2.extras.DictCursor)
    try:
        yield ps_cursor
    except Exception: # don't do this, catch specific errors instead.
        ps_connection.rollback()
        raise
    finally:
        ps_cursor.close()
        ps_connection.close()

Usage

with QuickConnection() as ps_cursor:
    data_tuple = (store, data["name"])
    ps_cursor.execute(sql_update_query, data_tuple)
    has_registered = ps_cursor.rowcount

With the context manager, you can reuse a connection, be sure that it is closed. Also, you can catch and handle errors related to your DB operations in the context manager. Usage of context managers is compact and readable, I believe in the end this is the goal.

user_6296850
Related