5 files changed, 250 insertions, 0 deletions

diff --git a/db/db.go b/db/db.go
new file mode 100644
index 0000000..bc4b8b3
--- /dev/null
+++ b/db/db.go
@@ -0,0 +1,75 @@
+package db
+
+import (
+        "fmt"
+        "strings"
+
+        "github.com/go-pg/pg"
+        "github.com/go-pg/pg/orm"
+        "github.com/jloup/utils"
+)
+
+var DB *pg.DB
+
+var log = utils.StandardL().WithField("module", "db")
+
+type DBConfig struct {
+        Address  string
+        Database string
+        User     string
+        Password string
+}
+
+func Init(config DBConfig) {
+        var err error
+
+        DB = connect(config)
+
+        err = createSchema(DB)
+        if err != nil {
+                log.Errorf("cannot create schemas %v\n", err)
+        }
+
+        err = createIndexes(DB)
+        if err != nil {
+                log.Errorf("cannot create indexes %v\n", err)
+        }
+}
+
+func connect(config DBConfig) *pg.DB {
+        return pg.Connect(&pg.Options{
+                User:     config.User,
+                Password: config.Password,
+                Database: config.Database,
+                Addr:     config.Address,
+        })
+}
+
+func createSchema(db *pg.DB) error {
+        for _, model := range []interface{}{&User{}, &MarketConfig{}} {
+                err := db.CreateTable(model, &orm.CreateTableOptions{IfNotExists: true})
+                if err != nil {
+                        return err
+                }
+        }
+        return nil
+}
+
+func createIndexes(db *pg.DB) error {
+        indexes := []struct {
+                TableName string
+                Name      string
+                Columns   []string
+        }{
+                {"market_configs", "market_name_user_id_idx", []string{"user_id", "market_name"}},
+        }
+
+        for _, index := range indexes {
+                _, err := db.Exec(fmt.Sprintf("CREATE UNIQUE INDEX IF NOT EXISTS %s ON %s (%s)", index.Name, index.TableName, strings.Join(index.Columns, ",")))
+                if err != nil {
+                        return err
+                }
+        }
+
+        return nil
+}
diff --git a/db/db_test.go b/db/db_test.go
new file mode 100644
index 0000000..0481915
--- /dev/null
+++ b/db/db_test.go
@@ -0,0 +1,35 @@
+package db
+
+import "testing"
+
+func TestInit(t *testing.T) {
+        Init(DBConfig{"localhost:5432", "cryptoportfolio", "cryptoportfolio", "cryptoportfolio-dev"})
+}
+
+func TestUpdateUser(t *testing.T) {
+        Init(DBConfig{"localhost:5432", "cryptoportfolio", "cryptoportfolio", "cryptoportfolio-dev"})
+        t.Log(InsertUser(&User{Email: "j@test.com", PasswordHash: "yp"}))
+        err := InsertUser(&User{Email: "t2@test.com", PasswordHash: "yp"})
+
+        t.Log(err, IsDup(err))
+
+        t.Log(GetUserByEmail("testyo"))
+}
+
+func TestMarketConfig(t *testing.T) {
+        Init(DBConfig{"localhost:5432", "cryptoportfolio", "cryptoportfolio", "cryptoportfolio-dev"})
+
+        config := MarketConfig{UserId: 1, MarketName: "poloniex"}
+        config.Config = make(map[string]string)
+
+        config.Config["secret"] = "key"
+
+        t.Log(InsertMarketConfig(&config))
+        t.Log(config)
+
+        t.Log(GetUserMarketConfig(1, "poloniex"))
+
+        config.Config["secret2"] = "key2"
+        t.Log(SetUserMarketConfig(1, "poloniex", config.Config))
+        t.Log(SetUserMarketConfig(1, "bifinance", config.Config))
+}
diff --git a/db/errors.go b/db/errors.go
new file mode 100644
index 0000000..ed5f371
--- /dev/null
+++ b/db/errors.go
@@ -0,0 +1,23 @@
+package db
+
+import (
+        "strings"
+
+        "github.com/go-pg/pg"
+)
+
+func PGCode(err error) string {
+        if _, ok := err.(pg.Error); !ok {
+                return ""
+        }
+
+        return err.(pg.Error).Field('C')
+}
+
+func IsDup(err error) bool {
+        return PGCode(err) == "23505"
+}
+
+func IsSQLError(err error) bool {
+        return strings.HasPrefix(err.Error(), "ERROR #")
+}
diff --git a/db/market_config.go b/db/market_config.go
new file mode 100644
index 0000000..b26c092
--- /dev/null
+++ b/db/market_config.go
@@ -0,0 +1,45 @@
+package db
+
+import "github.com/go-pg/pg"
+
+type MarketConfig struct {
+        Id         int64
+        MarketName string `sql:",notnull"`
+        UserId     int64  `sql:",notnull"`
+        Config     map[string]string
+}
+
+func InsertMarketConfig(config *MarketConfig) error {
+        return DB.Insert(config)
+}
+
+func GetUserMarketConfig(userId int64, market string) (*MarketConfig, error) {
+        var config MarketConfig
+
+        err := DB.Model(&config).Where("user_id = ?", userId).Where("market_name = ?", market).First()
+
+        if err != nil && err != pg.ErrNoRows {
+                return nil, err
+        }
+
+        if err == pg.ErrNoRows {
+                return nil, nil
+        } else {
+                return &config, nil
+        }
+}
+
+func SetUserMarketConfig(userId int64, market string, newConfig map[string]string) (*MarketConfig, error) {
+        config := MarketConfig{
+                UserId:     userId,
+                MarketName: market,
+                Config:     newConfig,
+        }
+
+        _, err := DB.Model(&config).
+                OnConflict("(user_id, market_name) DO UPDATE").
+                Set("config = ?", newConfig).
+                Insert()
+
+        return &config, err
+}
diff --git a/db/user.go b/db/user.go
new file mode 100644
index 0000000..aed0ac1
--- /dev/null
+++ b/db/user.go
@@ -0,0 +1,72 @@
+package db
+
+import (
+        "golang.org/x/crypto/bcrypt"
+)
+
+type UserStatus uint8
+
+const (
+        Confirmed UserStatus = iota + 1
+        AwaitingConfirmation
+)
+
+type User struct {
+        Id           int64
+        Email        string `sql:",unique,notnull"`
+        PasswordHash string `sql:",notnull"`
+        OtpSecret    string
+        IsOtpSetup   bool
+        Status       UserStatus
+}
+
+func HashPassword(password string) (string, error) {
+        b, err := bcrypt.GenerateFromPassword([]byte(password), 10)
+
+        return string(b), err
+}
+
+func ValidatePassword(password string, hash string) error {
+        return bcrypt.CompareHashAndPassword([]byte(hash), []byte(password))
+}
+
+func InsertUser(user *User) error {
+        return DB.Insert(user)
+}
+
+func ConfirmUserByEmail(email string) error {
+        _, err := DB.Model(&User{}).Set("status=?", Confirmed).Where("email=?", email).Returning("*").Update()
+
+        return err
+}
+
+func GetUserById(id int64) (*User, error) {
+        user := User{Id: id}
+
+        err := DB.Select(&user)
+
+        return &user, err
+}
+
+func GetUserByEmail(email string) (*User, error) {
+        var users []User
+
+        err := DB.Model(&users).Where("email = ?", email).Select()
+
+        if err != nil {
+                return nil, err
+        }
+
+        if len(users) == 0 {
+                return nil, nil
+        }
+
+        return &users[0], nil
+}
+
+func SetOtpSecret(user *User, secret string, temporary bool) error {
+        user.OtpSecret = secret
+        user.IsOtpSetup = !temporary
+
+        return DB.Update(user)
+}