2021-04-22 01:39:24 -07:00
|
|
|
package main
|
|
|
|
|
|
|
|
import (
|
|
|
|
"database/sql"
|
|
|
|
"fmt"
|
|
|
|
"os"
|
2021-04-22 04:40:56 -07:00
|
|
|
"regexp"
|
2021-04-22 01:39:24 -07:00
|
|
|
|
|
|
|
_ "github.com/lib/pq"
|
|
|
|
_ "github.com/mattn/go-sqlite3"
|
|
|
|
)
|
|
|
|
|
2021-04-22 03:33:50 -07:00
|
|
|
const (
|
|
|
|
DBTypeSQLite3 = iota
|
|
|
|
DBTypePSQL
|
|
|
|
)
|
|
|
|
|
2021-04-22 01:39:24 -07:00
|
|
|
type DB struct {
|
|
|
|
*sql.DB
|
2021-04-22 03:33:50 -07:00
|
|
|
dbType int
|
2021-04-22 01:39:24 -07:00
|
|
|
}
|
|
|
|
|
|
|
|
// OpenSQLite3 opens and returns a SQLite3 database
|
|
|
|
func OpenSQLite3(name, initSQL string) (*DB, error) {
|
|
|
|
os.Mkdir("storage", 0777)
|
|
|
|
|
|
|
|
db, err := sql.Open("sqlite3", "storage/"+name)
|
|
|
|
if err != nil {
|
|
|
|
return nil, err
|
|
|
|
}
|
|
|
|
|
|
|
|
if _, err := db.Exec(initSQL); err != nil {
|
|
|
|
db.Close()
|
|
|
|
return nil, err
|
|
|
|
}
|
|
|
|
|
2021-04-22 03:33:50 -07:00
|
|
|
return &DB{DB: db, dbType: DBTypeSQLite3}, nil
|
2021-04-22 01:39:24 -07:00
|
|
|
}
|
|
|
|
|
|
|
|
// OpenPSQL opens and returns a PostgreSQL database
|
2021-04-22 01:57:47 -07:00
|
|
|
func OpenPSQL(host, name, user, password, initSQL string, port int) (*DB, error) {
|
2021-04-22 01:39:24 -07:00
|
|
|
psqlconn := fmt.Sprintf("host=%s port=%d user=%s password=%s dbname=%s sslmode=disable", host, port, user, password, name)
|
|
|
|
|
|
|
|
db, err := sql.Open("postgres", psqlconn)
|
|
|
|
if err != nil {
|
|
|
|
return nil, err
|
|
|
|
}
|
|
|
|
|
|
|
|
err = db.Ping()
|
|
|
|
if err != nil {
|
|
|
|
db.Close()
|
|
|
|
return nil, err
|
|
|
|
}
|
|
|
|
|
2021-04-22 01:57:47 -07:00
|
|
|
if _, err := db.Exec(initSQL); err != nil {
|
|
|
|
db.Close()
|
|
|
|
return nil, err
|
|
|
|
}
|
|
|
|
|
2021-04-22 03:33:50 -07:00
|
|
|
return &DB{DB: db, dbType: DBTypePSQL}, nil
|
2021-04-22 01:39:24 -07:00
|
|
|
}
|
|
|
|
|
2021-04-22 03:33:50 -07:00
|
|
|
// Type returns the type of database that is being interacted with
|
|
|
|
func (db *DB) Type() int { return db.dbType }
|
|
|
|
|
2021-04-22 01:57:47 -07:00
|
|
|
// Exec executes a SQL statement
|
2021-04-22 02:00:49 -07:00
|
|
|
func (db *DB) Exec(sql string, values ...interface{}) (sql.Result, error) {
|
2021-04-22 04:40:56 -07:00
|
|
|
if db.Type() == DBTypeSQLite3 {
|
|
|
|
r, err := regexp.Compile("\\$+[0-9]")
|
|
|
|
if err != nil {
|
|
|
|
return nil, err
|
|
|
|
}
|
|
|
|
|
|
|
|
sql = r.ReplaceAllString(sql, "?")
|
2021-04-22 03:33:50 -07:00
|
|
|
}
|
|
|
|
return db.DB.Exec(sql, values...)
|
2021-04-22 01:57:47 -07:00
|
|
|
}
|
|
|
|
|
2021-04-22 04:40:56 -07:00
|
|
|
// QueryRow executes a SQL statement and stores the results
|
2021-04-22 02:17:03 -07:00
|
|
|
func (db *DB) QueryRow(sql string, values ...interface{}) *sql.Row {
|
2021-04-22 04:40:56 -07:00
|
|
|
if db.Type() == DBTypeSQLite3 {
|
|
|
|
r, err := regexp.Compile("\\$+[0-9]")
|
|
|
|
if err != nil {
|
|
|
|
return nil
|
|
|
|
}
|
|
|
|
|
|
|
|
sql = r.ReplaceAllString(sql, "?")
|
2021-04-22 03:33:50 -07:00
|
|
|
}
|
2021-04-22 02:17:03 -07:00
|
|
|
return db.DB.QueryRow(sql, values...)
|
2021-04-22 01:39:24 -07:00
|
|
|
}
|