Skip to content

Instantly share code, notes, and snippets.

@Hasstrup
Last active May 24, 2019 07:12
Show Gist options
  • Save Hasstrup/2629f0023afdd8045a1712d46ca578c0 to your computer and use it in GitHub Desktop.
Save Hasstrup/2629f0023afdd8045a1712d46ca578c0 to your computer and use it in GitHub Desktop.
package main
import (
"context"
"flag"
"fmt"
"log"
"time"
"github.com/mongodb/mongo-go-driver/bson"
"github.com/mongodb/mongo-go-driver/mongo/options"
"github.com/mongodb/mongo-go-driver/x/bsonx"
"github.com/mongodb/mongo-go-driver/mongo"
)
var (
cmd = flag.String("cmd", "", "list or add?")
address = flag.String("address", "", "mongodb address to connect to")
database = flag.String("db", "", "The name of the database to connect to")
collection = flag.String("collection", "", "The collection (in the db) to connect to")
key = flag.String("field", "", "The field you'd like to place an index on")
unique = flag.Bool("unique", false, "Would you like the index to be unique?")
value = flag.Int("type", 1, "would you like the index to be ascending (1) or descending (-1)?")
)
func main() {
flag.Parse()
switch {
case (*cmd != "add" && *cmd != "list"):
log.Fatalf("The first argument has to be 'add' or 'list :)")
case *database == "" || *address == "":
log.Fatalf("Please provide a valid database address and database name :)")
case *cmd == "add" && *key == "":
log.Fatalf("Please pass in the name of the field to place the index :)")
}
client := ConnectToTheMongoDB(*address)
if *cmd == "add" {
PopulateIndex(*database, *collection, client)
} else if *cmd == "list" {
ListIndexes(client, *database, *collection)
}
}
func ConnectToTheMongoDB(address string) *mongo.Client {
client, err := mongo.Connect(context.Background(), address)
if err != nil {
panic(err.Error())
}
log.Println("Successfully connected to the address provided")
return client
}
func PopulateIndex(database, collection string, client *mongo.Client) {
c := client.Database(database).Collection(collection)
opts := options.CreateIndexes().SetMaxTime(10 * time.Second)
index := yieldIndexModel()
c.Indexes().CreateOne(context.Background(), index, opts)
log.Println("Successfully create the index")
}
func yieldIndexModel() mongo.IndexModel {
keys := bsonx.Doc{{Key: *key, Value: bsonx.Int32(int32(*value))}}
index := mongo.IndexModel{}
index.Keys = keys
if *unique {
index.Options = bsonx.Doc{{Key: "unique", Value: bsonx.Boolean(true)}}
}
return index
}
func ListIndexes(client *mongo.Client, database, collection string) {
c := client.Database(database).Collection(collection)
duration := 10 * time.Second
batchSize := int32(10)
cur, err := c.Indexes().List(context.Background(), &options.ListIndexesOptions{&batchSize, &duration})
if err != nil {
log.Fatalf("Something went wrong listing %v", err)
}
for cur.Next(context.Background()) {
index := bson.D{}
cur.Decode(&index)
log.Println(fmt.Sprintf("index found %v", index))
}
}
@like-a-bause
Copy link

Here is an updates Version for mongo driver 1.0.2:

package main

import (
	"context"
	"flag"
	"fmt"
	"go.mongodb.org/mongo-driver/bson"
	"go.mongodb.org/mongo-driver/mongo"
	"go.mongodb.org/mongo-driver/mongo/options"
	"go.mongodb.org/mongo-driver/x/bsonx"
	"log"
	"time"
)

var (
	cmd        = flag.String("cmd", "", "list or add?")
	address    = flag.String("address", "", "mongodb address to connect to")
	database   = flag.String("db", "", "The name of the database to connect to")
	collection = flag.String("collection", "", "The collection (in the db) to connect to")
	key        = flag.String("field", "", "The field you'd like to place an index on")
	unique     = flag.Bool("unique", false, "Would you like the index to be unique?")
	value      = flag.Int("type", 1, "would you like the index to be ascending (1) or descending (-1)?")
)

func main() {
	flag.Parse()
	switch {
	case (*cmd != "add" && *cmd != "list"):
		log.Fatalf("The first argument has to be 'add' or 'list :)")
	case *database == "" || *address == "" || *collection == "":
		log.Fatalf("Please provide a valid database address, database name and collection name. :)")
	case *cmd == "add" && *key == "":
		log.Fatalf("Please pass in the name of the field to place the index :)")
	}
	client := ConnectToTheMongoDB(*address)
	if *cmd == "add" {
		PopulateIndex(*database, *collection, client)
	} else if *cmd == "list" {
		ListIndexes(client, *database, *collection)
	}
}

func ConnectToTheMongoDB(address string) *mongo.Client {
	client, err := mongo.Connect(context.Background(), options.Client().ApplyURI(address))
	if err != nil {
		panic(err.Error())
	}
	log.Println("Successfully connected to the address provided")
	return client
}

func PopulateIndex(database, collection string, client *mongo.Client) {
	c := client.Database(database).Collection(collection)
	opts := options.CreateIndexes().SetMaxTime(10 * time.Second)
	index := yieldIndexModel()
	c.Indexes().CreateOne(context.Background(), index, opts)
	log.Println("Successfully create the index")
}

func yieldIndexModel() mongo.IndexModel {
	keys := bsonx.Doc{{Key: *key, Value: bsonx.Int32(int32(*value))}}
	index := mongo.IndexModel{}
	index.Keys = keys
	if *unique {
		index.Options = options.Index().SetUnique(true)
	}
	return index
}

func ListIndexes(client *mongo.Client, database, collection string) {
	c := client.Database(database).Collection(collection)
	duration := 10 * time.Second
	batchSize := int32(10)
	cur, err := c.Indexes().List(context.Background(), &options.ListIndexesOptions{&batchSize, &duration})
	if err != nil {
		log.Fatalf("Something went wrong listing %v", err)
	}
	for cur.Next(context.Background()) {
		index := bson.D{}
		cur.Decode(&index)
		log.Println(fmt.Sprintf("index found %v", index))
	}
}

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment