From 04aeda21e0ad2f7e8dd2bad3328e6ce0ba38f6a9 Mon Sep 17 00:00:00 2001
From: Niall Sheridan <nsheridan@gmail.com>
Date: Fri, 29 Jul 2016 00:59:48 +0100
Subject: Support mongo datastores

---
 cmd/dbinit/dbinit.go | 60 ++++++++++++++++++++++++++++++++++++++++++++++------
 1 file changed, 53 insertions(+), 7 deletions(-)

(limited to 'cmd/dbinit/dbinit.go')

diff --git a/cmd/dbinit/dbinit.go b/cmd/dbinit/dbinit.go
index de06cbb..46484d1 100644
--- a/cmd/dbinit/dbinit.go
+++ b/cmd/dbinit/dbinit.go
@@ -7,6 +7,8 @@ import (
 	"log"
 	"strings"
 
+	mgo "gopkg.in/mgo.v2"
+
 	"github.com/go-sql-driver/mysql"
 )
 
@@ -14,16 +16,19 @@ var (
 	host        = flag.String("host", "localhost", "host[:port]")
 	adminUser   = flag.String("admin_user", "root", "Admin user")
 	adminPasswd = flag.String("admin_password", "", "Admin password")
-	dbUser      = flag.String("db_user", "root", "Database user")
-	dbPasswd    = flag.String("db_password", "", "Admin password")
+	dbUser      = flag.String("db_user", "user", "Database user")
+	dbPasswd    = flag.String("db_password", "passwd", "Admin password")
+	dbType      = flag.String("db_type", "mysql", "Database engine (\"mysql\" or \"mongo\")")
+
+	certsDB     = "certs"
+	issuedTable = "issued_certs"
 )
 
-func main() {
-	flag.Parse()
+func initMySQL() {
 	var createTableStmt = []string{
-		`CREATE DATABASE IF NOT EXISTS certs DEFAULT CHARACTER SET = 'utf8' DEFAULT COLLATE 'utf8_general_ci';`,
-		`USE certs;`,
-		`CREATE TABLE IF NOT EXISTS issued_certs (
+		`CREATE DATABASE IF NOT EXISTS ` + certsDB + ` DEFAULT CHARACTER SET = 'utf8' DEFAULT COLLATE 'utf8_general_ci';`,
+		`USE ` + certsDB + `;`,
+		`CREATE TABLE IF NOT EXISTS ` + issuedTable + ` (
 			key_id VARCHAR(255) NOT NULL,
 			principals VARCHAR(255) NULL,
 			created_at DATETIME NULL,
@@ -59,3 +64,44 @@ func main() {
 		}
 	}
 }
+
+func initMongo() {
+	di := &mgo.DialInfo{
+		Addrs:    strings.Split(*host, ","),
+		Username: *adminUser,
+		Password: *adminPasswd,
+	}
+	session, err := mgo.DialWithInfo(di)
+	if err != nil {
+		log.Fatalln(err)
+	}
+	defer session.Close()
+	d := session.DB(certsDB)
+	if err := d.UpsertUser(&mgo.User{
+		Username: *dbUser,
+		Password: *dbPasswd,
+		Roles:    []mgo.Role{mgo.RoleReadWrite},
+	}); err != nil {
+		log.Fatalln(err)
+	}
+	c := d.C(issuedTable)
+	i := mgo.Index{
+		Key:    []string{"keyid"},
+		Unique: true,
+	}
+	if err != c.EnsureIndex(i) {
+		log.Fatalln(err)
+	}
+}
+
+func main() {
+	flag.Parse()
+	switch *dbType {
+	case "mysql":
+		initMySQL()
+	case "mongo":
+		initMongo()
+	default:
+		log.Fatalf("Invalid database type")
+	}
+}
-- 
cgit v1.2.3