104 lines
3.0 KiB
Go
104 lines
3.0 KiB
Go
package storage
|
|
|
|
import (
|
|
"context"
|
|
"fmt"
|
|
|
|
"cloud.google.com/go/firestore"
|
|
"git.pengzhan.dev/aimaren/internal/crawler"
|
|
"google.golang.org/grpc/codes"
|
|
"google.golang.org/grpc/status"
|
|
)
|
|
|
|
const (
|
|
StateCollection = "hermes_state"
|
|
AppStateDocument = "main"
|
|
UserStateDocument = "users"
|
|
)
|
|
|
|
// FirestoreClient implements Storer
|
|
type FirestoreClient struct {
|
|
client *firestore.Client
|
|
}
|
|
|
|
func NewFirestoreClient(ctx context.Context, projectID string) (*FirestoreClient, error) {
|
|
client, err := firestore.NewClient(ctx, projectID)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
return &FirestoreClient{client: client}, nil
|
|
}
|
|
|
|
func (fs *FirestoreClient) Close() {
|
|
fs.client.Close()
|
|
}
|
|
|
|
// FetchAppState retrieves the entire application state from a single document.
|
|
func (fs *FirestoreClient) FetchAppState(ctx context.Context) (*AppState, error) {
|
|
doc, err := fs.client.Collection(StateCollection).Doc(AppStateDocument).Get(ctx)
|
|
if err != nil {
|
|
if status.Code(err) == codes.NotFound {
|
|
return &AppState{
|
|
Bags: make(map[string]crawler.Bag),
|
|
ChatIDs: []int64{},
|
|
}, nil
|
|
}
|
|
return nil, err
|
|
}
|
|
var state AppState
|
|
if err := doc.DataTo(&state); err != nil {
|
|
return nil, err
|
|
}
|
|
if state.Bags == nil {
|
|
state.Bags = make(map[string]crawler.Bag)
|
|
}
|
|
if state.ChatIDs == nil {
|
|
state.ChatIDs = []int64{}
|
|
}
|
|
return &state, nil
|
|
}
|
|
|
|
// UpdateAppState writes the entire application state back to the document.
|
|
func (fs *FirestoreClient) UpdateAppState(ctx context.Context, newState *AppState) error {
|
|
_, err := fs.client.Collection(StateCollection).Doc(AppStateDocument).Set(ctx, newState)
|
|
return err
|
|
}
|
|
|
|
// AddChatID atomically adds a new chat ID to the list in the main state document.
|
|
func (fs *FirestoreClient) AddChatID(ctx context.Context, chatID int64) error {
|
|
docRef := fs.client.Collection(StateCollection).Doc(AppStateDocument)
|
|
_, err := docRef.Update(ctx, []firestore.Update{
|
|
{Path: "chat_ids", Value: firestore.ArrayUnion(chatID)},
|
|
})
|
|
if status.Code(err) == codes.NotFound {
|
|
return fs.UpdateAppState(ctx, &AppState{
|
|
Bags: make(map[string]crawler.Bag),
|
|
ChatIDs: []int64{chatID},
|
|
})
|
|
}
|
|
return err
|
|
}
|
|
|
|
// FetchUserState retrieves the user state document and converts string keys back to int64.
|
|
func (fs *FirestoreClient) FetchUserState(ctx context.Context) (*UserState, error) {
|
|
docSnap, err := fs.client.Collection(StateCollection).Doc(UserStateDocument).Get(ctx)
|
|
if err != nil {
|
|
if status.Code(err) == codes.NotFound {
|
|
return &UserState{Users: make(map[string]ChatState)}, nil
|
|
}
|
|
return nil, err
|
|
}
|
|
var result UserState
|
|
if err := docSnap.DataTo(&result); err != nil {
|
|
return nil, fmt.Errorf("failed to decode user state: %w", err)
|
|
}
|
|
return &result, nil
|
|
}
|
|
|
|
// UpdateUserState writes the whole UserState back to Firestore.
|
|
// It converts int64 keys to string keys for Firestore compatibility.
|
|
func (fs *FirestoreClient) UpdateUserState(ctx context.Context, newUserState *UserState) error {
|
|
_, err := fs.client.Collection(StateCollection).Doc(UserStateDocument).Set(ctx, newUserState)
|
|
return err
|
|
}
|