mirror of
https://github.com/aykhans/bsky-feedgen.git
synced 2025-07-17 21:34:00 +00:00
🦋
This commit is contained in:
194
pkg/storage/mongodb/collections/feed_az.go
Normal file
194
pkg/storage/mongodb/collections/feed_az.go
Normal file
@@ -0,0 +1,194 @@
|
||||
package collections
|
||||
|
||||
import (
|
||||
"context"
|
||||
"time"
|
||||
|
||||
"github.com/aykhans/bsky-feedgen/pkg/config"
|
||||
"go.mongodb.org/mongo-driver/bson"
|
||||
"go.mongodb.org/mongo-driver/mongo"
|
||||
"go.mongodb.org/mongo-driver/mongo/options"
|
||||
)
|
||||
|
||||
type FeedAzCollection struct {
|
||||
Collection *mongo.Collection
|
||||
}
|
||||
|
||||
func NewFeedAzCollection(client *mongo.Client) (*FeedAzCollection, error) {
|
||||
client.Database(config.MongoDBBaseDB).Collection("")
|
||||
coll := client.Database(config.MongoDBBaseDB).Collection("feed_az")
|
||||
|
||||
_, err := coll.Indexes().CreateMany(
|
||||
context.Background(),
|
||||
[]mongo.IndexModel{
|
||||
{
|
||||
Keys: bson.D{{Key: "sequence", Value: -1}},
|
||||
},
|
||||
{
|
||||
Keys: bson.D{{Key: "created_at", Value: -1}},
|
||||
},
|
||||
},
|
||||
)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
return &FeedAzCollection{Collection: coll}, nil
|
||||
}
|
||||
|
||||
type FeedAz struct {
|
||||
ID string `bson:"_id"`
|
||||
Sequence int64 `bson:"sequence"`
|
||||
DID string `bson:"did"`
|
||||
RecordKey string `bson:"record_key"`
|
||||
CreatedAt time.Time `bson:"created_at"`
|
||||
}
|
||||
|
||||
func (f FeedAzCollection) GetByCreatedAt(ctx context.Context, skip int64, limit int64) ([]*FeedAz, error) {
|
||||
cursor, err := f.Collection.Find(
|
||||
ctx, bson.D{},
|
||||
options.Find().
|
||||
SetSort(bson.D{{Key: "created_at", Value: -1}}).
|
||||
SetSkip(skip).
|
||||
SetLimit(limit),
|
||||
)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
defer func() { _ = cursor.Close(ctx) }()
|
||||
|
||||
var feedAzItems []*FeedAz
|
||||
if err = cursor.All(ctx, &feedAzItems); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
return feedAzItems, nil
|
||||
}
|
||||
|
||||
func (f FeedAzCollection) GetMaxSequence(ctx context.Context) (*int64, error) {
|
||||
pipeline := mongo.Pipeline{
|
||||
{
|
||||
{Key: "$group", Value: bson.D{
|
||||
{Key: "_id", Value: nil},
|
||||
{Key: "maxSequence", Value: bson.D{
|
||||
{Key: "$max", Value: "$sequence"},
|
||||
},
|
||||
},
|
||||
},
|
||||
},
|
||||
},
|
||||
}
|
||||
|
||||
cursor, err := f.Collection.Aggregate(ctx, pipeline)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
defer func() { _ = cursor.Close(ctx) }()
|
||||
|
||||
var result struct {
|
||||
MaxSequence int64 `bson:"maxSequence"`
|
||||
}
|
||||
|
||||
if cursor.Next(ctx) {
|
||||
if err := cursor.Decode(&result); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return &result.MaxSequence, err
|
||||
}
|
||||
|
||||
return nil, nil
|
||||
}
|
||||
|
||||
func (f FeedAzCollection) Insert(ctx context.Context, overwrite bool, feedAz ...*FeedAz) error {
|
||||
switch len(feedAz) {
|
||||
case 0:
|
||||
return nil
|
||||
case 1:
|
||||
if overwrite == false {
|
||||
_, err := f.Collection.InsertOne(ctx, feedAz[0])
|
||||
return err
|
||||
}
|
||||
_, err := f.Collection.ReplaceOne(
|
||||
ctx,
|
||||
bson.M{"_id": feedAz[0].ID},
|
||||
feedAz[0],
|
||||
options.Replace().SetUpsert(true),
|
||||
)
|
||||
return err
|
||||
default:
|
||||
if overwrite == false {
|
||||
documents := make([]any, len(feedAz))
|
||||
for i, feed := range feedAz {
|
||||
documents[i] = feed
|
||||
}
|
||||
|
||||
_, err := f.Collection.InsertMany(ctx, documents)
|
||||
return err
|
||||
}
|
||||
var models []mongo.WriteModel
|
||||
|
||||
for _, feed := range feedAz {
|
||||
filter := bson.M{"_id": feed.ID}
|
||||
model := mongo.NewReplaceOneModel().
|
||||
SetFilter(filter).
|
||||
SetReplacement(feed).
|
||||
SetUpsert(true)
|
||||
models = append(models, model)
|
||||
}
|
||||
|
||||
opts := options.BulkWrite().SetOrdered(false)
|
||||
_, err := f.Collection.BulkWrite(ctx, models, opts)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
return nil
|
||||
}
|
||||
}
|
||||
|
||||
func (f FeedAzCollection) CutoffByCount(
|
||||
ctx context.Context,
|
||||
maxDocumentCount int64,
|
||||
) (int64, error) {
|
||||
count, err := f.Collection.CountDocuments(ctx, bson.M{})
|
||||
if err != nil {
|
||||
return 0, err
|
||||
}
|
||||
|
||||
if count <= maxDocumentCount {
|
||||
return 0, nil
|
||||
}
|
||||
|
||||
deleteCount := count - maxDocumentCount
|
||||
|
||||
findOpts := options.Find().
|
||||
SetSort(bson.D{{Key: "created_at", Value: 1}}).
|
||||
SetLimit(deleteCount)
|
||||
|
||||
cursor, err := f.Collection.Find(ctx, bson.M{}, findOpts)
|
||||
if err != nil {
|
||||
return 0, err
|
||||
}
|
||||
defer func() { _ = cursor.Close(ctx) }()
|
||||
|
||||
var docsToDelete []bson.M
|
||||
if err = cursor.All(ctx, &docsToDelete); err != nil {
|
||||
return 0, err
|
||||
}
|
||||
|
||||
if len(docsToDelete) == 0 {
|
||||
return 0, nil
|
||||
}
|
||||
|
||||
ids := make([]any, len(docsToDelete))
|
||||
for i := range docsToDelete {
|
||||
ids[i] = docsToDelete[i]["_id"]
|
||||
}
|
||||
|
||||
result, err := f.Collection.DeleteMany(ctx, bson.M{"_id": bson.M{"$in": ids}})
|
||||
if err != nil {
|
||||
return 0, err
|
||||
}
|
||||
|
||||
return result.DeletedCount, nil
|
||||
}
|
||||
183
pkg/storage/mongodb/collections/post.go
Normal file
183
pkg/storage/mongodb/collections/post.go
Normal file
@@ -0,0 +1,183 @@
|
||||
package collections
|
||||
|
||||
import (
|
||||
"context"
|
||||
"time"
|
||||
|
||||
"github.com/aykhans/bsky-feedgen/pkg/config"
|
||||
"go.mongodb.org/mongo-driver/bson"
|
||||
"go.mongodb.org/mongo-driver/mongo"
|
||||
"go.mongodb.org/mongo-driver/mongo/options"
|
||||
)
|
||||
|
||||
type PostCollection struct {
|
||||
Collection *mongo.Collection
|
||||
}
|
||||
|
||||
func NewPostCollection(client *mongo.Client) (*PostCollection, error) {
|
||||
client.Database(config.MongoDBBaseDB).Collection("")
|
||||
coll := client.Database(config.MongoDBBaseDB).Collection("post")
|
||||
_, err := coll.Indexes().CreateOne(
|
||||
context.Background(),
|
||||
mongo.IndexModel{
|
||||
Keys: bson.D{{Key: "sequence", Value: -1}},
|
||||
},
|
||||
)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
return &PostCollection{Collection: coll}, nil
|
||||
}
|
||||
|
||||
type Post struct {
|
||||
ID string `bson:"_id"`
|
||||
Sequence int64 `bson:"sequence"`
|
||||
DID string `bson:"did"`
|
||||
RecordKey string `bson:"record_key"`
|
||||
CreatedAt time.Time `bson:"created_at"`
|
||||
Langs []string `bson:"langs"`
|
||||
Tags []string `bson:"tags"`
|
||||
Text string `bson:"text"`
|
||||
Facets *Facets `bson:"facets"`
|
||||
Reply *Reply `bson:"reply"`
|
||||
}
|
||||
|
||||
type Facets struct {
|
||||
Tags []string `bson:"tags"`
|
||||
Links []string `bson:"links"`
|
||||
Mentions []string `bson:"mentions"`
|
||||
}
|
||||
|
||||
type Reply struct {
|
||||
RootURI string `bson:"root_uri"`
|
||||
ParentURI string `bson:"parent_uri"`
|
||||
}
|
||||
|
||||
func (p PostCollection) CutoffByCount(
|
||||
ctx context.Context,
|
||||
maxDocumentCount int64,
|
||||
) (int64, error) {
|
||||
count, err := p.Collection.CountDocuments(ctx, bson.M{})
|
||||
if err != nil {
|
||||
return 0, err
|
||||
}
|
||||
|
||||
if count <= maxDocumentCount {
|
||||
return 0, nil
|
||||
}
|
||||
|
||||
deleteCount := count - maxDocumentCount
|
||||
|
||||
findOpts := options.Find().
|
||||
SetSort(bson.D{{Key: "created_at", Value: 1}}).
|
||||
SetLimit(deleteCount)
|
||||
|
||||
cursor, err := p.Collection.Find(ctx, bson.M{}, findOpts)
|
||||
if err != nil {
|
||||
return 0, err
|
||||
}
|
||||
defer func() { _ = cursor.Close(ctx) }()
|
||||
|
||||
var docsToDelete []bson.M
|
||||
if err = cursor.All(ctx, &docsToDelete); err != nil {
|
||||
return 0, err
|
||||
}
|
||||
|
||||
if len(docsToDelete) == 0 {
|
||||
return 0, nil
|
||||
}
|
||||
|
||||
ids := make([]any, len(docsToDelete))
|
||||
for i := range docsToDelete {
|
||||
ids[i] = docsToDelete[i]["_id"]
|
||||
}
|
||||
|
||||
result, err := p.Collection.DeleteMany(ctx, bson.M{"_id": bson.M{"$in": ids}})
|
||||
if err != nil {
|
||||
return 0, err
|
||||
}
|
||||
|
||||
return result.DeletedCount, nil
|
||||
}
|
||||
|
||||
func (p PostCollection) GetMaxSequence(ctx context.Context) (*int64, error) {
|
||||
pipeline := mongo.Pipeline{
|
||||
{
|
||||
{Key: "$group", Value: bson.D{
|
||||
{Key: "_id", Value: nil},
|
||||
{Key: "maxSequence", Value: bson.D{
|
||||
{Key: "$max", Value: "$sequence"},
|
||||
},
|
||||
},
|
||||
},
|
||||
},
|
||||
},
|
||||
}
|
||||
|
||||
cursor, err := p.Collection.Aggregate(ctx, pipeline)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
defer func() { _ = cursor.Close(ctx) }()
|
||||
|
||||
var result struct {
|
||||
MaxSequence int64 `bson:"maxSequence"`
|
||||
}
|
||||
|
||||
if cursor.Next(ctx) {
|
||||
if err := cursor.Decode(&result); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return &result.MaxSequence, err
|
||||
}
|
||||
|
||||
return nil, nil
|
||||
}
|
||||
|
||||
func (p PostCollection) Insert(ctx context.Context, overwrite bool, posts ...*Post) error {
|
||||
switch len(posts) {
|
||||
case 0:
|
||||
return nil
|
||||
case 1:
|
||||
if overwrite == false {
|
||||
_, err := p.Collection.InsertOne(ctx, posts[0])
|
||||
return err
|
||||
}
|
||||
_, err := p.Collection.ReplaceOne(
|
||||
ctx,
|
||||
bson.M{"_id": posts[0].ID},
|
||||
posts[0],
|
||||
options.Replace().SetUpsert(true),
|
||||
)
|
||||
return err
|
||||
default:
|
||||
if overwrite == false {
|
||||
documents := make([]any, len(posts))
|
||||
for i, post := range posts {
|
||||
documents[i] = post
|
||||
}
|
||||
|
||||
_, err := p.Collection.InsertMany(ctx, documents)
|
||||
return err
|
||||
}
|
||||
var models []mongo.WriteModel
|
||||
|
||||
for _, post := range posts {
|
||||
filter := bson.M{"_id": post.ID}
|
||||
model := mongo.NewReplaceOneModel().
|
||||
SetFilter(filter).
|
||||
SetReplacement(post).
|
||||
SetUpsert(true)
|
||||
models = append(models, model)
|
||||
}
|
||||
|
||||
opts := options.BulkWrite().SetOrdered(false)
|
||||
_, err := p.Collection.BulkWrite(ctx, models, opts)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
return nil
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user