1// db holds the functions relating to the database
 2package main
 3
 4import (
 5	"gopkg.in/mgo.v2"
 6	"gopkg.in/mgo.v2/bson"
 7	"strconv"
 8)
 9
10// DB_CONFIG defines the interactions within the database
11type DB_CONFIG struct {
12	// DB_name is the name of the database
13	DB_name,
14	// CollectionName is the name of the collection within the database
15	CollectionName string
16}
17
18// DB_config stores the configuration data for the database
19var DB_config = DB_CONFIG{
20	DB_name:        "records",
21	CollectionName: "cyber"}
22
23// DBsession contains the session information for the connection
24// to the database
25var DBsession *mgo.Session
26
27// getCategoriesFromDB
28func getCategoriesFromDB() ([]string, error) {
29	if DBsession == nil {
30		if _, err := InitDBSession(); err != nil {
31			return nil, err
32		}
33	}
34	var dbCats []string
35	collection := DBsession.DB(DB_config.DB_name).C(DB_config.CollectionName)
36	err := collection.Find(bson.M{"path": bson.M{"$regex": ""}}).Distinct("path", &dbCats)
37	return dbCats, err
38}
39
40// initDBSession initializes the connection to the database containing
41// the list of records.
42// returns the *mgo.Session connected to the database
43func InitDBSession() (*mgo.Session, error) {
44	var err error
45	if serverConfig.DB_PORT != 0 {
46		DBsession, err = mgo.Dial(serverConfig.DB_ADDRESS + ":" + strconv.Itoa(serverConfig.DB_PORT))
47	}
48	if err != nil {
49		return nil, err
50	}
51	index := mgo.Index{Key: []string{"question"},
52		DropDups: true,
53		Unique:   true}
54	DBsession.DB(DB_config.DB_name).C(DB_config.CollectionName).EnsureIndex(index)
55	categories, err = getCategoriesFromDB()
56	if err != nil {
57		return nil, err
58	}
59	return DBsession, nil
60}
61
62// returnForQuery
63func returnForQuery(selectedCategories []string) ([]Record, error) {
64	if DBsession == nil {
65		if _, err := InitDBSession(); err != nil {
66			return nil, err
67		}
68	}
69	collection := DBsession.DB(DB_config.DB_name).C(DB_config.CollectionName)
70	var results []Record
71	var collector []Record
72	var err error
73	for _, category := range selectedCategories {
74		err = collection.Find(bson.M{"path": bson.M{"$regex": category}}).All(&collector)
75		if err != nil {
76			return nil, err
77		} else {
78			results = append(results, collector...)
79		}
80	}
81	return results, nil
82}
83
84// addRecordToDB
85func addRecordToDB(record *Record) error {
86	if DBsession == nil {
87		if _, err := InitDBSession(); err != nil {
88			return err
89		}
90	}
91	err := DBsession.DB(DB_config.DB_name).C(DB_config.CollectionName).Insert(record)
92	return err
93}