<commit_before>package db

import (
	"fmt"

	"gopkg.in/mgo.v2"
	"gopkg.in/mgo.v2/bson"
)

//MongodbConfig DB配置
type MongodbConfig struct {
	URL    string
	DBName string
}

//MongoDataInterface 存储的MongoData
type MongoDataInterface interface {
	GenObjectID() bson.ObjectId
	GetObjectID() bson.ObjectId
}

var mgoSession *mgo.Session
var dbConfig *MongodbConfig

//ConnectMongoDB 连接MongoDB
func ConnectMongoDB(url string, dbName string) error {
	dbConfig = &MongodbConfig{}
	dbConfig.URL = url
	dbConfig.DBName = dbName
	var err error
	mgoSession, err = mgo.Dial(url)

	if err != nil {
		return err
	}
	return nil
}

//GetMongoSession 获取session， 必须手动Close
func getMongoSession() *mgo.Session {
	if mgoSession == nil {
		return nil
	}
	//最大连接池默认为4096
	return mgoSession.Clone()
}

//QueryWithCollection 获取collection对象
func QueryWithCollection(collection string, s func(*mgo.Collection) error) error {
	session := getMongoSession()
	defer session.Close()
	c := session.DB(dbConfig.DBName).C(collection)
	return s(c)
}

//GetNextSeq 获取自增id
func GetNextSeq(num int) int {
	doc := struct{ Seq int }{}
	cid := "counterid"
	query := func(c *mgo.Collection) error {
		change := mgo.Change{
			Update:    bson.M{"$inc": bson.M{"seq": num}},
			Upsert:    true,
			ReturnNew: true,
		}
		_, err := c.Find(bson.M{"_id": cid}).Apply(change, &doc)
		if err != nil {
			panic(fmt.Errorf("get counter failed:", err.Error()))
		}
		return err
	}

	err := QueryWithCollection("counter", query)
	if err != nil {
		fmt.Println(err.Error())
		return -1
	}
	return doc.Seq
}<commit_msg>feat(mongo):add set seq doc<commit_after>package db

import (
	"fmt"

	"gopkg.in/mgo.v2"
	"gopkg.in/mgo.v2/bson"
)

//MongodbConfig DB配置
type MongodbConfig struct {
	URL    string
	DBName string
}

//MongoDataInterface 存储的MongoData
type MongoDataInterface interface {
	GenObjectID() bson.ObjectId
	GetObjectID() bson.ObjectId
}

var mgoSession *mgo.Session
var dbConfig *MongodbConfig

//ConnectMongoDB 连接MongoDB
func ConnectMongoDB(url string, dbName string) error {
	dbConfig = &MongodbConfig{}
	dbConfig.URL = url
	dbConfig.DBName = dbName
	var err error
	mgoSession, err = mgo.Dial(url)

	if err != nil {
		return err
	}
	return nil
}

//GetMongoSession 获取session， 必须手动Close
func getMongoSession() *mgo.Session {
	if mgoSession == nil {
		return nil
	}
	//最大连接池默认为4096
	return mgoSession.Clone()
}

//QueryWithCollection 获取collection对象
func QueryWithCollection(collection string, s func(*mgo.Collection) error) error {
	session := getMongoSession()
	defer session.Close()
	c := session.DB(dbConfig.DBName).C(collection)
	return s(c)
}

func SetNextSeqBaseNum(collection string, key string, base int) int {
	doc := struct{ Seq int }{}
	query := func(c *mgo.Collection) error {
		change := mgo.Change{
			Update: bson.M{
				"$setOnInsert": bson.M{"seq": base},
			},
			Upsert:    true,
			ReturnNew: true,
		}
		_, err := c.Find(bson.M{"_id": key}).Apply(change, &doc)
		if err != nil {
			panic(fmt.Errorf("set counter failed:", err.Error()))
		}
		return err
	}
	err := QueryWithCollection(collection, query)
	if err != nil {
		return -1
	}
	return doc.Seq
}

//GetNextSeq 获取自增id
func GetNextSeq(collection string, key string, num int) int {
	doc := struct{ Seq int }{}
	cid := key
	query := func(c *mgo.Collection) error {
		change := mgo.Change{
			Update:    bson.M{"$inc": bson.M{"seq": num}},
			Upsert:    true,
			ReturnNew: true,
		}
		_, err := c.Find(bson.M{"_id": cid}).Apply(change, &doc)
		if err != nil {
			panic(fmt.Errorf("get counter failed:", err.Error()))
		}
		return err
	}

	err := QueryWithCollection(collection, query)
	if err != nil {
		fmt.Println(err.Error())
		return -1
	}
	return doc.Seq
}
{'fraction_non_alphanumeric': 0.136117136659436, 'fraction_numerical': 0.004067245119305857, 'mean_word_length': 3.9055851063829787, 'pattern_counts': {'":': 9, '<': 3, '<?xml version=': 0, '>': 3, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '6665972', 'n_tokens_mistral': 1490, 'n_tokens_neox': 1447, 'n_words': 333}