package sqlite import ( "context" "database/sql" "errors" "fmt" "github.com/doug-martin/goqu/v9" "github.com/doug-martin/goqu/v9/exp" "github.com/jmoiron/sqlx" "gopkg.in/guregu/null.v4" "github.com/stashapp/stash/pkg/file" "github.com/stashapp/stash/pkg/logger" "github.com/stashapp/stash/pkg/models" "github.com/stashapp/stash/pkg/sliceutil/intslice" ) type table struct { table exp.IdentifierExpression idColumn exp.IdentifierExpression } type NotFoundError struct { ID int Table string } func (e *NotFoundError) Error() string { return fmt.Sprintf("id %d does not exist in %s", e.ID, e.Table) } func (t *table) insert(ctx context.Context, o interface{}) (sql.Result, error) { q := dialect.Insert(t.table).Prepared(true).Rows(o) ret, err := exec(ctx, q) if err != nil { return nil, fmt.Errorf("inserting into %s: %w", t.table.GetTable(), err) } return ret, nil } func (t *table) insertID(ctx context.Context, o interface{}) (int, error) { result, err := t.insert(ctx, o) if err != nil { return 0, err } ret, err := result.LastInsertId() if err != nil { return 0, err } return int(ret), nil } func (t *table) updateByID(ctx context.Context, id interface{}, o interface{}) error { q := dialect.Update(t.table).Prepared(true).Set(o).Where(t.byID(id)) if _, err := exec(ctx, q); err != nil { return fmt.Errorf("updating %s: %w", t.table.GetTable(), err) } return nil } func (t *table) byID(id interface{}) exp.Expression { return t.idColumn.Eq(id) } func (t *table) byIDInts(ids ...int) exp.Expression { ii := make([]interface{}, len(ids)) for i, id := range ids { ii[i] = id } return t.idColumn.In(ii...) } func (t *table) idExists(ctx context.Context, id interface{}) (bool, error) { q := dialect.Select(goqu.COUNT("*")).From(t.table).Where(t.byID(id)) var count int if err := querySimple(ctx, q, &count); err != nil { return false, err } return count == 1, nil } func (t *table) checkIDExists(ctx context.Context, id int) error { exists, err := t.idExists(ctx, id) if err != nil { return err } if !exists { return &NotFoundError{ID: id, Table: t.table.GetTable()} } return nil } func (t *table) destroyExisting(ctx context.Context, ids []int) error { for _, id := range ids { exists, err := t.idExists(ctx, id) if err != nil { return err } if !exists { return &NotFoundError{ ID: id, Table: t.table.GetTable(), } } } return t.destroy(ctx, ids) } func (t *table) destroy(ctx context.Context, ids []int) error { q := dialect.Delete(t.table).Where(t.idColumn.In(ids)) if _, err := exec(ctx, q); err != nil { return fmt.Errorf("destroying %s: %w", t.table.GetTable(), err) } return nil } // func (t *table) get(ctx context.Context, q *goqu.SelectDataset, dest interface{}) error { // tx, err := getTx(ctx) // if err != nil { // return err // } // sql, args, err := q.ToSQL() // if err != nil { // return fmt.Errorf("generating sql: %w", err) // } // return tx.GetContext(ctx, dest, sql, args...) // } type joinTable struct { table fkColumn exp.IdentifierExpression } func (t *joinTable) invert() *joinTable { return &joinTable{ table: table{ table: t.table.table, idColumn: t.fkColumn, }, fkColumn: t.table.idColumn, } } func (t *joinTable) get(ctx context.Context, id int) ([]int, error) { q := dialect.Select(t.fkColumn).From(t.table.table).Where(t.idColumn.Eq(id)) const single = false var ret []int if err := queryFunc(ctx, q, single, func(rows *sqlx.Rows) error { var fk int if err := rows.Scan(&fk); err != nil { return err } ret = append(ret, fk) return nil }); err != nil { return nil, fmt.Errorf("getting foreign keys from %s: %w", t.table.table.GetTable(), err) } return ret, nil } func (t *joinTable) insertJoins(ctx context.Context, id int, foreignIDs []int) error { // manually create SQL so that we can prepare once // ignore duplicates q := fmt.Sprintf("INSERT INTO %s (%s, %s) VALUES (?, ?) ON CONFLICT (%[2]s, %s) DO NOTHING", t.table.table.GetTable(), t.idColumn.GetCol(), t.fkColumn.GetCol()) tx := dbWrapper{} stmt, err := tx.Prepare(ctx, q) if err != nil { return err } defer stmt.Close() // eliminate duplicates foreignIDs = intslice.IntAppendUniques(nil, foreignIDs) for _, fk := range foreignIDs { if _, err := tx.ExecStmt(ctx, stmt, id, fk); err != nil { return err } } return nil } func (t *joinTable) replaceJoins(ctx context.Context, id int, foreignIDs []int) error { if err := t.destroy(ctx, []int{id}); err != nil { return err } return t.insertJoins(ctx, id, foreignIDs) } func (t *joinTable) addJoins(ctx context.Context, id int, foreignIDs []int) error { // get existing foreign keys fks, err := t.get(ctx, id) if err != nil { return err } // only add foreign keys that are not already present foreignIDs = intslice.IntExclude(foreignIDs, fks) return t.insertJoins(ctx, id, foreignIDs) } func (t *joinTable) destroyJoins(ctx context.Context, id int, foreignIDs []int) error { q := dialect.Delete(t.table.table).Where( t.idColumn.Eq(id), t.fkColumn.In(foreignIDs), ) if _, err := exec(ctx, q); err != nil { return fmt.Errorf("destroying %s: %w", t.table.table.GetTable(), err) } return nil } func (t *joinTable) modifyJoins(ctx context.Context, id int, foreignIDs []int, mode models.RelationshipUpdateMode) error { switch mode { case models.RelationshipUpdateModeSet: return t.replaceJoins(ctx, id, foreignIDs) case models.RelationshipUpdateModeAdd: return t.addJoins(ctx, id, foreignIDs) case models.RelationshipUpdateModeRemove: return t.destroyJoins(ctx, id, foreignIDs) } return nil } type stashIDTable struct { table } type stashIDRow struct { StashID null.String `db:"stash_id"` Endpoint null.String `db:"endpoint"` } func (r *stashIDRow) resolve() *models.StashID { return &models.StashID{ StashID: r.StashID.String, Endpoint: r.Endpoint.String, } } func (t *stashIDTable) get(ctx context.Context, id int) ([]*models.StashID, error) { q := dialect.Select("endpoint", "stash_id").From(t.table.table).Where(t.idColumn.Eq(id)) const single = false var ret []*models.StashID if err := queryFunc(ctx, q, single, func(rows *sqlx.Rows) error { var v stashIDRow if err := rows.StructScan(&v); err != nil { return err } ret = append(ret, v.resolve()) return nil }); err != nil { return nil, fmt.Errorf("getting stash ids from %s: %w", t.table.table.GetTable(), err) } return ret, nil } func (t *stashIDTable) insertJoin(ctx context.Context, id int, v models.StashID) (sql.Result, error) { q := dialect.Insert(t.table.table).Cols(t.idColumn.GetCol(), "endpoint", "stash_id").Vals( goqu.Vals{id, v.Endpoint, v.StashID}, ) ret, err := exec(ctx, q) if err != nil { return nil, fmt.Errorf("inserting into %s: %w", t.table.table.GetTable(), err) } return ret, nil } func (t *stashIDTable) insertJoins(ctx context.Context, id int, v []models.StashID) error { for _, fk := range v { if _, err := t.insertJoin(ctx, id, fk); err != nil { return err } } return nil } func (t *stashIDTable) replaceJoins(ctx context.Context, id int, v []models.StashID) error { if err := t.destroy(ctx, []int{id}); err != nil { return err } return t.insertJoins(ctx, id, v) } func (t *stashIDTable) addJoins(ctx context.Context, id int, v []models.StashID) error { // get existing foreign keys fks, err := t.get(ctx, id) if err != nil { return err } // only add values that are not already present var filtered []models.StashID for _, vv := range v { for _, e := range fks { if vv.Endpoint == e.Endpoint { continue } filtered = append(filtered, vv) } } return t.insertJoins(ctx, id, filtered) } func (t *stashIDTable) destroyJoins(ctx context.Context, id int, v []models.StashID) error { for _, vv := range v { q := dialect.Delete(t.table.table).Where( t.idColumn.Eq(id), t.table.table.Col("endpoint").Eq(vv.Endpoint), t.table.table.Col("stash_id").Eq(vv.StashID), ) if _, err := exec(ctx, q); err != nil { return fmt.Errorf("destroying %s: %w", t.table.table.GetTable(), err) } } return nil } func (t *stashIDTable) modifyJoins(ctx context.Context, id int, v []models.StashID, mode models.RelationshipUpdateMode) error { switch mode { case models.RelationshipUpdateModeSet: return t.replaceJoins(ctx, id, v) case models.RelationshipUpdateModeAdd: return t.addJoins(ctx, id, v) case models.RelationshipUpdateModeRemove: return t.destroyJoins(ctx, id, v) } return nil } type scenesMoviesTable struct { table } type moviesScenesRow struct { SceneID null.Int `db:"scene_id"` MovieID null.Int `db:"movie_id"` SceneIndex null.Int `db:"scene_index"` } func (r moviesScenesRow) resolve(sceneID int) models.MoviesScenes { return models.MoviesScenes{ MovieID: int(r.MovieID.Int64), SceneIndex: nullIntPtr(r.SceneIndex), } } func (t *scenesMoviesTable) get(ctx context.Context, id int) ([]models.MoviesScenes, error) { q := dialect.Select("movie_id", "scene_index").From(t.table.table).Where(t.idColumn.Eq(id)) const single = false var ret []models.MoviesScenes if err := queryFunc(ctx, q, single, func(rows *sqlx.Rows) error { var v moviesScenesRow if err := rows.StructScan(&v); err != nil { return err } ret = append(ret, v.resolve(id)) return nil }); err != nil { return nil, fmt.Errorf("getting scene movies from %s: %w", t.table.table.GetTable(), err) } return ret, nil } func (t *scenesMoviesTable) insertJoin(ctx context.Context, id int, v models.MoviesScenes) (sql.Result, error) { q := dialect.Insert(t.table.table).Cols(t.idColumn.GetCol(), "movie_id", "scene_index").Vals( goqu.Vals{id, v.MovieID, intFromPtr(v.SceneIndex)}, ) ret, err := exec(ctx, q) if err != nil { return nil, fmt.Errorf("inserting into %s: %w", t.table.table.GetTable(), err) } return ret, nil } func (t *scenesMoviesTable) insertJoins(ctx context.Context, id int, v []models.MoviesScenes) error { for _, fk := range v { if _, err := t.insertJoin(ctx, id, fk); err != nil { return err } } return nil } func (t *scenesMoviesTable) replaceJoins(ctx context.Context, id int, v []models.MoviesScenes) error { if err := t.destroy(ctx, []int{id}); err != nil { return err } return t.insertJoins(ctx, id, v) } func (t *scenesMoviesTable) addJoins(ctx context.Context, id int, v []models.MoviesScenes) error { // get existing foreign keys fks, err := t.get(ctx, id) if err != nil { return err } // only add values that are not already present var filtered []models.MoviesScenes for _, vv := range v { found := false for _, e := range fks { if vv.MovieID == e.MovieID { found = true break } } if !found { filtered = append(filtered, vv) } } return t.insertJoins(ctx, id, filtered) } func (t *scenesMoviesTable) destroyJoins(ctx context.Context, id int, v []models.MoviesScenes) error { for _, vv := range v { q := dialect.Delete(t.table.table).Where( t.idColumn.Eq(id), t.table.table.Col("movie_id").Eq(vv.MovieID), ) if _, err := exec(ctx, q); err != nil { return fmt.Errorf("destroying %s: %w", t.table.table.GetTable(), err) } } return nil } func (t *scenesMoviesTable) modifyJoins(ctx context.Context, id int, v []models.MoviesScenes, mode models.RelationshipUpdateMode) error { switch mode { case models.RelationshipUpdateModeSet: return t.replaceJoins(ctx, id, v) case models.RelationshipUpdateModeAdd: return t.addJoins(ctx, id, v) case models.RelationshipUpdateModeRemove: return t.destroyJoins(ctx, id, v) } return nil } type relatedFilesTable struct { table } // type scenesFilesRow struct { // SceneID int `db:"scene_id"` // Primary bool `db:"primary"` // FileID file.ID `db:"file_id"` // } func (t *relatedFilesTable) insertJoin(ctx context.Context, id int, primary bool, fileID file.ID) error { q := dialect.Insert(t.table.table).Cols(t.idColumn.GetCol(), "primary", "file_id").Vals( goqu.Vals{id, primary, fileID}, ) _, err := exec(ctx, q) if err != nil { return fmt.Errorf("inserting into %s: %w", t.table.table.GetTable(), err) } return nil } func (t *relatedFilesTable) insertJoins(ctx context.Context, id int, firstPrimary bool, fileIDs []file.ID) error { for i, fk := range fileIDs { if err := t.insertJoin(ctx, id, firstPrimary && i == 0, fk); err != nil { return err } } return nil } func (t *relatedFilesTable) replaceJoins(ctx context.Context, id int, fileIDs []file.ID) error { if err := t.destroy(ctx, []int{id}); err != nil { return err } const firstPrimary = true return t.insertJoins(ctx, id, firstPrimary, fileIDs) } func (t *relatedFilesTable) setPrimary(ctx context.Context, id int, fileID file.ID) error { table := t.table.table q := dialect.Update(table).Prepared(true).Set(goqu.Record{ "primary": 0, }).Where(t.idColumn.Eq(id), table.Col(fileIDColumn).Neq(fileID)) if _, err := exec(ctx, q); err != nil { return fmt.Errorf("unsetting primary flags in %s: %w", t.table.table.GetTable(), err) } q = dialect.Update(table).Prepared(true).Set(goqu.Record{ "primary": 1, }).Where(t.idColumn.Eq(id), table.Col(fileIDColumn).Eq(fileID)) if _, err := exec(ctx, q); err != nil { return fmt.Errorf("setting primary flag in %s: %w", t.table.table.GetTable(), err) } return nil } type sqler interface { ToSQL() (sql string, params []interface{}, err error) } func exec(ctx context.Context, stmt sqler) (sql.Result, error) { tx, err := getTx(ctx) if err != nil { return nil, err } sql, args, err := stmt.ToSQL() if err != nil { return nil, fmt.Errorf("generating sql: %w", err) } logger.Tracef("SQL: %s [%v]", sql, args) ret, err := tx.ExecContext(ctx, sql, args...) if err != nil { return nil, fmt.Errorf("executing `%s` [%v]: %w", sql, args, err) } return ret, nil } func count(ctx context.Context, q *goqu.SelectDataset) (int, error) { var count int if err := querySimple(ctx, q, &count); err != nil { return 0, err } return count, nil } func queryFunc(ctx context.Context, query *goqu.SelectDataset, single bool, f func(rows *sqlx.Rows) error) error { q, args, err := query.ToSQL() if err != nil { return err } wrapper := dbWrapper{} rows, err := wrapper.QueryxContext(ctx, q, args...) if err != nil && !errors.Is(err, sql.ErrNoRows) { return fmt.Errorf("querying `%s` [%v]: %w", q, args, err) } defer rows.Close() for rows.Next() { if err := f(rows); err != nil { return err } if single { break } } if err := rows.Err(); err != nil { return err } return nil } func querySimple(ctx context.Context, query *goqu.SelectDataset, out interface{}) error { q, args, err := query.ToSQL() if err != nil { return err } wrapper := dbWrapper{} rows, err := wrapper.QueryxContext(ctx, q, args...) if err != nil { return fmt.Errorf("querying `%s` [%v]: %w", q, args, err) } defer rows.Close() if rows.Next() { if err := rows.Scan(out); err != nil { return err } } if err := rows.Err(); err != nil { return err } return nil } // func cols(table exp.IdentifierExpression, cols []string) []interface{} { // var ret []interface{} // for _, c := range cols { // ret = append(ret, table.Col(c)) // } // return ret // }