Created flat2mongo, a utility script to move database from flatfile to mongodb

Automatically encode/decode stock tickers with a dot so they can be stored in mongodb
Remove print debug code
This commit is contained in:
socks 2021-03-08 14:20:15 +00:00
parent a893cc1c04
commit 689c32abf0
3 changed files with 50 additions and 6 deletions

26
flat2mongo.py Normal file
View file

@ -0,0 +1,26 @@
# Utility script to port from flatfile to database
import config as cfg
import flatfile as f
import mongo as m
import json
with open(cfg.database) as f:
data = json.loads(f.read())
for user in data.keys():
print("\n\nUser:", user)
print("Portfolio:")
for stock in data[user]["portfolio"].keys():
ticker = stock
amount = data[user]["portfolio"][stock]
m.add_stock(user, ticker, amount)
print("Watchlist:")
for stock in data[user]["watchlist"].keys():
ticker = stock
amount = data[user]["watchlist"][stock]
m.watch(user, ticker, amount)
print("[*] Done!")

View file

@ -6,7 +6,6 @@ import config as cfg
from decimal import Decimal
import typing
print(cfg.use_mongodb)
if cfg.use_mongodb:
import mongo as database

View file

@ -17,6 +17,14 @@ class Watchlist(Document):
stocks = DictField(default={})
def encodeKey(key):
return key.replace(".", "\\u002e")
def decodeKey(key):
return key.replace("\\u002e", ".")
def user_check(user):
# Will create empty portfolio and watchlist for new users
if not Portfolio.objects(user=user):
@ -32,7 +40,11 @@ def user_check(user):
def get_stocks(user):
user_check(user)
return Portfolio.objects.get(user=user).stocks
stocks = Portfolio.objects.get(user=user).stocks
decoded_stocks = {}
for stock, amount in stocks.items():
decoded_stocks[decodeKey(stock)] = amount.to_decimal()
return decoded_stocks
def add_stock(user, stock, amount):
@ -43,7 +55,7 @@ def add_stock(user, stock, amount):
if not (yfi.stock_exists(stock)):
# Stock ticker does not exist on yahoo finance
return False
stock = encodeKey(stock)
if stock in portfolio.stocks.keys():
portfolio.stocks[stock] = Decimal128(
amount + portfolio.stocks[stock].to_decimal()
@ -56,6 +68,7 @@ def add_stock(user, stock, amount):
def delete_stock(user, stock):
user_check(user)
stock = encodeKey(stock)
portfolio = Portfolio.objects.get(user=user)
if stock in portfolio.stocks.keys():
del portfolio.stocks[stock]
@ -67,24 +80,30 @@ def delete_stock(user, stock):
def get_watchlist(user):
user_check(user)
return Watchlist.objects.get(user=user).stocks
wlist = Watchlist.objects.get(user=user).stocks
decoded_stocks = {}
for stock, val in wlist.items():
decoded_stocks[decodeKey(stock)] = val.to_decimal()
return decoded_stocks
def watch(user, stock, est_price="0"):
user_check(user)
watchlist = Watchlist.objects.get(user=user)
if not (yfi.stock_exists(stock)):
# Stock ticker does not exist on yahoo finance
return False
watchlist.stocks[stock] = str(Decimal128(est_price))
stock = encodeKey(stock)
watchlist.stocks[stock] = Decimal128(est_price)
watchlist.save()
return True
def unwatch(user, stock):
user_check(user)
stock = stock.replace(".", "\\u002e")
watchlist = Watchlist.objects.get(user=user)
if stock in watchlist.stocks.keys():
del watchlist.stocks[stock]