cpolis/cmd/model/db.go

package model

import (
	"database/sql"
	"fmt"
	"log"
	"time"

	"github.com/go-sql-driver/mysql"
	"golang.org/x/crypto/bcrypt"
)

type DB struct {
	*sql.DB
}

func OpenDB(dbName string) (*DB, error) {
	var err error
	db := DB{DB: &sql.DB{}}

	cfg := mysql.NewConfig()
	cfg.DBName = dbName
	cfg.User, cfg.Passwd, err = getCredentials()
	if err != nil {
		return nil, fmt.Errorf("error reading user credentials for DB: %v", err)
	}

	db.DB, err = sql.Open("mysql", cfg.FormatDSN())
	if err != nil {
		return nil, fmt.Errorf("error opening DB: %v", err)
	}
	if err = db.Ping(); err != nil {
		return nil, fmt.Errorf("error pinging DB: %v", err)
	}

	return &db, nil
}

func (db *DB) UpdateAttribute(table string, id int64, attribute string, val interface{}) error {
	query := fmt.Sprintf(`
    UPDATE %s
    SET %s = ?
    WHERE id = ?
    `, table, attribute)
	if _, err := db.Exec(query, val, id); err != nil {
		return fmt.Errorf("error updating article in DB: %v", err)
	}
	return nil
}

func (db *DB) AddUser(user *User, pass string) error {
	hashedPass, err := bcrypt.GenerateFromPassword([]byte(pass), bcrypt.DefaultCost)
	if err != nil {
		return fmt.Errorf("error creating password hash: %v", err)
	}

	query := `
    INSERT INTO users
        (username, password, first_name, last_name, role)
    VALUES (?, ?, ?, ?, ?)
    `
	if _, err = db.Exec(query, user.UserName, string(hashedPass), user.FirstName, user.LastName, user.Role); err != nil {
		return fmt.Errorf("error inserting user into DB: %v", err)
	}

	return nil
}

func (db *DB) GetID(userName string) (int64, error) {
	var id int64

	query := `
    SELECT id
    FROM users
    WHERE username = ?
    `
	row := db.QueryRow(query, userName)
	if err := row.Scan(&id); err != nil {
		return 0, fmt.Errorf("user not in DB: %v", err)
	}

	return id, nil
}

func (db *DB) CheckPassword(id int64, pass string) error {
	var queriedPass string

	query := `
    SELECT password
    FROM users
    WHERE id = ?
    `
	row := db.QueryRow(query, id)
	if err := row.Scan(&queriedPass); err != nil {
		return fmt.Errorf("error reading password from DB: %v", err)
	}

	if err := bcrypt.CompareHashAndPassword([]byte(queriedPass), []byte(pass)); err != nil {
		return fmt.Errorf("incorrect password: %v", err)
	}

	return nil
}

func (db *DB) ChangePassword(id int64, oldPass, newPass string) error {
	tx, err := db.Begin()
	if err != nil {
		return fmt.Errorf("error starting transaction: %v", err)
	}

	var queriedPass string
	getQuery := `
    SELECT password
    FROM users
    WHERE id = ?
    `
	row := tx.QueryRow(getQuery, id)
	if err := row.Scan(&queriedPass); err != nil {
		if rollbackErr := tx.Rollback(); rollbackErr != nil {
			log.Fatalf("error: transaction error: %v, rollback error: %v", err, rollbackErr)
		}
		return fmt.Errorf("error reading password from DB: %v", err)
	}

	if err := bcrypt.CompareHashAndPassword([]byte(queriedPass), []byte(oldPass)); err != nil {
		if rollbackErr := tx.Rollback(); rollbackErr != nil {
			log.Fatalf("error: transaction error: %v, rollback error: %v", err, rollbackErr)
		}
		return fmt.Errorf("incorrect password: %v", err)
	}

	newHashedPass, err := bcrypt.GenerateFromPassword([]byte(newPass), bcrypt.DefaultCost)
	if err != nil {
		if rollbackErr := tx.Rollback(); rollbackErr != nil {
			log.Fatalf("error: transaction error: %v, rollback error: %v", err, rollbackErr)
		}
		return fmt.Errorf("error creating password hash: %v", err)
	}

	setQuery := `
    UPDATE users
    SET password = ?
    WHERE id = ?
    `
	if _, err = tx.Exec(setQuery, string(newHashedPass), id); err != nil {
		if rollbackErr := tx.Rollback(); rollbackErr != nil {
			log.Fatalf("error: transaction error: %v, rollback error: %v", err, rollbackErr)
		}
		return fmt.Errorf("error updating password in DB: %v", err)
	}

	if err = tx.Commit(); err != nil {
		return fmt.Errorf("error committing transaction: %v", err)
	}

	return nil
}

func (db *DB) CountEntries() (int64, error) {
	var count int64

	query := `SELECT COUNT(*) FROM users`
	row := db.QueryRow(query)
	if err := row.Scan(&count); err != nil {
		return 0, fmt.Errorf("error counting rows in user DB: %v", err)
	}

	return count, nil
}

