package store
import (
"context"
"database/sql"
"fmt"
"io/ioutil"
"os"
"path"
"path/filepath"
"strings"
// For SQLite support.
_ "github.com/glebarez/go-sqlite"
)
// ImageRecord holds all data regarding an image.
type ImageRecord struct {
UUID string
Filename string
OrigFilename string
Thumbname string
IsHidden bool
Tags []string
Timestamp int64
Title string
Description string
}
// ImageTag is a tag with related attributes.
type ImageTag struct {
Name string
Hidden bool
HiddenThumbnail string
PreviewEnabled bool
}
// ReadProvider specifies the necessary structure to read from a store.
type ReadProvider interface {
GetNumImageRecords(context.Context, ...string) (int, error)
GetImageRecords(context.Context, int, int, ...string) ([]ImageRecord, error)
GetImageRecord(
context.Context, string, ...string,
) (ImageRecord, string, string, error)
GetTags(context.Context) ([]ImageTag, error)
}
// WriteProvider specifies the necessary structure to write to a store.
type WriteProvider interface {
StoreImageRecords(context.Context, []ImageRecord) error
}
// SQLiteOption is the SQLiteProvider option func.
type SQLiteOption func(*SQLiteProvider)
// WithSchemaFile specifies what schemafile to use when setting up an empty db.
func WithSchemaFile(file string) SQLiteOption {
return func(s *SQLiteProvider) {
s.schemaFile = file
}
}
// SQLiteProvider connects to and handles SQLite calls.
type SQLiteProvider struct {
db *sql.DB
schemaFile string
}
// NewSQLiteProvider creates a new SQLite provider.
func NewSQLiteProvider(
databaseFile string, options ...SQLiteOption,
) (*SQLiteProvider, error) {
var s SQLiteProvider
for _, option := range options {
option(&s)
}
databaseDir := filepath.Dir(databaseFile)
if info, err := os.Stat(databaseDir); os.IsNotExist(err) {
if err := os.Mkdir(databaseDir, 0750); err != nil {
return nil, fmt.Errorf("database directory could not be created: %w", err)
}
} else if err == nil && !info.IsDir() {
return nil, fmt.Errorf(
"could not set up database, path exists but parent is not a folder",
)
}
var err error
s.db, err = sql.Open("sqlite", databaseFile)
if err != nil {
return nil, fmt.Errorf("could not open database: %w", err)
}
if _, err := os.Stat(databaseFile); os.IsNotExist(err) {
if s.schemaFile == "" {
return nil, fmt.Errorf("empty database and no schemafile defined")
}
schema, err := ioutil.ReadFile(s.schemaFile)
if err != nil {
return nil, fmt.Errorf(
"no database and schema file is missing/corrupt: %w", err,
)
}
_, err = s.db.Exec(string(schema))
if err != nil {
return nil, fmt.Errorf("could not execute schema: %w", err)
}
}
return &s, nil
}
// Close closes the database connection.
func (p *SQLiteProvider) Close() {
p.db.Close()
}
func stringInSlice(str string, slice []string) bool {
for _, s := range slice {
if str == s {
return true
}
}
return false
}
// Makes sure all of source is in target.
func sliceInSlice(source, target []string) bool {
var count int
for _, s := range source {
if stringInSlice(s, target) {
count++
}
}
return count == len(source)
}
// GetNumImageRecords does a simple image count lookup.
func (p *SQLiteProvider) GetNumImageRecords(
ctx context.Context, tags ...string,
) (int, error) {
var row *sql.Row
if len(tags) == 0 {
row = p.db.QueryRowContext(
ctx,
`SELECT COUNT(*) FROM
(
SELECT i.id
FROM images AS i
INNER JOIN imagetags AS it ON i.id=it.image_id
INNER JOIN tags AS t ON t.id=it.tag_id
GROUP BY i.id
)`,
)
} else {
row = p.db.QueryRowContext(
ctx,
`SELECT COUNT(*) FROM
(
SELECT i.filename, GROUP_CONCAT(t.name, ',')
FROM images AS i
INNER JOIN imagetags AS it ON i.id=it.image_id
INNER JOIN tags AS t ON t.id=it.tag_id
WHERE t.name IN ('`+strings.Join(tags, "','")+`')
GROUP BY i.id
)`,
)
}
var count int
if err := row.Scan(&count); err != nil {
return 0, fmt.Errorf("failed when getting count: %w", err)
}
return count, nil
}
// GetImageRecords gets all images that optionally matches tags.
func (p *SQLiteProvider) GetImageRecords(
ctx context.Context, offset, count int, tags ...string,
) ([]ImageRecord, error) {
var rows *sql.Rows
var err error
allTags, err := p.GetTags(ctx)
if err != nil {
return nil, err
}
if len(tags) == 0 {
rows, err = p.db.QueryContext(
ctx,
`SELECT i.uuid, i.filename, i.orig_filename, i.timestamp, GROUP_CONCAT(t.name, ','), IFNULL(md.title, ""), IFNULL(md.description, "")
FROM images AS i
LEFT JOIN metadata AS md ON i.id=md.image_id
INNER JOIN imagetags AS it ON i.id=it.image_id
INNER JOIN tags AS t ON t.id=it.tag_id
GROUP BY i.id ORDER BY timestamp DESC
LIMIT ? OFFSET ?`,
count, offset,
)
} else {
rows, err = p.db.QueryContext(
ctx, `SELECT i.uuid, i.filename, i.orig_filename, i.timestamp, GROUP_CONCAT(t.name, ','), IFNULL(md.title, ""), IFNULL(md.description, "")
FROM images AS i
LEFT JOIN metadata AS md ON i.id=md.image_id
INNER JOIN imagetags AS it ON i.id=it.image_id
INNER JOIN tags AS t ON t.id=it.tag_id
WHERE i.id IN (
SELECT itt.image_id FROM imagetags AS itt
INNER JOIN tags AS tt ON tt.id=itt.tag_id
WHERE i.id=itt.image_id
AND tt.name IN ('`+strings.Join(tags, "','")+`')
)
GROUP BY i.id ORDER BY timestamp DESC
LIMIT ? OFFSET ?`, count, offset,
)
}
if err != nil {
return nil, fmt.Errorf("could not query image list from db: %w", err)
}
defer rows.Close()
imageRecords := make([]ImageRecord, 0, 100)
for rows.Next() {
var timestamp int64
var uuid, filename, origFilename, tag, title, description string
if err := rows.Scan(&uuid, &filename, &origFilename, ×tamp, &tag, &title, &description); err != nil {
return nil, fmt.Errorf("could not scan image list from db: %w", err)
}
tagNames := strings.Split(tag, ",")
if len(tags) > 0 && !sliceInSlice(tags, tagNames) {
continue
}
var isHidden bool
skip:
for _, tag := range tagNames {
for i := range allTags {
if allTags[i].Hidden && tag == allTags[i].Name {
if !stringInSlice(tag, tags) {
isHidden = true
}
break skip
}
}
}
thumbname := strings.TrimSuffix(
filename, filepath.Ext(filename),
) + "_th.jpg"
imageRecords = append(imageRecords, ImageRecord{
UUID: uuid,
Filename: filename,
OrigFilename: origFilename,
Thumbname: thumbname,
IsHidden: isHidden,
Tags: tagNames,
Timestamp: timestamp,
Title: title,
Description: description,
})
}
if err := rows.Err(); err != nil {
return nil, fmt.Errorf("could not scan rows from db: %w", err)
}
return imageRecords, nil
}
// GetImageRecord fetches all data for a single image, and next/prev uuids.
func (p *SQLiteProvider) GetImageRecord(
ctx context.Context, uuid string, tags ...string,
) (ImageRecord, string, string, error) {
var record ImageRecord
var prev, next string
var rows *sql.Rows
var err error
row := p.db.QueryRowContext(
ctx,
`SELECT i.id, i.filename, i.orig_filename, i.timestamp, GROUP_CONCAT(t.name, ','), IFNULL(md.title, ""), IFNULL(md.description, "")
FROM images AS i
LEFT JOIN metadata AS md ON i.id=md.image_id
INNER JOIN imagetags AS it ON i.id=it.image_id
INNER JOIN tags AS t ON t.id=it.tag_id
WHERE i.uuid=?
GROUP BY i.id ORDER BY timestamp DESC`,
uuid,
)
var imageID int
var timestamp int64
var filename, origFilename, tag, title, description string
if err := row.Scan(&imageID, &filename, &origFilename, ×tamp, &tag, &title, &description); err != nil {
return ImageRecord{}, "", "",
fmt.Errorf("could not scan image from db: %w", err)
}
record.Filename = filename
record.OrigFilename = filepath.Base(origFilename)
record.UUID = uuid
record.Tags = strings.Split(tag, ",")
record.Timestamp = timestamp
record.Title = title
record.Description = description
if len(tags) == 0 {
rows, err = p.db.QueryContext(
ctx,
`SELECT i.id, i.filename
FROM images AS i
INNER JOIN imagetags AS it ON i.id=it.image_id
INNER JOIN tags AS t ON t.id=it.tag_id
GROUP BY i.id ORDER BY timestamp DESC`,
)
} else {
rows, err = p.db.QueryContext(
ctx, `SELECT i.id, i.filename
FROM images AS i
INNER JOIN imagetags AS it ON i.id=it.image_id
INNER JOIN tags AS t ON t.id=it.tag_id
WHERE i.id IN (
SELECT itt.image_id FROM imagetags AS itt
INNER JOIN tags AS tt ON tt.id=itt.tag_id
WHERE i.id=itt.image_id
AND tt.name IN ('`+strings.Join(tags, "','")+`')
)
GROUP BY i.id ORDER BY timestamp DESC`,
)
}
if err != nil {
return ImageRecord{}, "", "",
fmt.Errorf("could not query image list from db: %w", err)
}
defer rows.Close()
for rows.Next() {
var id int
var filename string
if err := rows.Scan(&id, &filename); err != nil {
return ImageRecord{}, "", "",
fmt.Errorf("could not scan image list from db: %w", err)
}
if id == imageID {
if rows.Next() {
if err := rows.Scan(&id, &next); err != nil {
return ImageRecord{}, "", "",
fmt.Errorf("could not scan next image from db: %w", err)
}
}
break
}
prev = filename
}
return record, prev, next, nil
}
// GetTags fetches all the tags from the database.
func (p *SQLiteProvider) GetTags(ctx context.Context) ([]ImageTag, error) {
rows, err := p.db.QueryContext(ctx, "SELECT name, hidden, hidden_thumbnail, preview_enabled FROM tags")
if err != nil {
return nil, fmt.Errorf("could not query image list from db: %w", err)
}
defer rows.Close()
tags := make([]ImageTag, 0, 10)
for rows.Next() {
var name string
var hidden bool
var thumbnail string
var preview bool
if err := rows.Scan(&name, &hidden, &thumbnail, &preview); err != nil {
return nil, fmt.Errorf("could not get tags from db: %w", err)
}
tags = append(tags, ImageTag{
Name: name,
Hidden: hidden,
HiddenThumbnail: thumbnail,
PreviewEnabled: preview,
})
}
return tags, nil
}
// StoreImageRecords stores images to the sqlite store.
// TODO: Call-specific data type instead of ImageRecord.
func (p *SQLiteProvider) StoreImageRecords(
ctx context.Context, records []ImageRecord,
) error {
tags := make(map[string]int64)
for _, record := range records {
for _, tag := range record.Tags {
if _, ok := tags[tag]; ok {
continue
}
var id int64
row := p.db.QueryRowContext(ctx, "SELECT id FROM tags WHERE name=?", tag)
if err := row.Scan(&id); err != nil {
if err != sql.ErrNoRows {
return fmt.Errorf("could not look up tag %s: %w", tag, err)
}
result, err := p.db.ExecContext(
ctx, "INSERT INTO tags (name, hidden) VALUES(?, false)", tag,
)
if err != nil {
return fmt.Errorf("could not find or create tag %s: %w", tag, err)
}
id, _ = result.LastInsertId()
}
tags[tag] = id
}
}
for _, record := range records {
result, err := p.db.ExecContext(
ctx,
`INSERT INTO images (uuid, filename, orig_filename, timestamp)
VALUES (?, ?, ?, ?)`,
record.UUID, path.Base(record.Filename), path.Base(record.OrigFilename),
record.Timestamp,
)
if err != nil {
return fmt.Errorf("could not execute store query: %w", err)
}
imageID, _ := result.LastInsertId()
for _, tag := range record.Tags {
tagID := tags[tag]
_, err := p.db.ExecContext(
ctx, "INSERT INTO imagetags (image_id, tag_id) VALUES(?, ?)",
imageID, tagID,
)
if err != nil {
return fmt.Errorf(
"could not link image %d with tag %d: %w", imageID, tagID, err,
)
}
}
_, err = p.db.ExecContext(
ctx, "INSERT INTO metadata (image_id, title, description) VALUES(?, ?, ?)",
imageID, record.Title, record.Description,
)
if err != nil {
return fmt.Errorf(
"could not link image %d with metadata %s, %s: %w", imageID, record.Title, record.Description, err,
)
}
}
return nil
}