mirror of
https://github.com/gofiber/storage.git
synced 2025-10-06 00:57:38 +08:00
Added original postgre implementation back
This commit is contained in:
213
postgres/postgres.go
Normal file
213
postgres/postgres.go
Normal file
@@ -0,0 +1,213 @@
|
||||
package postgres
|
||||
|
||||
import (
|
||||
"database/sql"
|
||||
"errors"
|
||||
"fmt"
|
||||
"net/url"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
_ "github.com/lib/pq"
|
||||
)
|
||||
|
||||
// Storage interface that is implemented by storage providers
|
||||
type Storage struct {
|
||||
db *sql.DB
|
||||
gcInterval time.Duration
|
||||
done chan struct{}
|
||||
|
||||
sqlSelect string
|
||||
sqlInsert string
|
||||
sqlDelete string
|
||||
sqlReset string
|
||||
sqlGC string
|
||||
}
|
||||
|
||||
var (
|
||||
checkSchemaMsg = "The `v` row has an incorrect data type. " +
|
||||
"It should be BYTEA but is instead %s. This will cause encoding-related panics if the DB is not migrated (see https://github.com/gofiber/storage/blob/main/MIGRATE.md)."
|
||||
dropQuery = `DROP TABLE IF EXISTS %s;`
|
||||
initQuery = []string{
|
||||
`CREATE TABLE IF NOT EXISTS %s (
|
||||
k VARCHAR(64) PRIMARY KEY NOT NULL DEFAULT '',
|
||||
v BYTEA NOT NULL,
|
||||
e BIGINT NOT NULL DEFAULT '0'
|
||||
);`,
|
||||
`CREATE INDEX IF NOT EXISTS e ON %s (e);`,
|
||||
}
|
||||
checkSchemaQuery = `SELECT DATA_TYPE FROM INFORMATION_SCHEMA.COLUMNS
|
||||
WHERE table_name = '%s' AND COLUMN_NAME = 'v';`
|
||||
)
|
||||
|
||||
// New creates a new storage
|
||||
func New(config ...Config) *Storage {
|
||||
// Set default config
|
||||
cfg := configDefault(config...)
|
||||
|
||||
// Create data source name
|
||||
var dsn string = "postgresql://"
|
||||
if cfg.Username != "" {
|
||||
dsn += url.QueryEscape(cfg.Username)
|
||||
}
|
||||
if cfg.Password != "" {
|
||||
dsn += ":" + cfg.Password
|
||||
}
|
||||
if cfg.Username != "" || cfg.Password != "" {
|
||||
dsn += "@"
|
||||
}
|
||||
dsn += fmt.Sprintf("%s:%d", url.QueryEscape(cfg.Host), cfg.Port)
|
||||
dsn += fmt.Sprintf("/%s?connect_timeout=%d&sslmode=%s",
|
||||
url.QueryEscape(cfg.Database),
|
||||
int64(cfg.timeout.Seconds()),
|
||||
cfg.SslMode,
|
||||
)
|
||||
|
||||
// Create db
|
||||
db, err := sql.Open("postgres", dsn)
|
||||
if err != nil {
|
||||
panic(err)
|
||||
}
|
||||
|
||||
// Set database options
|
||||
db.SetMaxOpenConns(cfg.maxOpenConns)
|
||||
db.SetMaxIdleConns(cfg.maxIdleConns)
|
||||
db.SetConnMaxLifetime(cfg.connMaxLifetime)
|
||||
|
||||
// Ping database
|
||||
if err := db.Ping(); err != nil {
|
||||
panic(err)
|
||||
}
|
||||
|
||||
// Drop table if set to true
|
||||
if cfg.Reset {
|
||||
if _, err = db.Exec(fmt.Sprintf(dropQuery, cfg.Table)); err != nil {
|
||||
_ = db.Close()
|
||||
panic(err)
|
||||
}
|
||||
}
|
||||
|
||||
// Init database queries
|
||||
for _, query := range initQuery {
|
||||
if _, err := db.Exec(fmt.Sprintf(query, cfg.Table)); err != nil {
|
||||
_ = db.Close()
|
||||
|
||||
panic(err)
|
||||
}
|
||||
}
|
||||
|
||||
// Create storage
|
||||
store := &Storage{
|
||||
db: db,
|
||||
gcInterval: cfg.GCInterval,
|
||||
done: make(chan struct{}),
|
||||
sqlSelect: fmt.Sprintf(`SELECT v, e FROM %s WHERE k=$1;`, cfg.Table),
|
||||
sqlInsert: fmt.Sprintf("INSERT INTO %s (k, v, e) VALUES ($1, $2, $3) ON CONFLICT (k) DO UPDATE SET v = $4, e = $5", cfg.Table),
|
||||
sqlDelete: fmt.Sprintf("DELETE FROM %s WHERE k=$1", cfg.Table),
|
||||
sqlReset: fmt.Sprintf("TRUNCATE TABLE %s;", cfg.Table),
|
||||
sqlGC: fmt.Sprintf("DELETE FROM %s WHERE e <= $1 AND e != 0", cfg.Table),
|
||||
}
|
||||
|
||||
store.checkSchema(cfg.Table)
|
||||
|
||||
// Start garbage collector
|
||||
go store.gcTicker()
|
||||
|
||||
return store
|
||||
}
|
||||
|
||||
var noRows = errors.New("sql: no rows in result set")
|
||||
|
||||
// Get value by key
|
||||
func (s *Storage) Get(key string) ([]byte, error) {
|
||||
if len(key) <= 0 {
|
||||
return nil, nil
|
||||
}
|
||||
row := s.db.QueryRow(s.sqlSelect, key)
|
||||
// Add db response to data
|
||||
var (
|
||||
data = []byte{}
|
||||
exp int64 = 0
|
||||
)
|
||||
if err := row.Scan(&data, &exp); err != nil {
|
||||
if err == sql.ErrNoRows {
|
||||
return nil, nil
|
||||
}
|
||||
return nil, err
|
||||
}
|
||||
|
||||
// If the expiration time has already passed, then return nil
|
||||
if exp != 0 && exp <= time.Now().Unix() {
|
||||
return nil, nil
|
||||
}
|
||||
|
||||
return data, nil
|
||||
}
|
||||
|
||||
// Set key with value
|
||||
func (s *Storage) Set(key string, val []byte, exp time.Duration) error {
|
||||
// Ain't Nobody Got Time For That
|
||||
if len(key) <= 0 || len(val) <= 0 {
|
||||
return nil
|
||||
}
|
||||
var expSeconds int64
|
||||
if exp != 0 {
|
||||
expSeconds = time.Now().Add(exp).Unix()
|
||||
}
|
||||
_, err := s.db.Exec(s.sqlInsert, key, val, expSeconds, val, expSeconds)
|
||||
return err
|
||||
}
|
||||
|
||||
// Delete entry by key
|
||||
func (s *Storage) Delete(key string) error {
|
||||
// Ain't Nobody Got Time For That
|
||||
if len(key) <= 0 {
|
||||
return nil
|
||||
}
|
||||
_, err := s.db.Exec(s.sqlDelete, key)
|
||||
return err
|
||||
}
|
||||
|
||||
// Reset all entries, including unexpired
|
||||
func (s *Storage) Reset() error {
|
||||
_, err := s.db.Exec(s.sqlReset)
|
||||
return err
|
||||
}
|
||||
|
||||
// Close the database
|
||||
func (s *Storage) Close() error {
|
||||
s.done <- struct{}{}
|
||||
return s.db.Close()
|
||||
}
|
||||
|
||||
// gcTicker starts the gc ticker
|
||||
func (s *Storage) gcTicker() {
|
||||
ticker := time.NewTicker(s.gcInterval)
|
||||
defer ticker.Stop()
|
||||
for {
|
||||
select {
|
||||
case <-s.done:
|
||||
return
|
||||
case t := <-ticker.C:
|
||||
s.gc(t)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// gc deletes all expired entries
|
||||
func (s *Storage) gc(t time.Time) {
|
||||
_, _ = s.db.Exec(s.sqlGC, t.Unix())
|
||||
}
|
||||
|
||||
func (s *Storage) checkSchema(tableName string) {
|
||||
var data []byte
|
||||
|
||||
row := s.db.QueryRow(fmt.Sprintf(checkSchemaQuery, tableName))
|
||||
if err := row.Scan(&data); err != nil {
|
||||
panic(err)
|
||||
}
|
||||
|
||||
if strings.ToLower(string(data)) != "bytea" {
|
||||
fmt.Printf(checkSchemaMsg, string(data))
|
||||
}
|
||||
}
|
Reference in New Issue
Block a user