pluto/handler/db.py

# Other library imports
from redis import StrictRedis

# Project imports
from settings import settings
import util

log = util.get_logger("db")

# Define the Redis endpoint to the socket
r = StrictRedis(unix_socket_path=settings.DB.RedisSocket, db=int(settings.DB.DB))


def get_refs():
    """
    Get all reference IDs for trades.
    :return: list of trade IDs
    :rtype: list
    """
    references = []
    ref_keys = r.keys("trade.*.reference")
    for key in ref_keys:
        references.append(r.get(key))
    return util.convert(references)


def tx_to_ref(tx):
    """
    Convert a trade ID to a reference.
    :param tx: trade ID
    :type tx: string
    :return: reference
    :rtype: string
    """
    refs = get_refs()
    for reference in refs:
        ref_data = util.convert(r.hgetall(f"trade.{reference}"))
        if not ref_data:
            continue
        if ref_data["id"] == tx:
            return reference


def ref_to_tx(reference):
    """
    Convert a reference to a trade ID.
    :param reference: trade reference
    :type reference: string
    :return: trade ID
    :rtype: string
    """
    ref_data = util.convert(r.hgetall(f"trade.{reference}"))
    if not ref_data:
        return False
    return ref_data["id"]


def get_ref_map():
    """
    Get all reference IDs for trades.
    :return: dict of references keyed by TXID
    :rtype: dict
    """
    references = {}
    ref_keys = r.keys("trade.*.reference")
    for key in ref_keys:
        tx = util.convert(key).split(".")[1]
        references[tx] = r.get(key)
    return util.convert(references)


def get_ref(reference):
    """
    Get the trade information for a reference.
    :param reference: trade reference
    :type reference: string
    :return: dict of trade information
    :rtype: dict
    """
    ref_data = r.hgetall(f"trade.{reference}")
    ref_data = util.convert(ref_data)
    if "subclass" not in ref_data:
        ref_data["subclass"] = "agora"
    if not ref_data:
        return False
    return ref_data


def get_tx(tx):
    """
    Get the transaction information for a transaction ID.
    :param reference: trade reference
    :type reference: string
    :return: dict of trade information
    :rtype: dict
    """
    tx_data = r.hgetall(f"tx.{tx}")
    tx_data = util.convert(tx_data)
    if not tx_data:
        return False
    return tx_data


def get_subclass(reference):
    obj = r.hget(f"trade.{reference}", "subclass")
    subclass = util.convert(obj)
    return subclass


def del_ref(reference):
    """
    Delete a given reference from the Redis database.
    :param reference: trade reference to delete
    :type reference: string
    """
    tx = ref_to_tx(reference)
    r.delete(f"trade.{reference}")
    r.delete(f"trade.{tx}.reference")


def cleanup(subclass, references):
    """
    Reconcile the internal reference database with a given list of references.
    Delete all internal references not present in the list and clean up artifacts.
    :param references: list of references to reconcile against
    :type references: list
    """
    for tx, reference in get_ref_map().items():
        if reference not in references:
            if get_subclass(reference) == subclass:
                log.info(f"Archiving trade reference: {reference} / TX: {tx}")
                r.rename(f"trade.{tx}.reference", f"archive.trade.{tx}.reference")
                r.rename(f"trade.{reference}", f"archive.trade.{reference}")