1 files changed, 293 insertions, 0 deletions
diff --git a/backend/internal/database/stocks.go b/backend/internal/database/stocks.go
new file mode 100644
index 0000000..24f5fe7
--- /dev/null
+++ b/backend/internal/database/stocks.go
@@ -0,0 +1,293 @@
+package database
+
+import (
+	"context"
+	"database/sql"
+	"database/sql/driver"
+	"errors"
+
+	pb "github.com/ansg191/ibd-trader-backend/api/gen/idb/stock/v1"
+	"github.com/ansg191/ibd-trader-backend/internal/analyzer"
+	"github.com/ansg191/ibd-trader-backend/internal/utils"
+
+	"github.com/Rhymond/go-money"
+)
+
+var ErrStockNotFound = errors.New("stock not found")
+
+func GetStock(ctx context.Context, exec Executor, symbol string) (Stock, error) {
+	row := exec.QueryRowContext(ctx, `
+SELECT symbol, name, ibd_url
+FROM stocks
+WHERE symbol = $1;
+`, symbol)
+
+	var stock Stock
+	if err := row.Scan(&stock.Symbol, &stock.Name, &stock.IBDUrl); err != nil {
+		if errors.Is(err, sql.ErrNoRows) {
+			return Stock{}, ErrStockNotFound
+		}
+		return Stock{}, err
+	}
+
+	return stock, nil
+}
+
+func AddStock(ctx context.Context, exec Executor, stock Stock) error {
+	_, err := exec.ExecContext(ctx, `
+INSERT INTO stocks (symbol, name, ibd_url)
+VALUES ($1, $2, $3)
+ON CONFLICT (symbol)
+    DO UPDATE SET name    = $2,
+                  ibd_url = $3;`, stock.Symbol, stock.Name, stock.IBDUrl)
+	return err
+}
+
+func AddRanking(ctx context.Context, exec Executor, symbol string, ibd50, cap20 int) error {
+	if ibd50 > 0 {
+		_, err := exec.ExecContext(ctx, `
+INSERT INTO stock_rank (symbol, rank_type, rank)
+VALUES ($1, $2, $3)`, symbol, "ibd50", ibd50)
+		if err != nil {
+			return err
+		}
+	}
+	if cap20 > 0 {
+		_, err := exec.ExecContext(ctx, `
+INSERT INTO stock_rank (symbol, rank_type, rank)
+VALUES ($1, $2, $3)`, symbol, "cap20", cap20)
+		if err != nil {
+			return err
+		}
+	}
+	return nil
+}
+
+func AddStockInfo(ctx context.Context, exec TransactionExecutor, info *StockInfo) (string, error) {
+	tx, err := exec.BeginTx(ctx, nil)
+	if err != nil {
+		return "", err
+	}
+	defer func(tx *sql.Tx) {
+		_ = tx.Rollback()
+	}(tx)
+
+	// Add raw chart analysis
+	row := tx.QueryRowContext(ctx, `
+INSERT INTO chart_analysis (raw_analysis)
+VALUES ($1)
+RETURNING id;`, info.ChartAnalysis)
+
+	var chartAnalysisID string
+	if err = row.Scan(&chartAnalysisID); err != nil {
+		return "", err
+	}
+
+	// Add stock info
+	row = tx.QueryRowContext(ctx,
+		`
+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;`,
+		info.Symbol,
+		info.Ratings.Composite,
+		info.Ratings.EPS,
+		info.Ratings.RelStr,
+		info.Ratings.GroupRelStr,
+		info.Ratings.SMR,
+		info.Ratings.AccDis,
+		chartAnalysisID,
+		info.Price.Display(),
+	)
+
+	var ratingsID string
+	if err = row.Scan(&ratingsID); err != nil {
+		return "", err
+	}
+
+	return ratingsID, tx.Commit()
+}
+
+func GetStockInfo(ctx context.Context, exec Executor, id string) (*StockInfo, error) {
+	row := exec.QueryRowContext(ctx, `
+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;`, id)
+
+	var info StockInfo
+	var priceStr string
+	err := row.Scan(
+		&info.Symbol,
+		&info.Name,
+		&info.ChartAnalysis,
+		&info.Ratings.Composite,
+		&info.Ratings.EPS,
+		&info.Ratings.RelStr,
+		&info.Ratings.GroupRelStr,
+		&info.Ratings.SMR,
+		&info.Ratings.AccDis,
+		&priceStr,
+	)
+	if err != nil {
+		return nil, err
+	}
+
+	info.Price, err = utils.ParseMoney(priceStr)
+	if err != nil {
+		return nil, err
+	}
+
+	return &info, nil
+}
+
+func AddAnalysis(
+	ctx context.Context,
+	exec Executor,
+	ratingId string,
+	analysis *analyzer.Analysis,
+) (id string, err error) {
+	err = exec.QueryRowContext(ctx, `
+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
+RETURNING ca.id;`,
+		ratingId,
+		analysis.Action,
+		analysis.Price.Display(),
+		analysis.Reason,
+		analysis.Confidence,
+	).Scan(&id)
+	return id, err
+}
+
+type Stock struct {
+	Symbol string
+	Name   string
+	IBDUrl string
+}
+
+type StockInfo struct {
+	Symbol        string
+	Name          string
+	ChartAnalysis string
+	Ratings       Ratings
+	Price         *money.Money
+}
+
+type Ratings struct {
+	Composite   uint8
+	EPS         uint8
+	RelStr      uint8
+	GroupRelStr LetterRating
+	SMR         LetterRating
+	AccDis      LetterRating
+}
+
+type LetterRating pb.LetterGrade
+
+func (r LetterRating) String() string {
+	switch pb.LetterGrade(r) {
+	case pb.LetterGrade_LETTER_GRADE_E:
+		return "E"
+	case pb.LetterGrade_LETTER_GRADE_E_PLUS:
+		return "E+"
+	case pb.LetterGrade_LETTER_GRADE_D_MINUS:
+		return "D-"
+	case pb.LetterGrade_LETTER_GRADE_D:
+		return "D"
+	case pb.LetterGrade_LETTER_GRADE_D_PLUS:
+		return "D+"
+	case pb.LetterGrade_LETTER_GRADE_C_MINUS:
+		return "C-"
+	case pb.LetterGrade_LETTER_GRADE_C:
+		return "C"
+	case pb.LetterGrade_LETTER_GRADE_C_PLUS:
+		return "C+"
+	case pb.LetterGrade_LETTER_GRADE_B_MINUS:
+		return "B-"
+	case pb.LetterGrade_LETTER_GRADE_B:
+		return "B"
+	case pb.LetterGrade_LETTER_GRADE_B_PLUS:
+		return "B+"
+	case pb.LetterGrade_LETTER_GRADE_A_MINUS:
+		return "A-"
+	case pb.LetterGrade_LETTER_GRADE_A:
+		return "A"
+	case pb.LetterGrade_LETTER_GRADE_A_PLUS:
+		return "A+"
+	default:
+		return "NA"
+	}
+}
+
+func LetterRatingFromString(str string) LetterRating {
+	switch str {
+	case "E":
+		return LetterRating(pb.LetterGrade_LETTER_GRADE_E)
+	case "E+":
+		return LetterRating(pb.LetterGrade_LETTER_GRADE_E_PLUS)
+	case "D-":
+		return LetterRating(pb.LetterGrade_LETTER_GRADE_D_MINUS)
+	case "D":
+		return LetterRating(pb.LetterGrade_LETTER_GRADE_D)
+	case "D+":
+		return LetterRating(pb.LetterGrade_LETTER_GRADE_D_PLUS)
+	case "C-":
+		return LetterRating(pb.LetterGrade_LETTER_GRADE_C_MINUS)
+	case "C":
+		return LetterRating(pb.LetterGrade_LETTER_GRADE_C)
+	case "C+":
+		return LetterRating(pb.LetterGrade_LETTER_GRADE_C_PLUS)
+	case "B-":
+		return LetterRating(pb.LetterGrade_LETTER_GRADE_B_MINUS)
+	case "B":
+		return LetterRating(pb.LetterGrade_LETTER_GRADE_B)
+	case "B+":
+		return LetterRating(pb.LetterGrade_LETTER_GRADE_B_PLUS)
+	case "A-":
+		return LetterRating(pb.LetterGrade_LETTER_GRADE_A_MINUS)
+	case "A":
+		return LetterRating(pb.LetterGrade_LETTER_GRADE_A)
+	case "A+":
+		return LetterRating(pb.LetterGrade_LETTER_GRADE_A_PLUS)
+	case "NA":
+		fallthrough
+	default:
+		return LetterRating(pb.LetterGrade_LETTER_GRADE_UNSPECIFIED)
+	}
+}
+
+func (r LetterRating) Value() (driver.Value, error) {
+	return r.String(), nil
+}
+
+func (r *LetterRating) Scan(src any) error {
+	var source string
+	switch v := src.(type) {
+	case string:
+		source = v
+	case []byte:
+		source = string(v)
+	default:
+		return errors.New("incompatible type for LetterRating")
+	}
+	*r = LetterRatingFromString(source)
+	return nil
+}