diff --git a/internal/fstore/fstore.go b/internal/fstore/fstore.go new file mode 100644 index 0000000..1c779e4 --- /dev/null +++ b/internal/fstore/fstore.go @@ -0,0 +1,77 @@ +// Copyright 2023 The Go Authors. All rights reserved. +// Use of this source code is governed by a BSD-style +// license that can be found in the LICENSE file. + +// Package fstore provides general support for Firestore. +// Its main feature is separate namespaces, to mimic separate +// databases for different purposes (prod, dev, test, etc.). +package fstore + +import ( + "context" + "errors" + + "cloud.google.com/go/firestore" + "golang.org/x/pkgsite-metrics/internal/derrors" +) + +const namespaceCollection = "Namespaces" + +// A Namespace is a top-level collection for partitioning a Firestore +// database into separate segments. +type Namespace struct { + client *firestore.Client + name string + doc *firestore.DocumentRef +} + +// OpenNamespace creates a new Firestore client whose collections will be located in the given namespace. +func OpenNamespace(ctx context.Context, projectID, name string) (_ *Namespace, err error) { + defer derrors.Wrap(&err, "OpenNamespace(%q, %q)", projectID, name) + + if name == "" { + return nil, errors.New("empty namespace") + } + client, err := firestore.NewClient(ctx, projectID) + if err != nil { + return nil, err + } + return &Namespace{ + client: client, + name: name, + doc: client.Collection(namespaceCollection).Doc(name), + }, nil +} + +// Name returns the Namespace's name. +func (ns *Namespace) Name() string { return ns.name } + +// Client returns the underlying Firestore client. +func (ns *Namespace) Client() *firestore.Client { return ns.client } + +// Close closes the underlying client. +func (ns *Namespace) Close() error { return ns.client.Close() } + +// Collection returns a reference to the named collection in the namespace. +func (ns *Namespace) Collection(name string) *firestore.CollectionRef { + return ns.doc.Collection(name) +} + +// Get gets the DocumentRef and decodes the result to a value of type T. +func Get[T any](ctx context.Context, dr *firestore.DocumentRef) (_ *T, err error) { + defer derrors.Wrap(&err, "fstore.Get(%q)", dr.Path) + docsnap, err := dr.Get(ctx) + if err != nil { + return nil, err + } + return Decode[T](docsnap) +} + +// Decode decodes a DocumentSnapshot into a value of type T. +func Decode[T any](ds *firestore.DocumentSnapshot) (*T, error) { + var t T + if err := ds.DataTo(&t); err != nil { + return nil, err + } + return &t, nil +} diff --git a/internal/jobs/firestore.go b/internal/jobs/firestore.go index 86194a9..8c3d647 100644 --- a/internal/jobs/firestore.go +++ b/internal/jobs/firestore.go @@ -6,42 +6,27 @@ package jobs import ( "context" - "errors" "time" "cloud.google.com/go/firestore" "golang.org/x/pkgsite-metrics/internal/derrors" + "golang.org/x/pkgsite-metrics/internal/fstore" "google.golang.org/api/iterator" ) -// A DB is a client for a database that stores Jobs. +const jobCollection = "Jobs" + type DB struct { - namespace string - client *firestore.Client - nsDoc *firestore.DocumentRef // the namespace for this db + ns *fstore.Namespace } -const ( - namespaceCollection = "Namespaces" - jobCollection = "Jobs" -) - // NewDB creates a new database client for jobs. func NewDB(ctx context.Context, projectID, namespace string) (_ *DB, err error) { - defer derrors.Wrap(&err, "job.NewDB(%q, %q)", projectID, namespace) - - if namespace == "" { - return nil, errors.New("empty namespace") - } - client, err := firestore.NewClient(ctx, projectID) + ns, err := fstore.OpenNamespace(ctx, projectID, namespace) if err != nil { return nil, err } - return &DB{ - namespace: namespace, - client: client, - nsDoc: client.Collection(namespaceCollection).Doc(namespace), - }, nil + return &DB{ns}, nil } // CreateJob creates a new job. It returns an error if a job with the same ID already exists. @@ -62,11 +47,7 @@ func (d *DB) DeleteJob(ctx context.Context, id string) (err error) { // GetJob retrieves the job with the given ID. It returns an error if the job does not exist. func (d *DB) GetJob(ctx context.Context, id string) (_ *Job, err error) { defer derrors.Wrap(&err, "job.DB.GetJob(%s)", id) - docsnap, err := d.jobRef(id).Get(ctx) - if err != nil { - return nil, err - } - return docsnapToJob(docsnap) + return fstore.Get[Job](ctx, d.jobRef(id)) } // UpdateJob gets the job with the given ID, which must exist, then calls f on @@ -74,13 +55,13 @@ func (d *DB) GetJob(ctx context.Context, id string) (_ *Job, err error) { // If f returns an error, that error is returned and no update occurs. func (d *DB) UpdateJob(ctx context.Context, id string, f func(*Job) error) (err error) { defer derrors.Wrap(&err, "job.DB.UpdateJob(%s)", id) - return d.client.RunTransaction(ctx, func(ctx context.Context, tx *firestore.Transaction) error { + return d.ns.Client().RunTransaction(ctx, func(ctx context.Context, tx *firestore.Transaction) error { docref := d.jobRef(id) docsnap, err := tx.Get(docref) if err != nil { return err } - j, err := docsnapToJob(docsnap) + j, err := fstore.Decode[Job](docsnap) if err != nil { return err } @@ -108,7 +89,7 @@ func (d *DB) Increment(ctx context.Context, id, name string, n int) (err error) func (d *DB) ListJobs(ctx context.Context, f func(_ *Job, lastUpdate time.Time) error) (err error) { defer derrors.Wrap(&err, "job.DB.ListJobs()") - q := d.nsDoc.Collection(jobCollection).OrderBy("StartedAt", firestore.Desc) + q := d.ns.Collection(jobCollection).OrderBy("StartedAt", firestore.Desc) iter := q.Documents(ctx) defer iter.Stop() for { @@ -119,7 +100,7 @@ func (d *DB) ListJobs(ctx context.Context, f func(_ *Job, lastUpdate time.Time) if err != nil { return err } - job, err := docsnapToJob(docsnap) + job, err := fstore.Decode[Job](docsnap) if err != nil { return err } @@ -132,14 +113,5 @@ func (d *DB) ListJobs(ctx context.Context, f func(_ *Job, lastUpdate time.Time) // jobRef returns the DocumentRef for a job with the given ID. func (d *DB) jobRef(id string) *firestore.DocumentRef { - return d.nsDoc.Collection(jobCollection).Doc(id) -} - -// docsnapToJob converts a DocumentSnapshot to a Job. -func docsnapToJob(ds *firestore.DocumentSnapshot) (*Job, error) { - var j Job - if err := ds.DataTo(&j); err != nil { - return nil, err - } - return &j, nil + return d.ns.Collection(jobCollection).Doc(id) }