diff --git a/src/migrations/add_direction_field.py b/src/migrations/add_direction_field.py new file mode 100644 index 0000000..cf4668c --- /dev/null +++ b/src/migrations/add_direction_field.py @@ -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()