add dgraph backend to benchmark suite with safe type assertions for multi-backend support

cmd/benchmark/dgraph_docker.go

@@ -0,0 +1,160 @@
+package main
+
+import (
+	"context"
+	"fmt"
+	"os"
+	"os/exec"
+	"path/filepath"
+	"time"
+)
+
+// DgraphDocker manages a dgraph instance via Docker Compose
+type DgraphDocker struct {
+	composeFile string
+	projectName string
+	running     bool
+}
+
+// NewDgraphDocker creates a new dgraph Docker manager
+func NewDgraphDocker() *DgraphDocker {
+	// Try to find the docker-compose file in the current directory first
+	composeFile := "docker-compose-dgraph.yml"
+
+	// If not found, try the cmd/benchmark directory (for running from project root)
+	if _, err := os.Stat(composeFile); os.IsNotExist(err) {
+		composeFile = filepath.Join("cmd", "benchmark", "docker-compose-dgraph.yml")
+	}
+
+	return &DgraphDocker{
+		composeFile: composeFile,
+		projectName: "orly-benchmark-dgraph",
+		running:     false,
+	}
+}
+
+// Start starts the dgraph Docker containers
+func (d *DgraphDocker) Start(ctx context.Context) error {
+	fmt.Println("Starting dgraph Docker containers...")
+
+	// Stop any existing containers first
+	d.Stop()
+
+	// Start containers
+	cmd := exec.CommandContext(
+		ctx,
+		"docker-compose",
+		"-f", d.composeFile,
+		"-p", d.projectName,
+		"up", "-d",
+	)
+	cmd.Stdout = os.Stdout
+	cmd.Stderr = os.Stderr
+
+	if err := cmd.Run(); err != nil {
+		return fmt.Errorf("failed to start dgraph containers: %w", err)
+	}
+
+	fmt.Println("Waiting for dgraph to be healthy...")
+
+	// Wait for health checks to pass
+	if err := d.waitForHealthy(ctx, 60*time.Second); err != nil {
+		d.Stop() // Clean up on failure
+		return err
+	}
+
+	d.running = true
+	fmt.Println("Dgraph is ready!")
+	return nil
+}
+
+// waitForHealthy waits for dgraph to become healthy
+func (d *DgraphDocker) waitForHealthy(ctx context.Context, timeout time.Duration) error {
+	deadline := time.Now().Add(timeout)
+
+	for time.Now().Before(deadline) {
+		// Check if alpha is healthy by checking docker health status
+		cmd := exec.CommandContext(
+			ctx,
+			"docker",
+			"inspect",
+			"--format={{.State.Health.Status}}",
+			"orly-benchmark-dgraph-alpha",
+		)
+
+		output, err := cmd.Output()
+		if err == nil && string(output) == "healthy\n" {
+			// Additional short wait to ensure full readiness
+			time.Sleep(2 * time.Second)
+			return nil
+		}
+
+		select {
+		case <-ctx.Done():
+			return ctx.Err()
+		case <-time.After(2 * time.Second):
+			// Continue waiting
+		}
+	}
+
+	return fmt.Errorf("dgraph failed to become healthy within %v", timeout)
+}
+
+// Stop stops and removes the dgraph Docker containers
+func (d *DgraphDocker) Stop() error {
+	if !d.running {
+		// Try to stop anyway in case of untracked state
+		cmd := exec.Command(
+			"docker-compose",
+			"-f", d.composeFile,
+			"-p", d.projectName,
+			"down", "-v",
+		)
+		cmd.Stdout = os.Stdout
+		cmd.Stderr = os.Stderr
+		_ = cmd.Run() // Ignore errors
+		return nil
+	}
+
+	fmt.Println("Stopping dgraph Docker containers...")
+
+	cmd := exec.Command(
+		"docker-compose",
+		"-f", d.composeFile,
+		"-p", d.projectName,
+		"down", "-v",
+	)
+	cmd.Stdout = os.Stdout
+	cmd.Stderr = os.Stderr
+
+	if err := cmd.Run(); err != nil {
+		return fmt.Errorf("failed to stop dgraph containers: %w", err)
+	}
+
+	d.running = false
+	fmt.Println("Dgraph containers stopped")
+	return nil
+}
+
+// GetGRPCEndpoint returns the dgraph gRPC endpoint
+func (d *DgraphDocker) GetGRPCEndpoint() string {
+	return "localhost:9080"
+}
+
+// IsRunning returns whether dgraph is running
+func (d *DgraphDocker) IsRunning() bool {
+	return d.running
+}
+
+// Logs returns the logs from dgraph containers
+func (d *DgraphDocker) Logs() error {
+	cmd := exec.Command(
+		"docker-compose",
+		"-f", d.composeFile,
+		"-p", d.projectName,
+		"logs",
+	)
+	cmd.Stdout = os.Stdout
+	cmd.Stderr = os.Stderr
+	return cmd.Run()
+}