Switching to using Postgres DB

Author: Breta01

.env.example

@@ -3,4 +3,7 @@ BRAIN_CREDENTIAL_PASSWORD=
 
 # Only google api key is used at the moment
 GOOGLE_API_KEY=
-# OPENAI_API_KEY=
+OPENAI_API_KEY=
+
+# Database URL
+DATABASE_URL=

.gitignore

@@ -135,3 +135,4 @@ brain/tools/data/delay_0.csv
 brain/tools/data/delay_1.csv
 brain/tools/data/TOP*_delay_*.csv
 alphas_database.db
+migrations/schema.sql

README.md

@@ -37,3 +37,19 @@ For development of code and commiting into git repository, run following command
 ```bash
 pre-commit install
 ```
+
+## Databse Migrations
+
+1. Set `DATABASE_URL` in `.env` file
+2. Create new migration:
+
+```bash
+./db_migrate.sh new my_migration_name
+```
+
+2. Update migrations file
+3. Execute migrations
+
+```bash
+./db_migrate.sh up
+```

brain/agent.py

@@ -38,9 +38,9 @@ def prompt(
 6. Always try uncommon data fields first before using close, volumn, open, close, etc.
 7. Propose the alpha and run the alpha simulation.
 
-- Experiment with different operators
 - Try applying different functions and changing parameter values.
 - For d parameter use one of the following values: 2, 5, 10, 20, 30, 50, 100, 120, 250
+- Create simple alphas that are easy to understand and explain
 - Start by combining two data fields and add more if alpha is good.
 - Apply vec_avg(x) or vec_sum(x) if data field is a type VECTOR.
 - Change alpha completely if you think it is not good enough.

brain/database.py

@@ -1,73 +1,45 @@
 """Database module for storing and managing alphas."""
 
-import sqlite3
-from pathlib import Path
+import os
 
 import numpy as np
+import psycopg
+from dotenv import load_dotenv
+from psycopg.rows import dict_row
 
 from brain.alpha_class import Alpha
 
-DB_PATH = Path(__file__).parent.parent / "alphas_database.db"
+load_dotenv()
 
-
-sqlite3.register_adapter(np.int64, lambda x: int(x))
+psycopg.adapters.register_dumper(np.int64, psycopg.types.numeric.IntDumper)
 
 
 class Database:
-    def __init__(self, db_path=DB_PATH):
-        self.conn = sqlite3.connect(db_path)
-        self.conn.row_factory = sqlite3.Row
+    def __init__(self, db_url: str = None):
+        """Initialize the database connection."""
+        self.db_url = db_url or os.environ.get("DATABASE_URL")
+        self.conn = psycopg.connect(self.db_url, row_factory=dict_row)
+        self.conn.set_autocommit(True)
         self.cursor = self.conn.cursor()
-        self._create_table()
-
-    def _create_table(self):
-        schema = """
-        CREATE TABLE IF NOT EXISTS alphas (
-            alpha_id TEXT PRIMARY KEY,
-            created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
-            regular TEXT NOT NULL,
-            region TEXT NOT NULL,
-            universe TEXT NOT NULL,
-            decay INTEGER NOT NULL,
-            delay INTEGER NOT NULL,
-            truncation REAL NOT NULL,
-            neutralization TEXT NOT NULL,
-            pasteurization TEXT NOT NULL,
-            nan_handling TEXT NOT NULL,
-            unit_handling TEXT NOT NULL,
-            fitness REAL,
-            sharpe REAL,
-            returns REAL,
-            drawdown REAL,
-            turnover REAL,
-            margin REAL,
-            long_count INTEGER,
-            short_count INTEGER,
-            self_correlation REAL,
-            failing_tests TEXT
-        );
-        """
-        self.cursor.execute(schema)
-        self.conn.commit()
 
     def insert_alpha(self, alpha: Alpha) -> int:
         """Insert a Alpha class instance into the table."""
         record = alpha.as_dict()
         cols = ", ".join(record)
-        bangs = ", ".join("?" for _ in record)
-
+        bangs = ", ".join("%s" for _ in record)
         sql = f"INSERT INTO alphas ({cols}) VALUES ({bangs})"
         self.cursor.execute(sql, tuple(record.values()))
-        self.conn.commit()
 
     def find_by_code(self, code: str, neutralization: str, delay: int) -> list[Alpha]:
         """Find an alpha by its code."""
         sql = (
-            "SELECT * FROM alphas WHERE regular = ? AND neutralization = ? AND delay = ? "
+            "SELECT * FROM alphas WHERE regular = %s AND neutralization = %s AND delay = %s "
             "ORDER BY created_at DESC"
         )
-        rows = self.cursor.execute(sql, (code, neutralization, delay))
+        self.cursor.execute(sql, (code, neutralization, delay))
+        rows = self.cursor.fetchall()
         return [Alpha.from_row(r) for r in rows]
 
     def close(self):
+        self.cursor.close()
         self.conn.close()

db_migrate.sh

@@ -0,0 +1,140 @@
+#!/bin/bash
+# Runs migrations on PostgreSQL
+# Based on:
+# https://github.com/Schniz/migrate.sh/edit/master/migrate.sh
+
+# Load environment variables from .env file
+set -a && source .env && set +a
+
+set -e
+
+function make_new_migration() {
+  FILE_DESC=$(echo $@ | sed -E "s/[[:space:]]+/_/g" | awk '{print tolower($0)}')
+  FILE_PATH=migrations/$(get_timestamp)_${FILE_DESC}.sql
+  mkdir -p migrations
+  touch $FILE_PATH
+  echo "  --> touched $FILE_PATH"
+}
+
+function verify_database_url() {
+  if [ "$DATABASE_URL" == "" ]; then
+    echo "  --> ERROR: variable DATABASE_URL is not defined"
+    exit 1
+  fi
+}
+
+function pending_migrations() {
+  UP_MIGRATIONS=$(bash -c 'psql $DATABASE_URL -t -c "select filename from migrations" | sed "s@[[:space:]]@@g" | grep . | cat' 2> /dev/null)
+  ALL_MIGRATIONS=$(bash -c 'cd migrations && ls *.sql')
+  STRINGIFIED_MIGRATIONS=" ${ALL_MIGRATIONS[*]} "
+  for item in ${UP_MIGRATIONS[@]}; do
+    STRINGIFIED_MIGRATIONS=${STRINGIFIED_MIGRATIONS/${item}/ }
+  done
+  MIGRATIONS_TO_RUN=( $STRINGIFIED_MIGRATIONS )
+  for item in ${MIGRATIONS_TO_RUN[@]}; do
+    echo $item
+  done
+}
+
+function run_migrations() {
+  PENDING_MIGRATIONS=$(pending_migrations)
+  if [ "$PENDING_MIGRATIONS" == "" ]; then
+    echo "  --> Nothing to migrate!"
+  else
+    echo "PENDING MIGRATIONS:"
+    echo ${PENDING_MIGRATIONS[*]}
+    echo "==================="
+    for PENDING_MIGRATION in $PENDING_MIGRATIONS; do
+      CONTENTS=$(cat migrations/$PENDING_MIGRATION)
+      CONTENTS_WITH_MIGRATION_RESULT="
+        begin;
+        $CONTENTS;
+        insert into migrations (filename) values ('$(basename $PENDING_MIGRATION)');
+        commit;
+      "
+      echo "  --> Running $PENDING_MIGRATION"
+      cat <<< "$CONTENTS_WITH_MIGRATION_RESULT"
+      psql $DATABASE_URL -v ON_ERROR_STOP=1  <<< $CONTENTS_WITH_MIGRATION_RESULT
+    done
+  fi
+#   dump_schema
+}
+
+function dump_schema() {
+  mkdir -p migrations
+
+  pg_dump \
+    --schema-only \
+    $DATABASE_URL \
+    > migrations/schema.sql.tmp
+  pg_dump \
+    --data-only \
+    -t migrations \
+    $DATABASE_URL \
+    >> migrations/schema.sql.tmp
+
+  cat migrations/schema.sql.tmp | grep -v -e "---\?\( .\+\)\$" | grep -v -e "^--\$" | grep -v -e "^\$" > migrations/schema.sql
+  rm migrations/schema.sql.tmp
+
+  echo "  --> Wrote migrations/schema.sql"
+}
+
+function reset_db() {
+  psql $DATABASE_URL <<< "
+    DROP SCHEMA public CASCADE;
+    CREATE SCHEMA public;
+    GRANT ALL ON SCHEMA public TO postgres;
+    GRANT ALL ON SCHEMA public TO public;
+    CREATE TABLE IF NOT EXISTS migrations (
+        filename TEXT PRIMARY KEY,
+        ran_at   TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP
+    );
+  "
+}
+
+function show_help() {
+  echo "Usage:"
+  echo "------"
+  echo ""
+  echo "$0 new add_something_table - create a new migration file"
+  echo "$0 schema:dump             - dump db schema"
+  echo "$0 schema:load             - load db schema"
+  echo "$0 up                      - run migrations"
+  echo "$0 danger:reset            - resets the database state"
+  echo "$0 help                    - show this message"
+}
+
+function get_timestamp() {
+  date +%Y%m%d%H%M%S
+}
+
+function main_migrate() {
+  ACTION=${1-:"help"}
+  shift
+  case $ACTION in
+    new)
+      make_new_migration $@
+      ;;
+    schema:dump)
+      verify_database_url
+      dump_schema $@
+      ;;
+    schema:load)
+      verify_database_url
+      psql $DATABASE_URL -f migrations/schema.sql
+      ;;
+    danger:reset)
+      verify_database_url
+      reset_db
+      ;;
+    up)
+      verify_database_url
+      run_migrations $@
+      ;;
+    *)
+      show_help
+      exit 1
+  esac
+}
+
+main_migrate $@

migrations/20250527124436_init.sql

@@ -0,0 +1,24 @@
+CREATE TABLE IF NOT EXISTS alphas (
+    alpha_id TEXT PRIMARY KEY,
+    created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
+    regular TEXT NOT NULL,
+    region TEXT NOT NULL,
+    universe TEXT NOT NULL,
+    decay INTEGER NOT NULL,
+    delay INTEGER NOT NULL,
+    truncation REAL NOT NULL,
+    neutralization TEXT NOT NULL,
+    pasteurization TEXT NOT NULL,
+    nan_handling TEXT NOT NULL,
+    unit_handling TEXT NOT NULL,
+    fitness REAL,
+    sharpe REAL,
+    returns REAL,
+    drawdown REAL,
+    turnover REAL,
+    margin REAL,
+    long_count INTEGER,
+    short_count INTEGER,
+    self_correlation REAL,
+    failing_tests TEXT
+);

requirements.txt

@@ -2,6 +2,7 @@ langchain[google-genai,openai]
 langgraph
 langmem
 pandas
+psycopg[binary]
 python-dotenv
 requests
 tqdm