pluto/handler/db.py

138 lines
3.5 KiB
Python

# 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
"""
messages = []
for tx, reference in get_ref_map().items():
if reference not in references:
if get_subclass(reference) == subclass:
logmessage = f"{reference}: Archiving trade reference. TX: {tx}"
messages.append(logmessage)
log.info(logmessage)
r.rename(f"trade.{tx}.reference", f"archive.trade.{tx}.reference")
r.rename(f"trade.{reference}", f"archive.trade.{reference}")
return messages