feat: Add migration script to include direction field in trades table

This commit is contained in:
Bobby (aider) 2025-02-11 07:07:28 -08:00
parent 483fcf7ba9
commit 8c11af9c05

View File

@ -0,0 +1,85 @@
from db.db_connection import create_client
from datetime import datetime
def migrate_trades():
"""Add direction field to existing trades"""
with create_client() as client:
# First get all trades
query = "SELECT * FROM stock_db.trades"
trades = client.query(query).result_rows
# Get column names
columns = ['id', 'position_id', 'ticker', 'entry_date', 'shares', 'entry_price',
'target_price', 'stop_loss', 'strategy', 'order_type', 'followed_rules',
'entry_reason', 'exit_price', 'exit_date', 'exit_reason', 'notes', 'created_at']
# Create new table with direction field
client.command("""
CREATE TABLE IF NOT EXISTS stock_db.trades_new (
id UInt32,
position_id String,
ticker String,
entry_date DateTime,
shares UInt32,
entry_price Float64,
target_price Nullable(Float64),
stop_loss Nullable(Float64),
strategy Nullable(String),
order_type String,
direction String,
followed_rules Nullable(UInt8),
entry_reason Nullable(String),
exit_price Nullable(Float64),
exit_date Nullable(DateTime),
exit_reason Nullable(String),
notes Nullable(String),
created_at DateTime DEFAULT now()
) ENGINE = MergeTree()
ORDER BY (position_id, id, entry_date)
""")
# Migrate data
for trade in trades:
trade_dict = dict(zip(columns, trade))
# Determine direction based on exit_price
direction = 'sell' if trade_dict['exit_price'] else 'buy'
# Insert into new table
query = f"""
INSERT INTO stock_db.trades_new (
id, position_id, ticker, entry_date, shares, entry_price, target_price, stop_loss,
strategy, order_type, direction, followed_rules, entry_reason, exit_price, exit_date,
exit_reason, notes, created_at
) VALUES (
{trade_dict['id']},
'{trade_dict['position_id']}',
'{trade_dict['ticker']}',
'{trade_dict['entry_date'].strftime('%Y-%m-%d %H:%M:%S')}',
{trade_dict['shares']},
{trade_dict['entry_price']},
{trade_dict['target_price'] if trade_dict['target_price'] else 'NULL'},
{trade_dict['stop_loss'] if trade_dict['stop_loss'] else 'NULL'},
{f"'{trade_dict['strategy']}'" if trade_dict['strategy'] else 'NULL'},
'{trade_dict['order_type']}',
'{direction}',
{1 if trade_dict['followed_rules'] else 0},
{f"'{trade_dict['entry_reason']}'" if trade_dict['entry_reason'] else 'NULL'},
{trade_dict['exit_price'] if trade_dict['exit_price'] else 'NULL'},
{f"'{trade_dict['exit_date'].strftime('%Y-%m-%d %H:%M:%S')}'" if trade_dict['exit_date'] else 'NULL'},
{f"'{trade_dict['exit_reason']}'" if trade_dict['exit_reason'] else 'NULL'},
{f"'{trade_dict['notes']}'" if trade_dict['notes'] else 'NULL'},
'{trade_dict['created_at'].strftime('%Y-%m-%d %H:%M:%S')}'
)
"""
client.command(query)
# Rename tables
client.command("RENAME TABLE stock_db.trades TO stock_db.trades_backup")
client.command("RENAME TABLE stock_db.trades_new TO stock_db.trades")
print("Migration completed successfully!")
print("Old table backed up as trades_backup")
if __name__ == "__main__":
migrate_trades()