Initial Commit

28 files changed, 262 insertions, 0 deletions

diff --git a/backend/db/embed.go b/backend/db/embed.go
new file mode 100644
index 0000000..4302300
--- /dev/null
+++ b/backend/db/embed.go
@@ -0,0 +1,17 @@
+package db
+
+import "embed"
+
+//go:embed migrations/*.sql
+var Migrations embed.FS
+
+//go:embed all:queries
+var queries embed.FS
+
+func GetQuery(name string) (string, error) {
+	query, err := queries.ReadFile("queries/" + name + ".sql")
+	if err != nil {
+		return "", err
+	}
+	return string(query), nil
+}
diff --git a/backend/db/migrations/000001_create_user_tables.down.sql b/backend/db/migrations/000001_create_user_tables.down.sql
new file mode 100644
index 0000000..f1ba1c0
--- /dev/null
+++ b/backend/db/migrations/000001_create_user_tables.down.sql
@@ -0,0 +1,5 @@
+DROP INDEX IF EXISTS idx_ibd_tokens_user_subject;
+
+DROP TABLE IF EXISTS ibd_tokens;
+DROP TABLE IF EXISTS users;
+DROP TABLE IF EXISTS keys;
diff --git a/backend/db/migrations/000001_create_user_tables.up.sql b/backend/db/migrations/000001_create_user_tables.up.sql
new file mode 100644
index 0000000..6681aa6
--- /dev/null
+++ b/backend/db/migrations/000001_create_user_tables.up.sql
@@ -0,0 +1,47 @@
+/**
+  * Keys Table
+  * The keys table stores the encryption keys used to encrypt user data.
+  * It holds encrypted AES keys and the KMS key used to encrypt them.
+ */
+CREATE TABLE IF NOT EXISTS keys
+(
+    id            SERIAL PRIMARY KEY,                             -- Unique ID for the key
+    kms_key_name  VARCHAR(255) NOT NULL,                          -- The name of the KMS key
+    encrypted_key BYTEA        NOT NULL,                          -- The encrypted AES key (encrypted with the KMS key)
+    created_at    TIMESTAMP    NOT NULL DEFAULT CURRENT_TIMESTAMP -- The datetime the key was created
+);
+
+/**
+ * Users Table
+ * The users table stores the user data for the app.
+ */
+CREATE TABLE IF NOT EXISTS users
+(
+    subject        TEXT PRIMARY KEY,            -- The unique subject of the user, retrieved from Auth0
+    ibd_username   VARCHAR(255),                -- The IBD username of the user
+    ibd_password   BYTEA,                       -- The encrypted IBD password. Encrypted with `encryption_key`
+    encryption_key INTEGER REFERENCES keys (id) -- The encryption key used to encrypt the IBD password
+);
+
+/**
+ * IBD Tokens Table
+ * The IBD tokens table stores the tokens/cookies used to authenticate with the IBD website.
+ * These are scraped using chromedb and stored in the database for future use via APIs.
+ */
+CREATE TABLE IF NOT EXISTS ibd_tokens
+(
+    id             SERIAL PRIMARY KEY,                            -- Unique ID for the token
+    user_subject   TEXT      NOT NULL REFERENCES users (subject), -- The user in the users table associated with the token
+    token          BYTEA     NOT NULL,                            -- The encrypted token/cookie
+    encryption_key INTEGER   NOT NULL REFERENCES keys (id),       -- The encryption key used to encrypt the token
+    created_at     TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP,  -- The datetime the token was created
+    expires_at     TIMESTAMP NOT NULL,                            -- The datetime the token expires
+    /**
+     * The `degraded` column is used to indicate that the token MAY be invalid.
+     * It is set to TRUE when a worker reports a failure using the token.
+     * This will cause the authentication cronjob to check the token and re-authenticate if necessary.
+     */
+    degraded       BOOLEAN   NOT NULL DEFAULT FALSE
+);
+
+CREATE INDEX IF NOT EXISTS idx_ibd_tokens_user_subject ON ibd_tokens (user_subject);
diff --git a/backend/db/migrations/000002_create_data_tables.down.sql b/backend/db/migrations/000002_create_data_tables.down.sql
new file mode 100644
index 0000000..640fe94
--- /dev/null
+++ b/backend/db/migrations/000002_create_data_tables.down.sql
@@ -0,0 +1,11 @@
+DROP INDEX IF EXISTS ratings_symbol_idx;
+
+DROP TABLE IF EXISTS stock_data;
+DROP TABLE IF EXISTS ratings;
+DROP TABLE IF EXISTS chart_analysis;
+DROP TABLE IF EXISTS stock_rank;
+DROP TABLE IF EXISTS stocks;
+
+DROP TYPE IF EXISTS CHART_ACTION;
+DROP TYPE IF EXISTS BAR_INTERVAL;
+DROP TYPE IF EXISTS RANK_TYPE;
\ No newline at end of file
diff --git a/backend/db/migrations/000002_create_data_tables.up.sql b/backend/db/migrations/000002_create_data_tables.up.sql
new file mode 100644
index 0000000..116c424
--- /dev/null
+++ b/backend/db/migrations/000002_create_data_tables.up.sql
@@ -0,0 +1,89 @@
+/**
+ * Stocks table
+ * The stocks table stores the basic information about each stock.
+ * This data is scraped periodically from the IBD website.
+ */
+CREATE TABLE IF NOT EXISTS stocks
+(
+    symbol  VARCHAR(16) PRIMARY KEY, -- The stock symbol/ticker
+    name    TEXT NOT NULL,           -- The full name of the stock
+    ibd_url TEXT                     -- The URL to the IBD page for the stock
+);
+
+/**
+ * Stock Rank table
+ * The stock rank table stores the rank information for each stock.
+ * This data is scraped periodically from the IBD website.
+ */
+CREATE TYPE RANK_TYPE AS ENUM ('ibd50', 'cap20');
+CREATE TABLE IF NOT EXISTS stock_rank
+(
+    symbol     VARCHAR(16) NOT NULL REFERENCES stocks (symbol), -- The stock symbol from the stocks table
+    rank_type  RANK_TYPE   NOT NULL,                            -- The type of rank (IBD 50 or CAP 20)
+    rank       SMALLINT,                                        -- The rank of the stock
+    created_at TIMESTAMP   NOT NULL DEFAULT CURRENT_TIMESTAMP   -- The datetime the rank was scraped/created
+);
+
+/**
+ * Chart Analysis table
+ * The chart analysis table stores the chart analysis for each stock.
+ * This data is scraped periodically from the IBD website then analyzed by OpenAI's GPT-4o.
+ */
+CREATE TYPE CHART_ACTION AS ENUM ('buy', 'sell', 'hold', 'unknown');
+CREATE TABLE IF NOT EXISTS chart_analysis
+(
+    id           UUID PRIMARY KEY DEFAULT gen_random_uuid(), -- Randomly generated UUID for the chart analysis
+    raw_analysis TEXT    NOT NULL,                           -- The raw chart analysis text (e.g. "Flat base, 3 weeks tight")
+
+    processed    BOOLEAN NOT NULL DEFAULT FALSE,             -- Whether the chart analysis has been processed by GPT-4o
+    action       CHART_ACTION,                               -- The action to take based on the chart analysis
+    price        MONEY,                                      -- The price to take the action at
+    reason       TEXT,                                       -- The reason for the action
+    confidence   SMALLINT                                    -- The confidence level of the action (0-100)
+);
+
+/**
+ * Ratings table
+ * The ratings table stores the ratings for each stock.
+ * This data is scraped periodically from the IBD website.
+ */
+CREATE TABLE IF NOT EXISTS ratings
+(
+    id             UUID PRIMARY KEY     DEFAULT gen_random_uuid(),  -- Randomly generated UUID for the rating
+    symbol         VARCHAR(16) NOT NULL REFERENCES stocks (symbol), -- The stock symbol from the stocks table
+
+    composite      SMALLINT,                                        -- The composite rating (0-99)
+    eps            SMALLINT,                                        -- The EPS rating (0-99)
+    rel_str        SMALLINT,                                        -- The relative strength rating (0-99)
+    group_rel_str  SMALLINT,                                        -- The group relative strength rating (E-A+) E=0, A+=13
+    smr            SMALLINT,                                        -- The SMR rating (E-A+)
+    acc_dis        SMALLINT,                                        -- The acc/dis rating (E-A+)
+    chart_analysis UUID REFERENCES chart_analysis (id),             -- The ID of the chart analysis for the stock
+    price          MONEY,                                           -- The price of the stock at the time of the rating
+
+    created_at     TIMESTAMP   NOT NULL DEFAULT CURRENT_TIMESTAMP   -- The datetime the rating was scraped/created
+);
+
+CREATE INDEX IF NOT EXISTS ratings_symbol_idx ON ratings (symbol);
+
+/**
+ * Stock Data table
+ * The stock data table stores the historical price data for each stock.
+ * This data should be retrieved from a brokerage API or other data source.
+ * These should be in the format of minute and daily bars.
+ */
+CREATE TYPE BAR_INTERVAL AS ENUM ('1m', '1d');
+CREATE TABLE IF NOT EXISTS stock_data
+(
+    id           UUID PRIMARY KEY DEFAULT gen_random_uuid(),       -- Randomly generated UUID for the stock data
+    symbol       VARCHAR(16)  NOT NULL REFERENCES stocks (symbol), -- The stock symbol from the stocks table
+
+    bar_interval BAR_INTERVAL NOT NULL,                            -- The interval of the bar data (1m, 1d)
+    time         TIMESTAMP    NOT NULL,                            -- The timestamp of the bar
+
+    open         MONEY        NOT NULL,                            -- The opening price of the bar
+    high         MONEY        NOT NULL,                            -- The high price of the bar
+    low          MONEY        NOT NULL,                            -- The low price of the bar
+    close        MONEY        NOT NULL,                            -- The closing price of the bar
+    volume       BIGINT       NOT NULL                             -- The volume of the bar
+);
diff --git a/backend/db/queries/cookies/add_cookie.sql b/backend/db/queries/cookies/add_cookie.sql
new file mode 100644
index 0000000..1519da4
--- /dev/null
+++ b/backend/db/queries/cookies/add_cookie.sql
@@ -0,0 +1,2 @@
+INSERT INTO ibd_tokens (token, expires_at, user_subject, encryption_key)
+VALUES ($1, $2, $3, $4)
\ No newline at end of file
diff --git a/backend/db/queries/cookies/get_any_cookie.sql b/backend/db/queries/cookies/get_any_cookie.sql
new file mode 100644
index 0000000..4e5c823
--- /dev/null
+++ b/backend/db/queries/cookies/get_any_cookie.sql
@@ -0,0 +1,7 @@
+SELECT ibd_tokens.id, token, encrypted_key, kms_key_name, expires_at
+FROM ibd_tokens
+         INNER JOIN keys ON encryption_key = keys.id
+WHERE expires_at > NOW()
+  AND degraded = FALSE
+ORDER BY random()
+LIMIT 1;
diff --git a/backend/db/queries/cookies/get_cookies.sql b/backend/db/queries/cookies/get_cookies.sql
new file mode 100644
index 0000000..3828ec3
--- /dev/null
+++ b/backend/db/queries/cookies/get_cookies.sql
@@ -0,0 +1,7 @@
+SELECT ibd_tokens.id, token, encrypted_key, kms_key_name, expires_at
+FROM ibd_tokens
+         INNER JOIN keys ON encryption_key = keys.id
+WHERE user_subject = $1
+  AND expires_at > NOW()
+  AND degraded = $2
+ORDER BY expires_at DESC;
\ No newline at end of file
diff --git a/backend/db/queries/cookies/set_cookie_degraded.sql b/backend/db/queries/cookies/set_cookie_degraded.sql
new file mode 100644
index 0000000..4fd8222
--- /dev/null
+++ b/backend/db/queries/cookies/set_cookie_degraded.sql
@@ -0,0 +1,3 @@
+UPDATE ibd_tokens
+SET degraded = $1
+WHERE id = $2;
\ No newline at end of file
diff --git a/backend/db/queries/keys/add_key.sql b/backend/db/queries/keys/add_key.sql
new file mode 100644
index 0000000..bb416c5
--- /dev/null
+++ b/backend/db/queries/keys/add_key.sql
@@ -0,0 +1,3 @@
+INSERT INTO keys (kms_key_name, encrypted_key)
+VALUES ($1, $2)
+RETURNING id;
\ No newline at end of file
diff --git a/backend/db/queries/keys/get_key.sql b/backend/db/queries/keys/get_key.sql
new file mode 100644
index 0000000..97d8367
--- /dev/null
+++ b/backend/db/queries/keys/get_key.sql
@@ -0,0 +1,3 @@
+SELECT id, kms_key_name, encrypted_key, created_at
+FROM keys
+WHERE id = $1;
\ No newline at end of file
diff --git a/backend/db/queries/sessions/check_state.sql b/backend/db/queries/sessions/check_state.sql
new file mode 100644
index 0000000..dac73e2
--- /dev/null
+++ b/backend/db/queries/sessions/check_state.sql
@@ -0,0 +1,3 @@
+SELECT 1
+FROM sessions
+where token = $1;
\ No newline at end of file
diff --git a/backend/db/queries/sessions/cleanup_sessions.sql b/backend/db/queries/sessions/cleanup_sessions.sql
new file mode 100644
index 0000000..5f2d22b
--- /dev/null
+++ b/backend/db/queries/sessions/cleanup_sessions.sql
@@ -0,0 +1,2 @@
+DELETE FROM sessions
+WHERE expires_at < NOW();
\ No newline at end of file
diff --git a/backend/db/queries/sessions/create_session.sql b/backend/db/queries/sessions/create_session.sql
new file mode 100644
index 0000000..44f8c56
--- /dev/null
+++ b/backend/db/queries/sessions/create_session.sql
@@ -0,0 +1,2 @@
+INSERT INTO sessions (token, user_subject, access_token, expires_at)
+VALUES ($1, $2, $3, $4);
\ No newline at end of file
diff --git a/backend/db/queries/sessions/create_state.sql b/backend/db/queries/sessions/create_state.sql
new file mode 100644
index 0000000..577ad7e
--- /dev/null
+++ b/backend/db/queries/sessions/create_state.sql
@@ -0,0 +1,2 @@
+INSERT INTO sessions (token, expires_at)
+VALUES ($1, CURRENT_TIMESTAMP + INTERVAL '1 hour');
\ No newline at end of file
diff --git a/backend/db/queries/sessions/get_session.sql b/backend/db/queries/sessions/get_session.sql
new file mode 100644
index 0000000..7da8bd0
--- /dev/null
+++ b/backend/db/queries/sessions/get_session.sql
@@ -0,0 +1,3 @@
+SELECT token, user_subject, access_token, expires_at
+FROM sessions
+WHERE token = $1;
\ No newline at end of file
diff --git a/backend/db/queries/stocks/add_analysis.sql b/backend/db/queries/stocks/add_analysis.sql
new file mode 100644
index 0000000..4bb4903
--- /dev/null
+++ b/backend/db/queries/stocks/add_analysis.sql
@@ -0,0 +1,9 @@
+UPDATE chart_analysis ca
+SET processed  = true,
+    action     = $2,
+    price      = $3,
+    reason     = $4,
+    confidence = $5
+FROM ratings r
+WHERE r.id = $1
+  AND r.chart_analysis = ca.id
\ No newline at end of file
diff --git a/backend/db/queries/stocks/add_rank.sql b/backend/db/queries/stocks/add_rank.sql
new file mode 100644
index 0000000..07f711e
--- /dev/null
+++ b/backend/db/queries/stocks/add_rank.sql
@@ -0,0 +1,2 @@
+INSERT INTO stock_rank (symbol, rank_type, rank)
+VALUES ($1, $2, $3);
\ No newline at end of file
diff --git a/backend/db/queries/stocks/add_rating.sql b/backend/db/queries/stocks/add_rating.sql
new file mode 100644
index 0000000..6c4baa0
--- /dev/null
+++ b/backend/db/queries/stocks/add_rating.sql
@@ -0,0 +1,3 @@
+INSERT INTO ratings (symbol, composite, eps, rel_str, group_rel_str, smr, acc_dis, chart_analysis, price)
+VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9)
+RETURNING id;
\ No newline at end of file
diff --git a/backend/db/queries/stocks/add_raw_chart_analysis.sql b/backend/db/queries/stocks/add_raw_chart_analysis.sql
new file mode 100644
index 0000000..a599d87
--- /dev/null
+++ b/backend/db/queries/stocks/add_raw_chart_analysis.sql
@@ -0,0 +1,3 @@
+INSERT INTO chart_analysis (raw_analysis)
+VALUES ($1)
+RETURNING id;
\ No newline at end of file
diff --git a/backend/db/queries/stocks/add_stock.sql b/backend/db/queries/stocks/add_stock.sql
new file mode 100644
index 0000000..180a9c3
--- /dev/null
+++ b/backend/db/queries/stocks/add_stock.sql
@@ -0,0 +1,5 @@
+INSERT INTO stocks (symbol, name, ibd_url)
+VALUES ($1, $2, $3)
+ON CONFLICT (symbol)
+    DO UPDATE SET name    = $2,
+                  ibd_url = $3;
\ No newline at end of file
diff --git a/backend/db/queries/stocks/get_stock.sql b/backend/db/queries/stocks/get_stock.sql
new file mode 100644
index 0000000..cecbd84
--- /dev/null
+++ b/backend/db/queries/stocks/get_stock.sql
@@ -0,0 +1,3 @@
+SELECT symbol, name, ibd_url
+FROM stocks
+WHERE symbol = $1;
\ No newline at end of file
diff --git a/backend/db/queries/stocks/get_stock_info.sql b/backend/db/queries/stocks/get_stock_info.sql
new file mode 100644
index 0000000..d4f1bf3
--- /dev/null
+++ b/backend/db/queries/stocks/get_stock_info.sql
@@ -0,0 +1,14 @@
+SELECT r.symbol,
+       s.name,
+       ca.raw_analysis,
+       r.composite,
+       r.eps,
+       r.rel_str,
+       r.group_rel_str,
+       r.smr,
+       r.acc_dis,
+       r.price
+FROM ratings r
+         INNER JOIN stocks s on r.symbol = s.symbol
+         INNER JOIN chart_analysis ca on r.chart_analysis = ca.id
+WHERE r.id = $1;
\ No newline at end of file
diff --git a/backend/db/queries/users/add_ibd_creds.sql b/backend/db/queries/users/add_ibd_creds.sql
new file mode 100644
index 0000000..054f328
--- /dev/null
+++ b/backend/db/queries/users/add_ibd_creds.sql
@@ -0,0 +1,5 @@
+UPDATE users
+SET ibd_username   = $2,
+    ibd_password   = $3,
+    encryption_key = $4
+WHERE subject = $1;
\ No newline at end of file
diff --git a/backend/db/queries/users/add_user.sql b/backend/db/queries/users/add_user.sql
new file mode 100644
index 0000000..bf97ad5
--- /dev/null
+++ b/backend/db/queries/users/add_user.sql
@@ -0,0 +1,3 @@
+INSERT INTO users (subject)
+VALUES ($1)
+ON CONFLICT DO NOTHING;
\ No newline at end of file
diff --git a/backend/db/queries/users/get_ibd_creds.sql b/backend/db/queries/users/get_ibd_creds.sql
new file mode 100644
index 0000000..271abcc
--- /dev/null
+++ b/backend/db/queries/users/get_ibd_creds.sql
@@ -0,0 +1,4 @@
+SELECT ibd_username, ibd_password, encrypted_key, kms_key_name
+FROM users
+INNER JOIN public.keys k on k.id = users.encryption_key
+WHERE subject = $1;
\ No newline at end of file
diff --git a/backend/db/queries/users/get_user.sql b/backend/db/queries/users/get_user.sql
new file mode 100644
index 0000000..567f988
--- /dev/null
+++ b/backend/db/queries/users/get_user.sql
@@ -0,0 +1,3 @@
+SELECT subject, ibd_username, ibd_password, encryption_key
+FROM users
+WHERE subject = $1;
\ No newline at end of file
diff --git a/backend/db/queries/users/list_users.sql b/backend/db/queries/users/list_users.sql
new file mode 100644
index 0000000..ceafeb2
--- /dev/null
+++ b/backend/db/queries/users/list_users.sql
@@ -0,0 +1,2 @@
+SELECT subject, ibd_username, ibd_password, encryption_key
+FROM users;
\ No newline at end of file