pantry-track/application/administration/administration_models.py

from application.database_postgres.UsersModel import UsersModel
from application.database_postgres.RolesModel import RolesModel
from application.database_postgres.BaseModel import DatabaseError, tupleDictionaryFactory
import config
import psycopg2

class ExtendedRolesModel(RolesModel):
    @classmethod
    def select_by_site_uuid(self, payload, convert=True, conn=None):
        roles = ()
        self_conn = False
        select_roles_sql = f"SELECT * FROM roles WHERE role_site_uuid = %(site_uuid)s::uuid;"
        try:
            if not conn:
                database_config = config.config()
                conn = psycopg2.connect(**database_config)
                conn.autocommit = True
                self_conn = True

            with conn.cursor() as cur:
                cur.execute(select_roles_sql, payload)
                rows = cur.fetchall()
                if rows and convert:
                    roles = [tupleDictionaryFactory(cur.description, role) for role in rows]
                elif rows and not convert:
                    roles = rows
            
            if self_conn:
                conn.close()

            return roles
        except Exception as error:
            raise DatabaseError(error, payload, select_roles_sql)

class ExtendedUsersModel(UsersModel):
    @classmethod
    def add_admin_user(self, payload:dict, convert=True, conn=None):
        admin_user = ()
        self_conn = False
        sql = f"""INSERT INTO users (user_name, user_password, user_email, user_row_type)  
        VALUES (%(user_name)s, %(user_password)s, %(user_email)s, %(user_row_type)s) ON CONFLICT (user_name) DO UPDATE SET user_name = excluded.user_name RETURNING *;"""
        try:
            if not conn:
                database_config = config.config()
                conn = psycopg2.connect(**database_config)
                conn.autocommit = True
                self_conn = True

            with conn.cursor() as cur:
                cur.execute(sql, payload)
                rows = cur.fetchone()
                if rows and convert:
                    admin_user = tupleDictionaryFactory(cur.description, rows)
                elif rows and not convert:
                    admin_user = rows   

            if self_conn:
                conn.commit()
                conn.close()

            return admin_user
        except Exception as e:
            DatabaseError(str(e), payload, sql)
    
    @classmethod
    def update_roles(self, payload, convert=True, conn=None):
        """ payload: {'role_uuid': x,} """
        self_conn = False
        try:
            if not conn:
                database_config = config.config()
                conn = psycopg2.connect(**database_config)
                conn.autocommit = True
                self_conn = True

            select_sql = f"SELECT users.user_uuid FROM users WHERE user_roles @> ARRAY[%(role_uuid)s::uuid];"
            with conn.cursor() as cur:
                cur.execute(select_sql, payload)
                users = tuple([row[0] for row in cur.fetchall()])

            update_sql = f"UPDATE users SET user_roles = array_remove(user_roles, %(role_uuid)s::uuid) WHERE user_uuid = %(user_uuid)s::uuid;"
            with conn.cursor() as cur:
                for user_uuid in users:
                    cur.execute(update_sql, {'role_uuid': payload['role_uuid'], 'user_uuid': user_uuid})
            
            if self_conn:
                conn.commit()
                conn.close()

        except Exception as error:
            raise error

    @classmethod
    def update_sites(self, payload, convert=True, conn=None):
        """ payload: {'site_uuid',} """
        self_conn = False
        try:
            if not conn:
                database_config = config.config()
                conn = psycopg2.connect(**database_config)
                conn.autocommit = True
                self_conn = True

            select_sql = f"SELECT users.user_uuid FROM users WHERE user_sites @> ARRAY[%(site_uuid)s::uuid];"
            with conn.cursor() as cur:
                cur.execute(select_sql, payload)
                user = tuple([row[0] for row in cur.fetchall()])

            update_sql = f"UPDATE users SET user_sites = array_remove(user_sites, %(site_uuid)s::uuid) WHERE user_uuid = %(user_uuid)s::uuid;"
            with conn.cursor() as cur:
                for user_uuid in user:
                    cur.execute(update_sql, {'site_uuid': payload['site_uuid'], 'user_uuid': user_uuid})
            
            if self_conn:
                conn.commit()
                conn.close()

        except Exception as error:
            raise error

    @classmethod
    def update_user_site_roles(self, payload, convert=True, conn=None):
        """ payload (tuple): (site_uuid, role_uuid, user_uuid) """
        sql = f"UPDATE users SET user_sites = user_sites || %(site_uuid)s::uuid, user_roles = user_roles || %(role_uuid)s::uuid WHERE user_uuid=%(user_uuid)s RETURNING *;"
        user = ()
        self_conn = False
        try:
            if not conn:
                database_config = config.config()
                conn = psycopg2.connect(**database_config)
                conn.autocommit = True
                self_conn = True
            
            with conn.cursor() as cur:
                cur.execute(sql, payload)
                rows = cur.fetchone()
                if rows and convert:
                    user = tupleDictionaryFactory(cur.description, rows)
                elif rows and not convert:
                    user = rows
            
            if self_conn:
                conn.commit()
                conn.close()

            return user
        except Exception as error:
            raise DatabaseError(error, payload, sql)