diff --git a/docker/scripts/init-shard.go b/docker/scripts/init-shard.go new file mode 100755 index 000000000..96435feb4 --- /dev/null +++ b/docker/scripts/init-shard.go @@ -0,0 +1,91 @@ +package main + +import ( + "fmt" + "net" + "os" + "os/exec" + "time" +) + +func resolveHost(hostname string) (string, error) { + addrs, err := net.LookupHost(hostname) + if err != nil { + fmt.Println("Error resolving hostname:", err) + return "", err + } + return addrs[0], nil +} + +func main() { + mongos := os.Getenv("MONGOS") + mongo11 := os.Getenv("MONGO11") + mongo12 := os.Getenv("MONGO12") + mongo13 := os.Getenv("MONGO13") + + mongo21 := os.Getenv("MONGO21") + mongo22 := os.Getenv("MONGO22") + mongo23 := os.Getenv("MONGO23") + + mongo31 := os.Getenv("MONGO31") + mongo32 := os.Getenv("MONGO32") + mongo33 := os.Getenv("MONGO33") + + port := os.Getenv("PORT") + if port == "" { + port = "27017" + } + + fmt.Println("Waiting for startup..") + for { + _, err := net.DialTimeout("tcp", net.JoinHostPort(mongos, port), 1*time.Second) + if err == nil { + break + } + fmt.Print(".") + time.Sleep(1 * time.Second) + } + fmt.Println("Started..") + + fmt.Printf("init-shard.sh time now: %s\n", time.Now().Format("15:04:05")) + + rs1 := os.Getenv("RS1") + rs2 := os.Getenv("RS2") + + mongodb11Resolved, _ := resolveHost(mongo11) + mongodb12Resolved, _ := resolveHost(mongo12) + mongodb13Resolved, _ := resolveHost(mongo13) + + mongodb21Resolved, _ := resolveHost(mongo21) + mongodb22Resolved, _ := resolveHost(mongo22) + mongodb23Resolved, _ := resolveHost(mongo23) + + mongodb31Resolved, _ := resolveHost(mongo31) + mongodb32Resolved, _ := resolveHost(mongo32) + mongodb33Resolved, _ := resolveHost(mongo33) + + cmd := exec.Command("mongo", "--host", fmt.Sprintf("%s:%s", mongos, port)) + cmd.Stdin = os.Stdin + cmd.Stdout = os.Stdout + cmd.Stderr = os.Stderr + + input := fmt.Sprintf(` + sh.addShard("%s/%s:%s,%s:%s,%s:%s"); + sh.addShard("%s/%s:%s,%s:%s,%s:%s"); + sh.status(); + quit(); + `, rs1, mongodb11Resolved, port, mongodb12Resolved, port, mongodb13Resolved, port, rs2, mongodb21Resolved, port, mongodb22Resolved, port, mongodb23Resolved, port) + + cmd.Stdin = os.Stdin + cmd.Stdout = os.Stdout + cmd.Stderr = os.Stderr + + err := cmd.Start() + if err != nil { + fmt.Println("Error:", err) + return + } + cmd.Stdin.Write([]byte(input)) + cmd.Wait() +} + diff --git a/docker/scripts/setup.go b/docker/scripts/setup.go new file mode 100755 index 000000000..33f89c692 --- /dev/null +++ b/docker/scripts/setup.go @@ -0,0 +1,109 @@ +package main + +import ( + "context" + "fmt" + "os" + "time" + + "go.mongodb.org/mongo-driver/mongo" + "go.mongodb.org/mongo-driver/mongo/options" +) + +func waitForStartup(client *mongo.Client, ctx context.Context) error { + for { + err := client.Ping(ctx, nil) + if err == nil { + break + } + fmt.Print(".") + time.Sleep(1 * time.Second) + } + fmt.Println("Started..") + return nil +} + +func cnfServers(client *mongo.Client, ctx context.Context, configServers []string) error { + fmt.Println("Setting up config servers") + rsConfig := mongo.Cmd{ + {"replSetInitiate", map[string]interface{}{ + "_id": "configReplSet", + "configsvr": true, + "members": configServers, + "protocolVersion": 1, + }}, + } + + _, err := client.Database("admin").RunCommand(ctx, rsConfig).DecodeBytes() + return err +} + +func generalServers(client *mongo.Client, ctx context.Context, replicaSet []string, arbiter string) error { + fmt.Println("Setting up general servers") + rsConfig := mongo.Cmd{ + {"replSetInitiate", map[string]interface{}{ + "_id": "rs0", + "members": replicaSet, + "protocolVersion": 1, + }}, + } + + _, err := client.Database("admin").RunCommand(ctx, rsConfig).DecodeBytes() + if err != nil { + return err + } + + arbCmd := fmt.Sprintf("rs.addArb(\"%s\")", arbiter) + _, err = client.Database("admin").RunCommand(ctx, mongo.Cmd{"eval", arbCmd}).DecodeBytes() + + return err +} + +func main() { + + ctx := context.TODO() + + clientOptions := options.Client().ApplyURI("mongodb://localhost:27017") + client, err := mongo.Connect(ctx, clientOptions) + if err != nil { + fmt.Println("Error connecting to MongoDB:", err) + os.Exit(1) + } + defer client.Disconnect(ctx) + + if len(os.Args) < 2 { + fmt.Println("Usage: go run main.go ") + return + } + + mongo1 := os.Getenv("MONGO1") + mongo2 := os.Getenv("MONGO2") + mongo3 := os.Getenv("MONGO3") + arbiter := os.Getenv("ARBITER") + + configServers := []string{mongo1, mongo2, mongo3} + replicaSet := []string{mongo1, mongo2, mongo3} + + err = waitForStartup(client, ctx) + if err != nil { + fmt.Println("Error waiting for startup:", err) + os.Exit(1) + } + + + + switch os.Args[1] { + case "cnf_servers": + err = cnfServers(client, ctx, configServers) + case "general_servers": + err = generalServers(client, ctx, replicaSet, arbiter) + default: + err = generalServers(client, ctx, replicaSet, arbiter) + } + + if err != nil { + fmt.Println("Error setting up MongoDB replica set:", err) + os.Exit(1) + } +} +