data-storage/src/piracyshield_data_storage/account/storage.py

from piracyshield_data_storage.database.arangodb.document import DatabaseArangodbDocument

from arango.cursor import Cursor

class AccountStorage(DatabaseArangodbDocument):

    collection_name = None

    collection_instance = None

    def __init__(self, collection_name: str):
        super().__init__()

        self.collection_name = collection_name

        self.collection_instance = self.collection(self.collection_name)

    def insert(self, document: dict) -> dict | Exception:
        try:
            return self.collection_instance.insert(document)

        except:
            raise AccountStorageCreateException()

    def get(self, account_id: str) -> Cursor | Exception:
        aql = f"""
            FOR document IN {self.collection_name}
            FILTER document.account_id == @account_id
            RETURN {{
                'account_id': document.account_id,
                'name': document.name,
                'email': document.email,
                'role': document.role,
                'is_active': document.is_active,
                'created_by': document.created_by
            }}
        """

        try:
            return self.query(aql, bind_vars = {
                'account_id': account_id
            })

        except:
            raise AccountStorageGetException()

    def get_complete(self, account_id: str) -> Cursor | Exception:
        aql = f"""
            FOR document IN {self.collection_name}

                FILTER document.account_id == @account_id

            RETURN document
        """

        try:
            return self.query(aql, bind_vars = {
                'account_id': account_id
            })

        except:
            raise AccountStorageGetException()

    def get_all(self) -> Cursor | Exception:
        aql = f"""
            FOR document IN {self.collection_name}
            RETURN {{
                'account_id': document.account_id,
                'name': document.name,
                'email': document.email,
                'is_active': document.is_active,
                'role': document.role
            }}
        """

        try:
            return self.query(aql)

        except:
            raise AccountStorageGetException()

    def get_total(self) -> int | Exception:
        try:
            return self.collection_instance.count()

        except:
            raise AccountStorageGetException()

    def exists_by_identifier(self, identifier: str) -> Cursor | Exception:
        """
        Checks if an account with this identifier is in the collection.

        :param value: a valid account identifier.
        :return: cursor with the requested data.
        """

        aql = f"""
            FOR document IN {self.collection_name}

            FILTER document.account_id == @identifier

            LIMIT 1

            RETURN document
        """

        try:
            return self.query(aql, bind_vars = {
                'identifier': identifier
            })

        except:
            raise AccountStorageGetException()

    def set_flag(self, account_id: str, flag: str, value: any) -> Cursor | Exception:
        """
        Sets a flag with a status.

        :param account_id: account identifier.
        :param flag: the flag to update.
        :param value: any value.
        :return: the number of updated rows.
        """

        aql = f"""
            FOR document IN {self.collection_name}

                FILTER document.account_id == @account_id

            UPDATE {{
                '_key': document._key,
                'flags': {{
                    @flag: @value
                }}
            }} IN {self.collection_name}

            RETURN NEW
        """

        try:
            affected_rows = self.query(
                aql,
                bind_vars = {
                    'account_id': account_id,
                    'flag': flag,
                    'value': value
                },
                count = True
            )

            return affected_rows

        except:
            raise AccountStorageUpdateException()

    def change_password(self, account_id: str, password: str) -> Cursor | Exception:
        """
        Changes the accounts' password.

        :param account_id: account identifier.
        :param password: the encoded password.
        :return: the number of updated rows.
        """

        aql = f"""
            FOR document IN {self.collection_name}

                FILTER document.account_id == @account_id

            UPDATE {{
                '_key': document._key,
                'password': @password
            }} IN {self.collection_name}

            RETURN NEW
        """

        try:
            affected_rows = self.query(
                aql,
                bind_vars = {
                    'account_id': account_id,
                    'password': password
                },
                count = True
            )

            return affected_rows

        except:
            raise AccountStorageUpdateException()

    def update_status(self, account_id: str, value: bool) -> Cursor | Exception:
        """
        Updates the account status.

        :param account_id: account identifier.
        :param is_active: account status.
        :return: the number of updated rows.
        """

        aql = f"""
            FOR document IN {self.collection_name}

                FILTER document.account_id == @account_id

            UPDATE {{
                '_key': document._key,
                'is_active': @value
            }} IN {self.collection_name}

            RETURN NEW
        """

        try:
            affected_rows = self.query(
                aql,
                bind_vars = {
                    'account_id': account_id,
                    'value': value
                },
                count = True
            )

            return affected_rows

        except:
            raise AccountStorageUpdateException()

    def remove(self, account_id: str) -> Cursor | Exception:
        aql = f"""
            FOR document IN {self.collection_name}

            FILTER document.account_id == @account_id
            REMOVE document IN {self.collection_name}

            RETURN OLD
        """

        try:
            affected_rows = self.query(
                aql,
                bind_vars = {
                    'account_id': account_id
                },
                count = True
            )

            return affected_rows

        except:
            raise AccountStorageRemoveException()

class AccountStorageCreateException(Exception):

    """
    Cannot create the account.
    """

    pass

class AccountStorageGetException(Exception):

    """
    Cannot get the account.
    """

    pass

class AccountStorageUpdateException(Exception):

    """
    Cannot update the account.
    """

    pass

class AccountStorageRemoveException(Exception):

    """
    Cannot remove the account.
    """

    pass