Queries for storing and get information from postgreSQL

Question

I have been working with postgresql where I am trying to get information from database, update, exists and much more. I have created a context manager for commit and to close whenever it is in finished.

This project is about storing some variables (mostly regarding store and link)

#!/usr/bin/python3
# -*- coding: utf-8 -*-
from typing import List, Set, Tuple

import pendulum
import psycopg2
import psycopg2.extras

from config import configuration

DATABASE_CONNECTION = {
    "host": configuration.path.database.environment,
    "database": configuration.postgresql.database,
    "user": configuration.postgresql.user,
    "password": configuration.postgresql.password,
}

ps_connection = psycopg2.connect(**DATABASE_CONNECTION)


class QuickConnection:
    """
        Function that commits and closes when its done.
        Rollbacks if an error happens
    """

    def __init__(self):
        self.ps_cursor = ps_connection.cursor(cursor_factory=psycopg2.extras.DictCursor)
        ps_connection.autocommit = True

    def __enter__(self):
        return self.ps_cursor

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


# ------------------------------------------------------------------------------- #

# Exists functions

# ------------------------------------------------------------------------------- #

def link_exists(store: str, link: str) -> bool:
    query_params = {
        "store": store,
        "link": link
    }

    sql_query = "SELECT EXISTS (SELECT 1 FROM store_items WHERE store=%(store)s AND link=%(link)s);"

    with QuickConnection() as ps_cursor:
        ps_cursor.execute(sql_query, query_params)
        exists, = ps_cursor.fetchone()
        return exists


# ------------------------------------------------------------------------------- #

# Insert / Update / Delete / Return value from database

# ------------------------------------------------------------------------------- #

def store_exists(store: str) -> bool:
    query_params = {
        "store": store
    }

    sql_query = "INSERT INTO store_config (store) SELECT %(store)s WHERE NOT EXISTS (SELECT 1 FROM store_config WHERE store = %(store)s);"

    with QuickConnection() as ps_cursor:
        ps_cursor.execute(sql_query, query_params)
        return bool(ps_cursor.rowcount)


def register_product(store: str, pageData) -> bool:
    query_params = {
        "store": store,
        "name": pageData.name,
        "link": pageData.link,
        "image": pageData.image,
        "visible": "yes",
        "added_date": pendulum.now('Europe/Stockholm').format('YYYY-MM-DD HH:mm:ss.SSSSSS')
    }

    sql_query = "INSERT INTO store_items (store, name, link, image, visible, added_date) VALUES (%(store)s, %(name)s, %(link)s, %(image)s, %(visible)s, %(added_date)s);"

    with QuickConnection() as ps_cursor:
        ps_cursor.execute(sql_query, query_params)
        return bool(ps_cursor.rowcount)


# ------------------------------------------------------------------------------- #

# Get all functions

# ------------------------------------------------------------------------------- #

def get_all_active_links(store: str) -> List[str]:
    query_params = {
        "store": store,
        "visible": "yes"
    }

    sql_query = "SELECT link FROM store_items WHERE store = %(store)s AND visible = %(visible)s;"

    with QuickConnection() as ps_cursor:
        ps_cursor.execute(sql_query, query_params)
        return [links["link"] for links in ps_cursor]


def get_feed_urls(store: str) -> List[str]:
    query_params = {
        "store": store
    }

    sql_query = "SELECT link FROM feed_urls WHERE store = %(store)s;"

    with QuickConnection() as ps_cursor:
        ps_cursor.execute(sql_query, query_params)
        return [links["link"] for links in ps_cursor]


def get_all_stores() -> List[str]:
    sql_query = "SELECT store FROM store_config;"

    with QuickConnection() as ps_cursor:
        ps_cursor.execute(sql_query)
        return [stores["store"] for stores in ps_cursor]

I wonder if there is anything I could improve :)

Answer 1

ps_connection existing as a global is ungood. QuickConnection on its own seems like a reasonable way to spin up a cursor, but would be better-modelled with a connection instance that it owns as a member, not referring to a global. This would also (reasonably) require that your query methods accept a QuickConnection as a parameter.

Consider using a connection pool instead of a single connection.

    ps_connection.autocommit = True

seems counter to your goal of having a transaction-managed cursor that supports rollback. Don't you want autocommit to be disabled, and have a commit in your exit function if no errors are found?

Whether by threading or coroutines, etc., it may be possible to have two __enter__ calls on one QuickConnection before any exit is seen. You should protect against this by throwing on an enter when there is already an active cursor. If you need more than one cursor at a time, this context management approach on its own won't be practical.

link_exists is reasonable, but compare that to the similarly-named store_exists that does a very different operation. The latter needs to be renamed to indicate that it's effectively a fail-safe insert.

return [links["link"] for links in ps_cursor]

doesn't need to be a list comprehension, and you can make get_all_active_links a generator:

for link in ps_cursor:
    yield link['link']

This should stream results to the caller as they're loaded from the database, rather than jamming them all into memory at once.

Answer 2

The more I look at the code and the more I think you could have used an ORM, for instance SQL Alchemy. And nothing stops you from building the same functions on top of the ORM. Your functions could be replaced by one-liner filter clauses. You would write something like:

store = session.query(Store).filter(Store.id=10)

This would get you one row (or none). The benefit of this approach is that you will always be dealing with rows whereas your functions have different declarations and return different types. The keyword here is consistency.

To count stores you could do something like:

session.query(Store).count()

and you can still derive a boolean depending on whether count is zero or greater.

All that code does nothing fancy, just filter data from tables on certain attributes or insert records. As pointed out by @Reinderien, transaction management is important and the whole point of using a DB class.

But this is important when adding/updating more than one row in a single batch of related data. If you are adding only one row at a time, it doesn't make a difference - autocommit is implied.

When inserting a new row, it would make sense to return the ID of newly created row... provided you have ID columns. If you're creating a new store, you might want to retrieve that ID to then add a bunch of products linked to that store by ID. And of course wrap the whole operation in a transaction, so that in case of failure the whole batch is rollbacked and you have no orphan records polluting the DB.

But there doesn't seem to be a primary key in your tables and that is a problem. At least that's my assumption. How will you unambiguously differentiate stores or products that have the same name ? A good data model is important because performance depends on good design. A big DB with lots of records will not scale well. Searching will become slower. Indexes are always welcome but a sound table structure is still important.

The point of using functions is to avoid repetition but you have not achieved this, in fact they all repeat the same stuff more or less:

with QuickConnection() as ps_cursor:
        ps_cursor.execute(sql_query, query_params)

So these functions do not bring any added value really. Or you should refactor your code to wrap the stuff that is repetitive in one function and reuse it. I am afraid you are going to add more and more functions, each designed to address a specific need. But that will only amplify repetition and make the code bigger for no benefit.

You use pendulum.now('Europe/Stockholm') to record the creation date. But timestamps without timezone context are ambiguous. I would consider using a timestamptz datatype for the column. Then PG will store timestamps as UTC and you can render datetime values to the client in the desired timezone by applying proper offset (consider daylight saving time too).

At some point you might change your time zone for whatever reason. For instance if you have a server and decide to host your site at a new location. The server will usually be preconfigured to the local TZ. Then the new datetime values will not have the same meaning as the older ones because the reference TZ is not the same. Hence the benefit of storing all datetime values as UTC.

PG has a boolean type which would be a better candidate than string for the field "visible".

Answer 3

Given that you're using Python 3, you can safely leave out the encoding line since the default encoding is UTF-8.

-*- coding: utf-8 -*-

Unless you work in a multi-encoding environment or the default encoding on your system is not utf8.