fediarchiver/db.go

package main

import (
	"database/sql"
	"encoding/json"
	"fmt"

	_ "github.com/lib/pq"
	"github.com/rs/zerolog/log"
)

/* Statements */
var stmtGetUserIDsByDomain *sql.Stmt
var stmtCreateUser *sql.Stmt
var stmtCheckUserByName *sql.Stmt
var stmtGetLastStatus *sql.Stmt
var stmtSetLastStatus *sql.Stmt
var stmtCreateStatus *sql.Stmt
var stmtCheckStatus *sql.Stmt
var stmtGetDomains *sql.Stmt
var stmtUpdateArchiveOrg *sql.Stmt
var stmtUpdateArchiveIs *sql.Stmt
var stmtUpdateArchiveOrgJSON *sql.Stmt
var stmtUpdateArchiveIsJSON *sql.Stmt
var stmtGetNullArchiveIs *sql.Stmt
var stmtGetNullArchiveOrg *sql.Stmt
var stmtGetNullArchiveIsJSON *sql.Stmt
var stmtGetNullArchiveOrgJSON *sql.Stmt

const (
	msgOpenDBFailed     = "Failed to open DB"
	msgPrepareSQLFailed = "Failed to prepare SQL statement"
)

func openDB(cfg *Config) *sql.DB {
	connStr := fmt.Sprintf("postgres://%s:%s@%s:%d/%s", cfg.DBUser, cfg.DBPass, cfg.DBHost, cfg.DBPort, cfg.DBDB)
	db, err := sql.Open("postgres", connStr)
	db.SetMaxOpenConns(cfg.DBMaxOpenConn)
	db.SetMaxIdleConns(cfg.DBMaxIdleConn)
	if err != nil {
		log.Panic().Err(err).Msg(msgOpenDBFailed)
	}

	/* Create Tables */
	_, err = db.Exec("CREATE TABLE IF NOT EXISTS statuses (id TEXT PRIMARY KEY,user_id TEXT NOT NULL,domain TEXT NOT NULL,status JSONB NOT NULL,archiveorg TEXT DEFAULT NULL,archiveis TEXT DEFAULT NULL,archiveorg_json TEXT DEFAULT NULL,archiveis_json TEXT DEFAULT NULL)")
	if err != nil {
		log.Panic().Err(err).Msg(msgPrepareSQLFailed)
	}

	_, err = db.Exec("CREATE TABLE IF NOT EXISTS users (id TEXT PRIMARY KEY,user_id TEXT NOT NULL,username TEXT NOT NULL,domain TEXT NOT NULL,last_status TEXT DEFAULT 0 NOT NULL)")
	if err != nil {
		log.Panic().Err(err).Msg(msgPrepareSQLFailed)
	}

	/* Prepare Statements */
	stmtCreateStatus, err = db.Prepare("INSERT INTO statuses(id,user_id,domain,status) VALUES($1,$2,$3,$4)")
	if err != nil {
		log.Panic().Err(err).Msg(msgPrepareSQLFailed)
	}
	stmtCheckStatus, err = db.Prepare("SELECT EXISTS(SELECT 1 FROM statuses WHERE id=$1)")
	if err != nil {
		log.Panic().Err(err).Msg(msgPrepareSQLFailed)
	}
	stmtGetUserIDsByDomain, err = db.Prepare("SELECT user_id FROM users WHERE domain=$1")
	if err != nil {
		log.Panic().Err(err).Msg(msgPrepareSQLFailed)
	}
	stmtCreateUser, err = db.Prepare("INSERT INTO users(id,user_id,username,domain) VALUES($1,$2,$3,$4)")
	if err != nil {
		log.Panic().Err(err).Msg(msgPrepareSQLFailed)
	}
	stmtCheckUserByName, err = db.Prepare("SELECT EXISTS(SELECT 1 FROM users WHERE username=$1 AND domain=$2)")
	if err != nil {
		log.Panic().Err(err).Msg(msgPrepareSQLFailed)
	}
	stmtGetLastStatus, err = db.Prepare("SELECT last_status FROM users WHERE id=$1")
	if err != nil {
		log.Panic().Err(err).Msg(msgPrepareSQLFailed)
	}
	stmtSetLastStatus, err = db.Prepare("UPDATE users SET last_status=$1 WHERE id=$2")
	if err != nil {
		log.Panic().Err(err).Msg(msgPrepareSQLFailed)
	}
	stmtGetDomains, err = db.Prepare("SELECT DISTINCT domain FROM users")
	if err != nil {
		log.Panic().Err(err).Msg(msgPrepareSQLFailed)
	}
	stmtUpdateArchiveIs, err = db.Prepare("UPDATE statuses SET archiveis=$1 WHERE id=$2")
	if err != nil {
		log.Panic().Err(err).Msg(msgPrepareSQLFailed)
	}
	stmtUpdateArchiveOrg, err = db.Prepare("UPDATE statuses SET archiveorg=$1 WHERE id=$2")
	if err != nil {
		log.Panic().Err(err).Msg(msgPrepareSQLFailed)
	}
	stmtUpdateArchiveIsJSON, err = db.Prepare("UPDATE statuses SET archiveis_json=$1 WHERE id=$2")
	if err != nil {
		log.Panic().Err(err).Msg(msgPrepareSQLFailed)
	}
	stmtUpdateArchiveOrgJSON, err = db.Prepare("UPDATE statuses SET archiveorg_json=$1 WHERE id=$2")
	if err != nil {
		log.Panic().Err(err).Msg(msgPrepareSQLFailed)
	}
	stmtGetNullArchiveIs, err = db.Prepare("SELECT id,domain,status->'url' AS url FROM statuses WHERE archiveis is NULL")
	if err != nil {
		log.Panic().Err(err).Msg(msgPrepareSQLFailed)
	}
	stmtGetNullArchiveOrg, err = db.Prepare("SELECT id,domain,status->'url' AS url FROM statuses WHERE archiveorg is NULL")
	if err != nil {
		log.Panic().Err(err).Msg(msgPrepareSQLFailed)
	}
	stmtGetNullArchiveIsJSON, err = db.Prepare("SELECT id,domain,status->'id' AS id FROM statuses WHERE archiveis_json is NULL")
	if err != nil {
		log.Panic().Err(err).Msg(msgPrepareSQLFailed)
	}
	stmtGetNullArchiveOrgJSON, err = db.Prepare("SELECT id,domain,status->'id' AS id FROM statuses WHERE archiveorg_json is NULL")
	if err != nil {
		log.Panic().Err(err).Msg(msgPrepareSQLFailed)
	}

	return db
}