// TODO: No need for ID field in general
func (db *DB) GetUser(id int64) (*User, error) {
	user := new(User)
	query := `
    SELECT id, username, first_name, last_name, role
    FROM users
    WHERE id = ?
    `

	row := db.QueryRow(query, id)
	if err := row.Scan(&user.ID, &user.UserName, &user.FirstName, &user.LastName, &user.Role); err != nil {
		return nil, fmt.Errorf("error reading user information: %v", err)
	}

	return user, nil
}

func (db *DB) AddTag(tagName string) error {
	query := "INSERT INTO tags (name) VALUES (?)"
	if _, err := db.Exec(query, tagName); err != nil {
		return fmt.Errorf("error inserting tag into DB: %v", err)
	}
	return nil
}

func (db *DB) GetTagList() ([]*Tag, error) {
	query := "SELECT id, name FROM tags"
	rows, err := db.Query(query)
	if err != nil {
		return nil, fmt.Errorf("error querying tags: %v", err)
	}

	tagList := make([]*Tag, 0)
	for rows.Next() {
		tag := new(Tag)
		if err = rows.Scan(&tag.ID, &tag.Name); err != nil {
			return nil, fmt.Errorf("error scanning tag row: %v", err)
		}
		tagList = append(tagList, tag)
	}

	return tagList, nil
}

func (db *DB) AddArticle(a *Article) (int64, error) {
	query := `
    INSERT INTO articles
        (title, description, content, published, author_id)
    VALUES
        (?, ?, ?, ?, ?)
    `

	result, err := db.Exec(query, a.Title, a.Desc, a.Content, a.Published, a.AuthorID)
	if err != nil {
		return 0, fmt.Errorf("error inserting article into DB: %v", err)
	}
	id, err := result.LastInsertId()
	if err != nil {
		return 0, fmt.Errorf("error retrieving last ID: %v", err)
	}

	return id, nil
}

func (db *DB) GetArticle(id int64) (*Article, error) {
	query := `
    SELECT title, created, description, content, published, author_id
    FROM articles
    WHERE id = ?
    `
	row := db.QueryRow(query, id)

	article := new(Article)
	var created []byte
	var err error

	if err := row.Scan(&article.Title, &created, &article.Desc,
		&article.Content, &article.Published, &article.AuthorID); err != nil {
		return nil, fmt.Errorf("error scanning article row: %v", err)
	}

	article.ID = id
	article.Created, err = time.Parse("2006-01-02 15:04:05", string(created))
	if err != nil {
		return nil, fmt.Errorf("error parsing created: %v", err)
	}

	return article, nil
}

func (db *DB) GetCertainArticles(published bool) ([]*Article, error) {
	query := `
    SELECT id, title, created, description, content, author_id
    FROM articles
    WHERE published = ?
    `
	rows, err := db.Query(query, published)
	if err != nil {
		return nil, fmt.Errorf("error querying articles: %v", err)
	}

	articleList := make([]*Article, 0)
	for rows.Next() {
		article := new(Article)
		var created []byte

		if err = rows.Scan(&article.ID, &article.Title, &created, &article.Desc,
			&article.Content, &article.AuthorID); err != nil {
			return nil, fmt.Errorf("error scanning article row: %v", err)
		}

		article.Published = false
		article.Created, err = time.Parse("2006-01-02 15:04:05", string(created))
		if err != nil {
			return nil, fmt.Errorf("error parsing created: %v", err)
		}

		articleList = append(articleList, article)
	}

	return articleList, nil
}

func (db *DB) WriteArticleTags(articleID int64, tagIDs []int64) error {
	tx, err := db.Begin()
	if err != nil {
		return fmt.Errorf("error starting transaction: %v", err)
	}

	for _, tagID := range tagIDs {
		query := `
        INSERT INTO articles_tags (article_id, tag_id)
        VALUES (?, ?)
        `
		if _, err := tx.Exec(query, articleID, tagID); err != nil {
			if rollbackErr := tx.Rollback(); rollbackErr != nil {
				log.Fatalf("error: transaction error: %v, rollback error: %v", err, rollbackErr)
			}
			return fmt.Errorf("error inserting into articles_tags: %v", err)
		}
	}

	if err = tx.Commit(); err != nil {
		return fmt.Errorf("error committing transaction: %v", err)
	}
	return nil
}

func (db *DB) GetArticleTags(articleID int64) ([]*Tag, error) {
	query := `
    SELECT t.id, t.name
    FROM articles a
    INNER JOIN articles_tags at ON a.id = at.article_id
    INNER JOIN tags t ON at.tag_id = t.id
    WHERE a.id = ?
    `
	rows, err := db.Query(query, articleID)
	if err != nil {
		return nil, fmt.Errorf("error querying articles_tags: %v", err)
	}

	tags := make([]*Tag, 0)
	for rows.Next() {
		tag := new(Tag)
		if err = rows.Scan(&tag.ID, &tag.Name); err != nil {
			return nil, fmt.Errorf("error scanning rows: %v", err)
		}
		tags = append(tags, tag)
	}

	return tags, nil
}