From ed582ae31da503c03d968513775bd32255ebbb3c Mon Sep 17 00:00:00 2001
From: Eric Bower <me@erock.io>
Date: Thu, 29 Feb 2024 14:37:20 -0500
Subject: [PATCH] chore: pico+ go script

---
 Makefile                      |  4 +++
 cmd/scripts/pico-plus/main.go | 32 ++++++++++++++++++
 db/db.go                      | 22 +++++++++++++
 db/postgres/storage.go        | 62 +++++++++++++++++++++++++++++++++++
 4 files changed, 120 insertions(+)
 create mode 100644 cmd/scripts/pico-plus/main.go

diff --git a/Makefile b/Makefile
index 99e328d2..3bc7f1dd 100644
--- a/Makefile
+++ b/Makefile
@@ -73,6 +73,10 @@ store-clean:
 	WRITE=$(WRITE) go run ./cmd/scripts/clean-object-store/clean.go
 .PHONY: store-clean
 
+pico-plus:
+	go run ./cmd/scripts/pico-plus/main.go $(USER) $(TXID)
+.PHONY: pico-plus
+
 scripts:
 	# might need to set MINIO_URL
 	docker run --rm -it --env-file .env -v $(shell pwd):/app -w /app golang:1.21 /bin/bash
diff --git a/cmd/scripts/pico-plus/main.go b/cmd/scripts/pico-plus/main.go
new file mode 100644
index 00000000..0701101c
--- /dev/null
+++ b/cmd/scripts/pico-plus/main.go
@@ -0,0 +1,32 @@
+package main
+
+import (
+	"log/slog"
+	"os"
+
+	"github.com/picosh/pico/db/postgres"
+)
+
+func main() {
+	logger := slog.Default()
+	DbURL := os.Getenv("DATABASE_URL")
+	dbpool := postgres.NewDB(DbURL, logger)
+
+	args := os.Args
+	username := args[1]
+	txId := args[2]
+
+	logger.Info(
+		"Upgrading user to pico+",
+		"username", username,
+		"txId", txId,
+	)
+
+	err := dbpool.AddPicoPlusUser(username, txId)
+	if err != nil {
+		logger.Error("Failed to add pico+ user", "err", err)
+		os.Exit(1)
+	} else {
+		logger.Info("Successfully added pico+ user")
+	}
+}
diff --git a/db/db.go b/db/db.go
index 324bfbc1..7e31bf1a 100644
--- a/db/db.go
+++ b/db/db.go
@@ -235,6 +235,27 @@ func (p *FeatureFlagData) Scan(value interface{}) error {
 	return json.Unmarshal(b, &p)
 }
 
+type PaymentHistoryData struct {
+	Notes string `json:"notes"`
+	TxID  string `json:"tx_id"`
+}
+
+// Make the Attrs struct implement the driver.Valuer interface. This method
+// simply returns the JSON-encoded representation of the struct.
+func (p PaymentHistoryData) Value() (driver.Value, error) {
+	return json.Marshal(p)
+}
+
+// Make the Attrs struct implement the sql.Scanner interface. This method
+// simply decodes a JSON-encoded value into the struct fields.
+func (p *PaymentHistoryData) Scan(value interface{}) error {
+	b, ok := value.([]byte)
+	if !ok {
+		return errors.New("type assertion to []byte failed")
+	}
+	return json.Unmarshal(b, &p)
+}
+
 type ErrMultiplePublicKeys struct{}
 
 func (m *ErrMultiplePublicKeys) Error() string {
@@ -303,6 +324,7 @@ type DB interface {
 
 	AddViewCount(postID string) (int, error)
 
+	AddPicoPlusUser(username string, txId string) error
 	FindFeatureForUser(userID string, feature string) (*FeatureFlag, error)
 	HasFeatureForUser(userID string, feature string) bool
 	FindTotalSizeForUser(userID string) (int, error)
diff --git a/db/postgres/storage.go b/db/postgres/storage.go
index 1928ebf5..41d8e180 100644
--- a/db/postgres/storage.go
+++ b/db/postgres/storage.go
@@ -1540,3 +1540,65 @@ func (me *PsqlDB) FindTokensForUser(userID string) ([]*db.Token, error) {
 	}
 	return keys, nil
 }
+
+func (me *PsqlDB) AddPicoPlusUser(username, txId string) error {
+	user, err := me.FindUserForName(username)
+	if err != nil {
+		return err
+	}
+
+	ctx := context.Background()
+	tx, err := me.Db.BeginTx(ctx, nil)
+	if err != nil {
+		return err
+	}
+	defer func() {
+		err = tx.Rollback()
+	}()
+
+	if txId != "" {
+		data := db.PaymentHistoryData{
+			Notes: "",
+			TxID:  txId,
+		}
+		_, err := tx.Exec(
+			`INSERT INTO payment_history (user_id, payment_type, amount, data) VALUES ($1, 'stripe', 20 * 1000000, $2);`,
+			user.ID,
+			data,
+		)
+		if err != nil {
+			fmt.Println(err)
+			return err
+		}
+	}
+
+	pgsQuery := `INSERT INTO feature_flags (user_id, name, data, expires_at)
+	VALUES ($1, 'pgs', '{"storage_max":10000000000, "file_max":50000000}'::jsonb, now() + '1 year'::interval);`
+	_, err = tx.Exec(pgsQuery, user.ID)
+	if err != nil {
+		return err
+	}
+
+	imgsQuery := `INSERT INTO feature_flags (user_id, name, data, expires_at)
+	VALUES ($1, 'imgs', '{"storage_max":2000000000}'::jsonb, now() + '1 year'::interval);`
+	_, err = tx.Exec(imgsQuery, user.ID)
+	if err != nil {
+		return err
+	}
+
+	proseQuery := `INSERT INTO feature_flags (user_id, name, data, expires_at)
+	VALUES ($1, 'prose', '{"storage_max":1000000000, "file_max":50000000}'::jsonb, now() + '1 year'::interval);`
+	_, err = tx.Exec(proseQuery, user.ID)
+	if err != nil {
+		return err
+	}
+
+	tunsQuery := `INSERT INTO feature_flags (user_id, name, expires_at)
+	VALUES ($1, 'tuns', now() + '1 year'::interval);`
+	_, err = tx.Exec(tunsQuery, user.ID)
+	if err != nil {
+		return err
+	}
+
+	return tx.Commit()
+}