/* User Table */
func sqlCreateUser(userID string, username string, domain string) error {
	// log.Debug().Str("Username", username).Str("UserID", userID).Str("Domain", domain).Msg("sqlCreateUser")
	_, err := stmtCreateUser.Exec(fmt.Sprintf("%s@%s", userID, domain), userID, username, domain)
	return err
}
func sqlCheckUserByName(username string, domain string) (bool, error) {
	var ok bool
	row := stmtCheckUserByName.QueryRow(username, domain)
	err := row.Scan(&ok)
	return ok, err
}
func sqlGetUserIDsByDomain(domain string) ([]QueryUserIDs, error) {
	var results []QueryUserIDs
	rows, err := stmtGetUserIDsByDomain.Query(domain)
	if err != nil {
		return nil, err
	}
	defer rows.Close()
	for rows.Next() {
		var result QueryUserIDs
		err = rows.Scan(&result.UserID)
		if err != nil {
			return nil, err
		}
		results = append(results, result)
	}
	return results, nil
}

/* Status Table */
func sqlGetLastStatus(userID string, domain string) (string, error) {
	log.Debug().Str("UserID", userID).Str("Domain", domain).Msg("sqlGetLastStatus")
	var lastStatus string
	row := stmtGetLastStatus.QueryRow(fmt.Sprintf("%s@%s", userID, domain))
	err := row.Scan(&lastStatus)
	return lastStatus, err
}

func sqlSetLastStatus(userID string, domain string, id string) {
	log.Debug().Str("UserID", userID).Str("Domain", domain).Msg("sqlSetLastStatus")
	_, err := stmtSetLastStatus.Exec(id, fmt.Sprintf("%s@%s", userID, domain))
	if err != nil {
		log.Error().Str("UserID", userID).Str("Domain", domain).Str("Latest", id).Err(err).Msg("sqlSetLastStatus")
	}
}

func sqlCheckStatus(statusID string, domain string) (bool, error) {
	// log.Debug().Str("StatusID", statusID).Str("Domain", domain).Msg("sqlCheckStatus")
	var ok bool
	row := stmtCheckStatus.QueryRow(fmt.Sprintf("%s@%s", statusID, domain))
	err := row.Scan(&ok)
	return ok, err
}

