mirror of
https://github.com/ent/ent.git
synced 2026-05-22 09:31:45 +03:00
123 lines
3.5 KiB
Go
123 lines
3.5 KiB
Go
// Copyright 2019-present Facebook Inc. All rights reserved.
|
|
// This source code is licensed under the Apache 2.0 license found
|
|
// in the LICENSE file in the root directory of this source tree.
|
|
|
|
package ent
|
|
|
|
import (
|
|
"context"
|
|
"errors"
|
|
"io"
|
|
"io/fs"
|
|
)
|
|
|
|
// Blob defines the interface for blob storage operations.
|
|
// Implementations should return [io/fs.ErrNotExist] (or an error wrapping it)
|
|
// from NewReader when the requested key does not exist.
|
|
type Blob interface {
|
|
// NewReader opens a reader for the given key.
|
|
NewReader(ctx context.Context, key string) (io.ReadCloser, error)
|
|
// NewWriter opens a writer for the given key.
|
|
NewWriter(ctx context.Context, key string) (io.WriteCloser, error)
|
|
// Close releases any resources held by the bucket.
|
|
Close() error
|
|
}
|
|
|
|
// BlobReader returns a reader for the given key from the blob bucket.
|
|
// The returned reader closes both the underlying reader and the bucket.
|
|
// Returns nil, nil if the blob does not exist (fs.ErrNotExist).
|
|
func BlobReader(ctx context.Context, b Blob, key string) (io.ReadCloser, error) {
|
|
switch r, err := b.NewReader(ctx, key); {
|
|
case errors.Is(err, fs.ErrNotExist):
|
|
return nil, b.Close()
|
|
case err != nil:
|
|
return nil, errors.Join(err, b.Close())
|
|
default:
|
|
return &blobReadCloser{ReadCloser: r, bucket: b}, nil
|
|
}
|
|
}
|
|
|
|
// blobReadCloser wraps an io.ReadCloser to also close the bucket on Close.
|
|
type blobReadCloser struct {
|
|
io.ReadCloser
|
|
bucket Blob
|
|
}
|
|
|
|
func (r *blobReadCloser) Close() error {
|
|
return errors.Join(r.ReadCloser.Close(), r.bucket.Close())
|
|
}
|
|
|
|
// BlobWriter returns a writer for the given key in the blob bucket.
|
|
// The returned writer closes both the underlying writer and the bucket.
|
|
func BlobWriter(ctx context.Context, b Blob, key string) (io.WriteCloser, error) {
|
|
w, err := b.NewWriter(ctx, key)
|
|
if err != nil {
|
|
return nil, errors.Join(err, b.Close())
|
|
}
|
|
return &blobWriteCloser{WriteCloser: w, bucket: b}, nil
|
|
}
|
|
|
|
// blobWriteCloser wraps an io.WriteCloser to also close the bucket on Close.
|
|
type blobWriteCloser struct {
|
|
io.WriteCloser
|
|
bucket Blob
|
|
}
|
|
|
|
func (w *blobWriteCloser) Close() error {
|
|
return errors.Join(w.WriteCloser.Close(), w.bucket.Close())
|
|
}
|
|
|
|
// BlobBulkWriter manages blob bucket lifecycles for write operations.
|
|
// It lazily opens buckets per field and provides a Close method
|
|
// to release all resources when done.
|
|
type BlobBulkWriter struct {
|
|
opener func(context.Context, string) (Blob, error)
|
|
buckets map[string]Blob
|
|
}
|
|
|
|
// NewBlobBulkWriter creates a writer that uses opener to lazily open buckets.
|
|
func NewBlobBulkWriter(opener func(context.Context, string) (Blob, error)) *BlobBulkWriter {
|
|
return &BlobBulkWriter{
|
|
buckets: make(map[string]Blob),
|
|
opener: opener,
|
|
}
|
|
}
|
|
|
|
// Close closes all open buckets.
|
|
func (w *BlobBulkWriter) Close() error {
|
|
var errs []error
|
|
for _, b := range w.buckets {
|
|
errs = append(errs, b.Close())
|
|
}
|
|
return errors.Join(errs...)
|
|
}
|
|
|
|
// Write writes r to the blob at key for the given field. The bucket is opened
|
|
// lazily on first use and reused for subsequent writes to the same field.
|
|
func (w *BlobBulkWriter) Write(ctx context.Context, field, key string, r io.Reader) error {
|
|
b, err := w.bucket(ctx, field)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
wr, err := b.NewWriter(ctx, key)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
if _, err := io.Copy(wr, r); err != nil {
|
|
return errors.Join(err, wr.Close())
|
|
}
|
|
return wr.Close()
|
|
}
|
|
|
|
func (w *BlobBulkWriter) bucket(ctx context.Context, field string) (Blob, error) {
|
|
if b, ok := w.buckets[field]; ok {
|
|
return b, nil
|
|
}
|
|
b, err := w.opener(ctx, field)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
w.buckets[field] = b
|
|
return b, nil
|
|
}
|