paste-abroad/internel/storage/sqlite/sqlite.go

package sqlite

import (
    "github.com/leafee98/paste-abroad/internel/storage"
    "github.com/leafee98/paste-abroad/internel/utils"

    "github.com/mattn/go-sqlite3"

    "database/sql"
    "errors"
    "time"
)

type StorageSqlite struct {
    db *sql.DB
    idLength int
}

func New(db_path string, idLength int) (*StorageSqlite, error) {
    var s StorageSqlite
    var err error

    s.idLength = idLength

    s.db, err = sql.Open("sqlite3", db_path)
    if err != nil {
        return &s, err
    }

    if err := initDatabase(s.db); err != nil {
        return nil, err
    }

    return &s, nil
}

func initDatabase(db *sql.DB) error {
    sqlStmt := `
        create table if not exists paste (
            id text unique,
            content blob,
            plain integer,
            expire integer
        );

        create index if not exists paste_id
            on paste ( id );

        create index if not exists paste_expire
            on paste ( expire );
    `

    if _, err := db.Exec(sqlStmt); err != nil {
        return err
    }

    return nil
}

func (s *StorageSqlite) Close() error {
    return s.db.Close()
}

func (s *StorageSqlite) Save(p *storage.Paste) (string, error) {
    stmt, err := s.db.Prepare(`insert into paste (id, content, plain, expire) values (?, ?, ?, ?);`)
    if err != nil {
        return "", err
    }
    defer stmt.Close()

    var id string
    var plain = 0

    if p.Plain {
        plain = 1
    }

    for true {
        id = utils.GenerateId(s.idLength)
        _, err = stmt.Exec(id, p.Content, plain, p.Expire)
        if err == nil {
            break
        } else if errors.Is(err, sqlite3.ErrConstraint) {
            continue
        } else {
            return "", err
        }
    }

    return id, nil
}

func (s *StorageSqlite) Get(id string) (*storage.Paste, error) {
    stmt, err := s.db.Prepare(`select content, plain, expire from paste where id = ?`)
    if err != nil {
        return nil, err
    }
    rows, err := stmt.Query(id)

    var p storage.Paste

    if !rows.Next() {
        return nil, nil
    }

    var plain = 0

    if err = rows.Scan(&p.Content, &plain, &p.Expire); err != nil {
        return nil, err
    }

    if plain != 0 {
        p.Plain = true
    }

    return &p, nil
}

func (s *StorageSqlite) Purge() (int64, error) {
    stmt, err := s.db.Prepare(`delete from paste where expire < ?`)
    if err != nil {
        return 0, err
    }

    affected, err := stmt.Exec(time.Now().UnixMilli())
    if err != nil {
        return 0, err
    }

    if affectedRows, err := affected.RowsAffected() ; err != nil {
        return 0, err
    } else {
        return affectedRows, nil
    }
}