func sqlCreateStatus(userID string, domain string, status map[string]interface{}) error {
	// log.Debug().Str("UserID", userID).Str("Domain", domain).Msg("sqlCreateStatus")
	statusStr, err := json.Marshal(status)
	if err != nil {
		return err
	}
	_, err = stmtCreateStatus.Exec(fmt.Sprintf("%s@%s", status["id"], domain), userID, domain, statusStr)
	return err
}

func sqlGetDomains() ([]QueryDomain, error) {
	var results []QueryDomain
	rows, err := stmtGetDomains.Query()
	if err != nil {
		return nil, err
	}
	defer rows.Close()
	for rows.Next() {
		var result QueryDomain
		err = rows.Scan(&result.Domain)
		if err != nil {
			return nil, err
		}
		results = append(results, result)
	}
	return results, nil
}

func sqlUpdateArchiveIs(ID string, url string) {
	log.Info().Str("ID", ID).Str("URL", url).Msg("Archive.is Recorded")
	_, err := stmtUpdateArchiveIs.Exec(url, ID)
	if err != nil {
		log.Error().Str("ID", ID).Str("URL", url).Err(err).Msg("sqlUpdateArchiveIs")
	}
}

func sqlUpdateArchiveOrg(ID string, url string) {
	log.Info().Str("ID", ID).Str("URL", url).Msg("Archive.org Recorded")
	_, err := stmtUpdateArchiveOrg.Exec(url, ID)
	if err != nil {
		log.Error().Str("ID", ID).Str("URL", url).Err(err).Msg("sqlUpdateArchiveOrg")
	}
}

func sqlUpdateArchiveIsJSON(statusID string, domain string, url string) {
	log.Info().Str("StatusID", statusID).Str("Domain", domain).Str("URL", url).Msg("Archive.is JSON Recorded")
	_, err := stmtUpdateArchiveIsJSON.Exec(url, statusID)
	if err != nil {
		log.Error().Str("StatusID", statusID).Str("Domain", domain).Str("URL", url).Err(err).Msg("sqlUpdateArchiveIsJSON")
	}
}

func sqlUpdateArchiveOrgJSON(statusID string, domain string, url string) {
	log.Info().Str("StatusID", statusID).Str("Domain", domain).Str("URL", url).Msg("Archive.org JSON Recorded")
	_, err := stmtUpdateArchiveOrgJSON.Exec(url, statusID)
	if err != nil {
		log.Error().Str("StatusID", statusID).Str("Domain", domain).Str("URL", url).Err(err).Msg("sqlUpdateArchiveOrgJSON")
	}
}

func sqlGetNullArchiveIs() ([]QueryURL, error) {
	var results []QueryURL
	rows, err := stmtGetNullArchiveIs.Query()
	if err != nil {
		return nil, err
	}
	defer rows.Close()
	for rows.Next() {
		var result QueryURL
		err = rows.Scan(&result.ID, &result.Domain, &result.URL)
		if err != nil {
			return nil, err
		}
		results = append(results, result)
	}
	return results, nil
}

func sqlGetNullArchiveOrg() ([]QueryURL, error) {
	var results []QueryURL
	rows, err := stmtGetNullArchiveOrg.Query()
	if err != nil {
		return nil, err
	}
	defer rows.Close()
	for rows.Next() {
		var result QueryURL
		err = rows.Scan(&result.ID, &result.Domain, &result.URL)
		if err != nil {
			return nil, err
		}
		results = append(results, result)
	}
	return results, nil
}

func sqlGetNullArchiveIsJSON() ([]QueryStatus, error) {
	var results []QueryStatus
	rows, err := stmtGetNullArchiveIsJSON.Query()
	if err != nil {
		return nil, err
	}
	defer rows.Close()
	for rows.Next() {
		var result QueryStatus
		err = rows.Scan(&result.ID, &result.Domain, &result.StatusID)
		if err != nil {
			return nil, err
		}
		results = append(results, result)
	}
	return results, nil
}

func sqlGetNullArchiveOrgJSON() ([]QueryStatus, error) {
	var results []QueryStatus
	rows, err := stmtGetNullArchiveOrgJSON.Query()
	if err != nil {
		return nil, err
	}
	defer rows.Close()
	for rows.Next() {
		var result QueryStatus
		err = rows.Scan(&result.ID, &result.Domain, &result.StatusID)
		if err != nil {
			return nil, err
		}
		results = append(results, result)
	}
	return results, nil
}