timetable-mgr/dbAccess/corpus.go

package dbAccess

import (
	"context"
	"strings"

	"git.fjla.uk/owlboard/go-types/pkg/database"
	"git.fjla.uk/owlboard/timetable-mgr/log"
	"go.mongodb.org/mongo-driver/bson"
	"go.mongodb.org/mongo-driver/mongo"
	"go.mongodb.org/mongo-driver/mongo/options"
	"go.uber.org/zap"
)

// Puts an array of Corpus Documents into the database
func PutManyCorpus(corpusData *[]database.CorpusEntry) error {
	collection := MongoClient.Database(DatabaseName).Collection(CorpusCollection)

	documents := convertCorpusToInterfaceSlice(corpusData)

	_, err := collection.InsertMany(context.Background(), *documents)
	if err != nil {
		return err
	}

	SetUpdateTime(CorpusCollection)
	return nil
}

// Puts an array of Stations documents into the database
func PutManyStations(stationsData *[]database.StationEntry) error {
	collection := MongoClient.Database(DatabaseName).Collection(StationsCollection)

	documents := convertStationsToInterfaceSlice(stationsData)

	_, err := collection.InsertMany(context.Background(), *documents)
	if err != nil {
		return err
	}

	SetUpdateTime(StationsCollection)
	return nil
}

// Converts []database.CorpusEntry types into interface slices required to put them into the database
func convertCorpusToInterfaceSlice(corpusData *[]database.CorpusEntry) *[]interface{} {
	interfaceSlice := make([]interface{}, len(*corpusData))
	for i, doc := range *corpusData {
		interfaceSlice[i] = doc
	}
	return &interfaceSlice
}

// Converts []database.StationEntry types into interface slices required to put them into the database
func convertStationsToInterfaceSlice(stationsData *[]database.StationEntry) *[]interface{} {
	interfaceSlice := make([]interface{}, len(*stationsData))
	for i, doc := range *stationsData {
		interfaceSlice[i] = doc
	}
	return &interfaceSlice
}

func CreateCorpusIndexes() error {
	coll := MongoClient.Database(DatabaseName).Collection(CorpusCollection)
	indexModels := []mongo.IndexModel{
		{
			Keys: bson.M{
				"tiploc": 1,
			},
			Options: options.Index().SetName("tiploc"),
		},
		{
			Keys: bson.M{
				"3alpha": 1,
			},
			Options: options.Index().SetName("3alpha"),
		},
	}

	_, err := coll.Indexes().CreateMany(context.Background(), indexModels)
	if err != nil {
		return err
	}
	return nil
}

func GetTiplocFromCrs(crs string) (tiploc string, err error) {
	// Return TIPLOC from CRS code

	// TEMP FIX: WPH does not return a tiploc
	if crs == "WPH" {
		return "WORCPHL", nil
	}

	crs = strings.ToUpper(crs)
	// PIPELINE:
	pipeline := bson.A{
		bson.D{{"$match", bson.D{{"3ALPHA", crs}}}},
		bson.D{
			{"$project",
				bson.D{
					{"TIPLOC", 1},
					{"_id", 0},
				},
			},
		},
	}

	coll := MongoClient.Database(DatabaseName).Collection(CorpusCollection)
	cursor, err := coll.Aggregate(context.Background(), pipeline)
	if err != nil {
		return "", err
	}

	defer cursor.Close(context.Background())

	var result struct {
		TIPLOC string `bson:"TIPLOC"`
	}
	if cursor.Next(context.Background()) {
		if err := cursor.Decode(&result); err != nil {
			return "", err
		}
		return result.TIPLOC, nil
	}

	log.Warn("No TIPLOC Found", zap.String("CRS", crs))
	return "", nil
}

func GetStanoxFromCrs(crs string) (stanox string, err error) {
	// Return STANOX from CRS code
	crs = strings.ToUpper(crs)
	// PIPELINE:
	pipeline := bson.A{
		bson.D{{"$match", bson.D{{"3ALPHA", crs}}}},
		bson.D{
			{"$project",
				bson.D{
					{"STANOX", 1},
					{"_id", 0},
				},
			},
		},
	}

	coll := MongoClient.Database(DatabaseName).Collection(CorpusCollection)
	cursor, err := coll.Aggregate(context.Background(), pipeline)
	if err != nil {
		return "", err
	}

	defer cursor.Close(context.Background())

	var result struct {
		STANOX string `bson:"STANOX"`
	}
	if cursor.Next(context.Background()) {
		if err := cursor.Decode(&result); err != nil {
			return "", err
		}
		return result.STANOX, nil
	}

	log.Warn("No STANOX Found", zap.String("CRS", crs))
	return "", nil
}