<commit_before>package database

import (
	"github.com/spf13/viper"
	"gopkg.in/mgo.v2"
)

type Database struct {
	session  *mgo.Session
	Updates  *mgo.Collection
	Vehicles *mgo.Collection
	Routes   *mgo.Collection
	Stops    *mgo.Collection
	Users    *mgo.Collection
}

type Config struct {
	MongoURL string
}

func New(cfg Config) (*Database, error) {
	db := &Database{}

	session, err := mgo.Dial(cfg.MongoURL)
	if err != nil {
		return nil, err
	}
	db.session = session

	db.Updates = db.session.DB("").C("updates")
	db.Vehicles = db.session.DB("").C("vehicles")
	db.Routes = db.session.DB("").C("routes")
	db.Stops = db.session.DB("").C("stops")
	db.Users = db.session.DB("").C("users")

	// Ensure unique vehicle identification
	vehicleIndex := mgo.Index{
		Key:      []string{"vehicleID"},
		Unique:   true,
		DropDups: true}
	db.Vehicles.EnsureIndex(vehicleIndex)

	// Create index on update create time to quickly find the most recent updates
	db.Updates.EnsureIndexKey("created")

	// Index on update vehicle ID to quickly find updates for specific vehicles.
	db.Updates.EnsureIndexKey("vehicleID")

	return db, nil
}

func NewConfig(v *viper.Viper) *Config {
	cfg := &Config{
		MongoURL: "localhost:27017",
	}
	v.SetDefault("database.mongourl", cfg.MongoURL)
	return cfg
}<commit_msg>Add updates vehicle ID and created time, enabled vehicles indices<commit_after>package database

import (
	"github.com/spf13/viper"
	"gopkg.in/mgo.v2"
)

type Database struct {
	session  *mgo.Session
	Updates  *mgo.Collection
	Vehicles *mgo.Collection
	Routes   *mgo.Collection
	Stops    *mgo.Collection
	Users    *mgo.Collection
}

type Config struct {
	MongoURL string
}

func New(cfg Config) (*Database, error) {
	db := &Database{}

	session, err := mgo.Dial(cfg.MongoURL)
	if err != nil {
		return nil, err
	}
	db.session = session

	db.Updates = db.session.DB("").C("updates")
	db.Vehicles = db.session.DB("").C("vehicles")
	db.Routes = db.session.DB("").C("routes")
	db.Stops = db.session.DB("").C("stops")
	db.Users = db.session.DB("").C("users")

	// Ensure unique vehicle identification
	vehicleIndex := mgo.Index{
		Key:      []string{"vehicleID"},
		Unique:   true,
		DropDups: true}
	db.Vehicles.EnsureIndex(vehicleIndex)

	// Create index on update vehicle ID and creation time to quickly find the most recent updates for specific vehicles.
	db.Updates.EnsureIndexKey("created")
	db.Updates.EnsureIndexKey("vehicleID")
	db.Updates.EnsureIndexKey("vehicleID", "created")

	// Index on enabled vehicles
	db.Vehicles.EnsureIndexKey("enabled")

	return db, nil
}

func NewConfig(v *viper.Viper) *Config {
	cfg := &Config{
		MongoURL: "localhost:27017",
	}
	v.SetDefault("database.mongourl", cfg.MongoURL)
	return cfg
}
{'fraction_non_alphanumeric': 0.14411764705882352, 'fraction_numerical': 0.0058823529411764705, 'mean_word_length': 4.143667296786389, 'pattern_counts': {'":': 0, '<': 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': '6571294', 'n_tokens_mistral': 1023, 'n_tokens_neox': 994, 'n_words': 226}