Connection and Cursor inside a class in psycopg2

Question

I would like to use the connection and cursor classes inside a custom made class. I would also like to inherit all the methods associated with the connection and cursor classes. I have done some research on it and found docs and this question related to my problem. I have some code that is working partly. i.e I can insert and update the database. However, I cannot select from the database because doing so returns none even when the row is in the database. Here is my code

from datetime import datetime, timedelta
import psycopg2
import psycopg2.extensions


class DataBase():
    """A class used to create tables in the database. Inherits from 
    psycopg2.extensions.connection in order to gain access to the cursor,
    commit, close, and many other features from the pyscopg module.
    """
    def __init__(self):
        self.my_connection = psycopg2.connect(database="public", user="public",
                                  password="general", host="127.0.0.1",
                                  port="5432")
        self.my_cursor = self.my_connection.cursor()

    def query_database(self, sql_statement, *args):
        return self.my_cursor.execute(sql_statement, *args)

    def commit_query(self):
        return self.my_connection.commit()

    def fetch_one(self, sql_statement, *args):
        result = self.query_database(sql_statement, *args)
        if result is None:
            return False
        return result.fetchone()

    def fetch_all(self, sql_statement, *args):
        result = self.query_database(sql_statement, *args)
        if result is None:
            return False
        return result.fetchall()

    def __del__(self):
        self.my_cursor.close()
        self.my_connection.close()


############################################################################
class CreateTables(DataBase):
    def create_user_table(self):
        """Helper function used to create the user_table"""
        sql_statement = '''CREATE TABLE IF NOT EXISTS USERS
                (ID                 SERIAL    PRIMARY KEY,
                FIRSTNAME           TEXT      NOT NULL,
                LASTNAME            TEXT      NOT NULL,
                USERNAME            TEXT      NOT NULL UNIQUE,
                EMAIL               TEXT      NOT NULL UNIQUE,
                PASSWORD            TEXT      NOT NULL,
                DATETIMEREGISTERED  TIMESTAMP NOT NULL);'''
        user_table = DataBase.query_database(self, sql_statement)
        DataBase.commit_query(self)
        return user_table

    def create_entries_table(self):
        """Helper function used to create an entries table."""
        sql_statement = '''CREATE TABLE IF NOT EXISTS ENTRIES
                        (ID             SERIAL      PRIMARY KEY,
                        TITLE           TEXT        NOT NULL,
                        DRINK           TEXT        NOT NULL,
                        DATEOFORDER     TIMESTAMP   NOT NULL,
                        TIMETODELIVERY  TIMESTAMP   NOT NULL,
                        SETREMINDER     TIMESTAMP   NOT NULL,
                        USERID      INT REFERENCES USERS ON DELETE CASCADE);'''
        entries_table = DataBase.query_database(self, sql_statement)
        DataBase.commit_query(self)
        print("entries table created.")
        return entries_table

# test = CreateTables() This is working well
# print(test.create_entries_table())

#####################################################################

class OperateDatabase(CreateTables):

    def create_user(self, email, username, *args):
        """Helper function used to create a user"""
        sql_statement = """SELECT ID FROM USERS WHERE EMAIL = %s OR
                         USERNAME = %s;"""
        user_in_database = CreateTables.fetch_one(self, sql_statement,
                                                  (email, username,))
        print("the user in database is :>>", user_in_database)

        sql_statement2 = """INSERT INTO USERS (FIRSTNAME, LASTNAME, USERNAME,
                         EMAIL, PASSWORD, DATETIMEREGISTERED)
                         VALUES (%s, %s, %s, %s, %s, %s);"""
        if not user_in_database:
            CreateTables.query_database(self,sql_statement2, *args)
            CreateTables.commit_query(self)
            return True
        return False

data = ("Jkdai", "Jkdal", "Jkdai", "jkdai@gmail.com", "password", datetime.now())
test = OperateDatabase()
print(test.create_user("jkdai@gmail.com", "jkdai", data))
#Inserts the user the very first time implying the insert statement is working
#raises an integrity error the second time implying the select statement is not working. 
#Also the print statement defaults to false when it is supposed to return the user's id.

score 5 · Accepted Answer · answered Aug 17 '18 at 12:36

cursor.execute() returns a vendor-defined value (it's not specified in the db-api spec), and for pyscopg2 it is actually documented as returning None indeed, so this:

def query_database(self, sql_statement, *args):
    return self.my_cursor.execute(sql_statement, *args)

def fetch_one(self, sql_statement, *args):
    result = self.query_database(sql_statement, *args)
    if result is None:
        return False
    return result.fetchone()

will obviously not work as you expect. You could return self.my_cursor() from query_database() instead, ie:

def query_database(self, sql_statement, *args):
    self.my_cursor.execute(sql_statement, *args)
    return self.my_cursor

def fetch_one(self, sql_statement, *args):
    cursor = self.query_database(sql_statement, *args)
    return cursor.fetchone()

but there's a fundamental flaw with this code, which is that it's not reentrant (nor thread safe FWIW). Actually, your class should NOT store a cursor as part of it's state and reuse it over and over (db-api cursors are not meant to be used that way) but create a new cursor for every operation (which is the indented use).

Also, you should not rely on __del__(self) to close your connection. The __del__() method is not a proper C++/Java style finalizer and is not even garanteed to be called when the object is collected. Actually, trying to wrap a db connection and cursor in a class is usually not a good idea, at least not this way.

Connection and Cursor inside a class in psycopg2

1 Answers1