feat: add scoped assets and storage core
This commit is contained in:
zn-admin committed
1 parent
5e60bb40e7
commit
7847c95538
10 files changed
+1345
No files matched your search
@@ -0,0 +1,193 @@
|
||||
package postgres
|
||||
|
||||
import (
|
||||
"context"
|
||||
"database/sql"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"os"
|
||||
"time"
|
||||
|
||||
"git.nianxx.cn/wangxuming/NianAIGC/backend/internal/assets"
|
||||
)
|
||||
|
||||
const assetFields = `a.id,
|
||||
a.owner_id,
|
||||
a.kind,
|
||||
a.name,
|
||||
a.url,
|
||||
a.storage_path,
|
||||
a.source,
|
||||
a.tags,
|
||||
a.metadata,
|
||||
a.created_at,
|
||||
a.updated_at`
|
||||
|
||||
const ListOwnerAssetsSQL = `SELECT ` + assetFields + `
|
||||
FROM public.assets AS a
|
||||
WHERE a.owner_id = $1::text
|
||||
ORDER BY a.created_at DESC`
|
||||
const GetOwnerAssetSQL = `SELECT ` + assetFields + `
|
||||
FROM public.assets AS a
|
||||
WHERE a.owner_id = $1::text AND a.id = $2::text
|
||||
LIMIT 1`
|
||||
const listPublicAccessibleAssetIDs = `SELECT unnest(j.input_asset_ids || j.output_asset_ids) AS asset_id
|
||||
FROM (
|
||||
SELECT input_asset_ids, output_asset_ids
|
||||
FROM public.generation_jobs
|
||||
WHERE owner_id = $1::text AND external_client_id = $3::text
|
||||
ORDER BY created_at DESC
|
||||
LIMIT $4::integer
|
||||
) AS j`
|
||||
const ListPublicAssetsSQL = `SELECT ` + assetFields + `
|
||||
FROM public.assets AS a
|
||||
WHERE a.owner_id = $1::text
|
||||
AND ($2::text = ANY(a.tags) OR a.id IN (` + listPublicAccessibleAssetIDs + `))
|
||||
ORDER BY a.created_at DESC`
|
||||
const getPublicAccessibleAssetIDs = `SELECT unnest(j.input_asset_ids || j.output_asset_ids) AS asset_id
|
||||
FROM (
|
||||
SELECT input_asset_ids, output_asset_ids
|
||||
FROM public.generation_jobs
|
||||
WHERE owner_id = $1::text AND external_client_id = $4::text
|
||||
ORDER BY created_at DESC
|
||||
LIMIT $5::integer
|
||||
) AS j`
|
||||
const GetPublicAssetSQL = `SELECT ` + assetFields + `
|
||||
FROM public.assets AS a
|
||||
WHERE a.owner_id = $1::text
|
||||
AND a.id = $2::text
|
||||
AND ($3::text = ANY(a.tags) OR a.id IN (` + getPublicAccessibleAssetIDs + `))
|
||||
LIMIT 1`
|
||||
const CreateAssetSQL = `INSERT INTO public.assets (
|
||||
id, owner_id, kind, name, url, storage_path, source, tags, metadata, created_at, updated_at
|
||||
) VALUES ($1::text, $2::text, $3::text, $4::text, $5::text, $6::text, $7::text, $8::text[], $9::jsonb, $10::timestamptz, $11::timestamptz)
|
||||
RETURNING id, owner_id, kind, name, url, storage_path, source, tags, metadata, created_at, updated_at`
|
||||
const DeleteOwnerAssetSQL = `DELETE FROM public.assets
|
||||
WHERE owner_id = $1::text AND id = $2::text
|
||||
RETURNING id, owner_id, kind, name, url, storage_path, source, tags, metadata, created_at, updated_at`
|
||||
|
||||
func (db *Database) ListOwner(ctx context.Context, owner string) ([]assets.Asset, error) {
|
||||
return db.listAssets(ctx, ListOwnerAssetsSQL, owner)
|
||||
}
|
||||
func (db *Database) GetOwner(ctx context.Context, owner, id string) (assets.Asset, bool, error) {
|
||||
return db.oneAsset(ctx, GetOwnerAssetSQL, owner, id)
|
||||
}
|
||||
func (db *Database) ListPublic(ctx context.Context, owner, client string, limit int) ([]assets.Asset, error) {
|
||||
return db.listAssets(ctx, ListPublicAssetsSQL, owner, assets.ClientTag(client), client, limit)
|
||||
}
|
||||
func (db *Database) GetPublic(ctx context.Context, owner, client, id string, limit int) (assets.Asset, bool, error) {
|
||||
return db.oneAsset(ctx, GetPublicAssetSQL, owner, id, assets.ClientTag(client), client, limit)
|
||||
}
|
||||
func (db *Database) Create(ctx context.Context, a assets.Asset) (assets.Asset, error) {
|
||||
if err := db.available(); err != nil {
|
||||
return assets.Asset{}, err
|
||||
}
|
||||
if a.Tags == nil {
|
||||
a.Tags = []string{}
|
||||
}
|
||||
if a.Metadata == nil {
|
||||
a.Metadata = map[string]any{}
|
||||
}
|
||||
metadata, err := json.Marshal(a.Metadata)
|
||||
if err != nil {
|
||||
return assets.Asset{}, fmt.Errorf("encode asset metadata: %w", err)
|
||||
}
|
||||
return db.requiredAsset(ctx, CreateAssetSQL, a.ID, a.OwnerID, string(a.Kind), a.Name, a.URL, optionalDatabaseText(a.StoragePath), string(a.Source), a.Tags, metadata, a.CreatedAt, a.UpdatedAt)
|
||||
}
|
||||
func (db *Database) DeleteOwner(ctx context.Context, owner, id string) (assets.Asset, bool, error) {
|
||||
return db.oneAsset(ctx, DeleteOwnerAssetSQL, owner, id)
|
||||
}
|
||||
func (db *Database) available() error {
|
||||
if db.config.Backend != BackendPostgres || db.querier == nil {
|
||||
return fmt.Errorf("PostgreSQL is unavailable when ZHINIAN_DATA_BACKEND=%s", db.config.Backend)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
func (db *Database) listAssets(ctx context.Context, query string, args ...any) ([]assets.Asset, error) {
|
||||
if err := db.available(); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
rows, err := db.querier.Query(ctx, query, args...)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("query assets: %w", err)
|
||||
}
|
||||
defer rows.Close()
|
||||
out := []assets.Asset{}
|
||||
for rows.Next() {
|
||||
a, err := scanAsset(rows)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
out = append(out, a)
|
||||
}
|
||||
if err = rows.Err(); err != nil {
|
||||
return nil, fmt.Errorf("read assets: %w", err)
|
||||
}
|
||||
return out, nil
|
||||
}
|
||||
func (db *Database) oneAsset(ctx context.Context, query string, args ...any) (assets.Asset, bool, error) {
|
||||
if err := db.available(); err != nil {
|
||||
return assets.Asset{}, false, err
|
||||
}
|
||||
rows, err := db.querier.Query(ctx, query, args...)
|
||||
if err != nil {
|
||||
return assets.Asset{}, false, fmt.Errorf("query asset: %w", err)
|
||||
}
|
||||
defer rows.Close()
|
||||
if !rows.Next() {
|
||||
if err = rows.Err(); err != nil {
|
||||
return assets.Asset{}, false, fmt.Errorf("read asset: %w", err)
|
||||
}
|
||||
return assets.Asset{}, false, nil
|
||||
}
|
||||
a, err := scanAsset(rows)
|
||||
if err != nil {
|
||||
return assets.Asset{}, false, err
|
||||
}
|
||||
if err = rows.Err(); err != nil {
|
||||
return assets.Asset{}, false, fmt.Errorf("read asset: %w", err)
|
||||
}
|
||||
return a, true, nil
|
||||
}
|
||||
func (db *Database) requiredAsset(ctx context.Context, query string, args ...any) (assets.Asset, error) {
|
||||
a, found, err := db.oneAsset(ctx, query, args...)
|
||||
if err != nil {
|
||||
return assets.Asset{}, err
|
||||
}
|
||||
if !found {
|
||||
return assets.Asset{}, fmt.Errorf("asset write returned no row")
|
||||
}
|
||||
return a, nil
|
||||
}
|
||||
func scanAsset(rows Rows) (assets.Asset, error) {
|
||||
var a assets.Asset
|
||||
var kind, source string
|
||||
var storage sql.NullString
|
||||
var metadata []byte
|
||||
var created, updated time.Time
|
||||
if err := rows.Scan(&a.ID, &a.OwnerID, &kind, &a.Name, &a.URL, &storage, &source, &a.Tags, &metadata, &created, &updated); err != nil {
|
||||
return assets.Asset{}, fmt.Errorf("scan asset: %w", err)
|
||||
}
|
||||
a.Kind = assets.Kind(kind)
|
||||
a.Source = assets.Source(source)
|
||||
if storage.Valid {
|
||||
a.StoragePath = storage.String
|
||||
}
|
||||
a.CreatedAt = created
|
||||
a.UpdatedAt = updated
|
||||
if len(metadata) == 0 {
|
||||
metadata = []byte(`{}`)
|
||||
}
|
||||
if err := json.Unmarshal(metadata, &a.Metadata); err != nil {
|
||||
return assets.Asset{}, fmt.Errorf("decode asset metadata: %w", err)
|
||||
}
|
||||
if a.Tags == nil {
|
||||
a.Tags = []string{}
|
||||
}
|
||||
if a.Metadata == nil {
|
||||
a.Metadata = map[string]any{}
|
||||
}
|
||||
return a, nil
|
||||
}
|
||||
|
||||
func readContractFixture(path string) ([]byte, error) { return os.ReadFile(path) }
|
||||
@@ -0,0 +1,172 @@
|
||||
package postgres
|
||||
|
||||
import (
|
||||
"context"
|
||||
"database/sql"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"reflect"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"git.nianxx.cn/wangxuming/NianAIGC/backend/internal/assets"
|
||||
)
|
||||
|
||||
func TestAssetCatalogUsesExplicitOwnerScopedQueries(t *testing.T) {
|
||||
now := time.Date(2026, 8, 13, 8, 0, 0, 0, time.UTC)
|
||||
row := []any{"asset-1", "owner-1", "image", "poster.png", "https://cdn.test/a.png", "uploads/a.png", "upload", []string{"upload"}, []byte(`{"size":4}`), now, now}
|
||||
tests := []struct {
|
||||
name string
|
||||
call func(*Database) (assets.Asset, bool, error)
|
||||
wantSQL string
|
||||
wantArgs []any
|
||||
}{
|
||||
{name: "get owner", call: func(db *Database) (assets.Asset, bool, error) {
|
||||
return db.GetOwner(context.Background(), "owner-1", "asset-1")
|
||||
}, wantSQL: GetOwnerAssetSQL, wantArgs: []any{"owner-1", "asset-1"}},
|
||||
{name: "get public", call: func(db *Database) (assets.Asset, bool, error) {
|
||||
return db.GetPublic(context.Background(), "api:agent-a", "agent-a", "asset-1", 200)
|
||||
}, wantSQL: GetPublicAssetSQL, wantArgs: []any{"api:agent-a", "asset-1", "api-client:agent-a", "agent-a", 200}},
|
||||
}
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
q := &assetQuerier{rows: &assetRows{rows: [][]any{row}}}
|
||||
db := NewDatabase(Config{Backend: BackendPostgres}, q)
|
||||
got, found, err := tt.call(db)
|
||||
if err != nil || !found {
|
||||
t.Fatalf("result = %#v,%v,%v", got, found, err)
|
||||
}
|
||||
if q.sql != tt.wantSQL || !reflect.DeepEqual(q.args, tt.wantArgs) {
|
||||
t.Fatalf("query=%q args=%#v", q.sql, q.args)
|
||||
}
|
||||
if got.Metadata["size"] != float64(4) || got.StoragePath != "uploads/a.png" {
|
||||
t.Fatalf("asset=%#v", got)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestAssetCatalogCreateDefaultsCollectionsAndUsesAllFields(t *testing.T) {
|
||||
now := time.Date(2026, 8, 13, 8, 0, 0, 0, time.UTC)
|
||||
q := &assetQuerier{rows: &assetRows{rows: [][]any{{"a", "o", "image", "n", "u", nil, "external", []string{}, []byte(`{}`), now, now}}}}
|
||||
db := NewDatabase(Config{Backend: BackendPostgres}, q)
|
||||
got, err := db.Create(context.Background(), assets.Asset{ID: "a", OwnerID: "o", Kind: assets.KindImage, Name: "n", URL: "u", Source: assets.SourceExternal, CreatedAt: now, UpdatedAt: now})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if q.sql != CreateAssetSQL || got.Tags == nil || got.Metadata == nil {
|
||||
t.Fatalf("query=%q asset=%#v", q.sql, got)
|
||||
}
|
||||
if !reflect.DeepEqual(q.args[7], []string{}) || string(q.args[8].([]byte)) != "{}" {
|
||||
t.Fatalf("args=%#v", q.args)
|
||||
}
|
||||
}
|
||||
|
||||
func TestAssetCatalogDeleteIsOwnerScopedAndNotFound(t *testing.T) {
|
||||
q := &assetQuerier{rows: &assetRows{}}
|
||||
db := NewDatabase(Config{Backend: BackendPostgres}, q)
|
||||
_, found, err := db.DeleteOwner(context.Background(), "owner", "asset")
|
||||
if err != nil || found {
|
||||
t.Fatalf("found=%v err=%v", found, err)
|
||||
}
|
||||
if q.sql != DeleteOwnerAssetSQL || !reflect.DeepEqual(q.args, []any{"owner", "asset"}) {
|
||||
t.Fatalf("query=%q args=%#v", q.sql, q.args)
|
||||
}
|
||||
}
|
||||
|
||||
func TestAssetCatalogFailsClosedWithoutPostgres(t *testing.T) {
|
||||
db := NewDatabase(Config{Backend: BackendLocal}, nil)
|
||||
if _, err := db.ListOwner(context.Background(), "owner"); err == nil {
|
||||
t.Fatal("expected error")
|
||||
}
|
||||
}
|
||||
|
||||
var _ assets.Catalog = (*Database)(nil)
|
||||
|
||||
type assetQuerier struct {
|
||||
rows *assetRows
|
||||
err error
|
||||
sql string
|
||||
args []any
|
||||
}
|
||||
|
||||
func (q *assetQuerier) Query(_ context.Context, sql string, args ...any) (Rows, error) {
|
||||
q.sql = sql
|
||||
q.args = args
|
||||
return q.rows, q.err
|
||||
}
|
||||
|
||||
type assetRows struct {
|
||||
rows [][]any
|
||||
idx int
|
||||
err error
|
||||
}
|
||||
|
||||
func (r *assetRows) Close() {}
|
||||
func (r *assetRows) Err() error { return r.err }
|
||||
func (r *assetRows) Next() bool { return r.idx < len(r.rows) }
|
||||
func (r *assetRows) Scan(dest ...any) error {
|
||||
if r.idx >= len(r.rows) {
|
||||
return errors.New("past end")
|
||||
}
|
||||
row := r.rows[r.idx]
|
||||
r.idx++
|
||||
if len(row) != len(dest) {
|
||||
return errors.New("arity")
|
||||
}
|
||||
for i, v := range row {
|
||||
switch d := dest[i].(type) {
|
||||
case *string:
|
||||
if v == nil {
|
||||
*d = ""
|
||||
} else {
|
||||
*d = v.(string)
|
||||
}
|
||||
case *sql.NullString:
|
||||
if v == nil {
|
||||
*d = sql.NullString{}
|
||||
} else {
|
||||
*d = sql.NullString{String: v.(string), Valid: true}
|
||||
}
|
||||
case *[]string:
|
||||
*d = append([]string(nil), v.([]string)...)
|
||||
case *[]byte:
|
||||
*d = append([]byte(nil), v.([]byte)...)
|
||||
case *time.Time:
|
||||
*d = v.(time.Time)
|
||||
default:
|
||||
return errors.New("unsupported")
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func TestAssetWireContractFixtureDecodes(t *testing.T) {
|
||||
raw, err := readContractFixture("../../../contracts/assets/assets-v1.json")
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
var fixture struct {
|
||||
Asset assets.Asset `json:"asset"`
|
||||
PublicVisibility struct {
|
||||
JobLookupLimit int `json:"jobLookupLimit"`
|
||||
} `json:"publicVisibility"`
|
||||
}
|
||||
if err = json.Unmarshal(raw, &fixture); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if fixture.Asset.ID != "asset-contract-1" || fixture.PublicVisibility.JobLookupLimit != assets.PublicJobLookupLimit {
|
||||
t.Fatalf("fixture=%#v", fixture)
|
||||
}
|
||||
wire, err := json.Marshal(fixture.Asset)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
var encoded map[string]any
|
||||
if err := json.Unmarshal(wire, &encoded); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if encoded["createdAt"] != "2026-08-13T08:00:00.000Z" || encoded["updatedAt"] != "2026-08-13T08:00:00.000Z" {
|
||||
t.Fatalf("wire timestamps = %q, %q", encoded["createdAt"], encoded["updatedAt"])
|
||||
}
|
||||
}
|
||||
Reference in new issue
Block a user