102 lines
2.0 KiB
Go
102 lines
2.0 KiB
Go
// Package db opens a *sql.DB for either SQLite or PostgreSQL and rewrites
|
|
// the shared `?` placeholders into PostgreSQL's `$n` form.
|
|
package db
|
|
|
|
import (
|
|
"database/sql"
|
|
"fmt"
|
|
"os"
|
|
"path/filepath"
|
|
"strings"
|
|
|
|
_ "github.com/lib/pq"
|
|
_ "modernc.org/sqlite"
|
|
)
|
|
|
|
type Dialect int
|
|
|
|
const (
|
|
SQLite Dialect = iota
|
|
Postgres
|
|
)
|
|
|
|
type DB struct {
|
|
*sql.DB
|
|
Dialect Dialect
|
|
}
|
|
|
|
func Open(driver, dsn string) (*DB, error) {
|
|
var d Dialect
|
|
switch driver {
|
|
case "sqlite":
|
|
d = SQLite
|
|
if err := os.MkdirAll(filepath.Dir(dsn), 0o755); err != nil {
|
|
return nil, err
|
|
}
|
|
dsn = addSQLiteParams(dsn)
|
|
case "postgres":
|
|
d = Postgres
|
|
default:
|
|
return nil, fmt.Errorf("unsupported driver %q", driver)
|
|
}
|
|
|
|
pool, err := sql.Open(driverName(driver), dsn)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
if d == SQLite {
|
|
// SQLite is single-writer; keep a small pool to avoid "database is locked".
|
|
pool.SetMaxOpenConns(1)
|
|
} else {
|
|
pool.SetMaxOpenConns(10)
|
|
}
|
|
if err := pool.Ping(); err != nil {
|
|
return nil, fmt.Errorf("connect %s: %w", driver, err)
|
|
}
|
|
return &DB{DB: pool, Dialect: d}, nil
|
|
}
|
|
|
|
func driverName(driver string) string {
|
|
if driver == "postgres" {
|
|
return "postgres"
|
|
}
|
|
return "sqlite"
|
|
}
|
|
|
|
func addSQLiteParams(dsn string) string {
|
|
if strings.HasPrefix(dsn, "file:") || strings.Contains(dsn, "?") {
|
|
return dsn
|
|
}
|
|
return dsn + "?_pragma=busy_timeout(5000)&_pragma=journal_mode(WAL)&_pragma=foreign_keys(1)"
|
|
}
|
|
|
|
// Rebind converts `?` placeholders to `$1..$n` on PostgreSQL.
|
|
func (d *DB) Rebind(q string) string {
|
|
if d.Dialect != Postgres {
|
|
return q
|
|
}
|
|
var b strings.Builder
|
|
b.Grow(len(q) + 8)
|
|
n := 0
|
|
for _, r := range q {
|
|
if r == '?' {
|
|
n++
|
|
b.WriteString("$")
|
|
b.WriteString(fmt.Sprint(n))
|
|
continue
|
|
}
|
|
b.WriteRune(r)
|
|
}
|
|
return b.String()
|
|
}
|
|
|
|
func (d *DB) Q(q string) string { return d.Rebind(q) }
|
|
|
|
// AutoInc returns the column definition for an auto-incrementing primary key.
|
|
func (d *DB) AutoInc() string {
|
|
if d.Dialect == Postgres {
|
|
return "BIGSERIAL PRIMARY KEY"
|
|
}
|
|
return "INTEGER PRIMARY KEY AUTOINCREMENT"
|
|
}
|