Refactor guacgql command to support out-of-tree backends

Signed-off-by: robert-cronin <[email protected]>

cmd/guacgql/cmd/root.go

@@ -20,6 +20,7 @@ import (
 	"os"
 	"strings"
 
+	"github.com/guacsec/guac/pkg/assembler/backends"
 	"github.com/guacsec/guac/pkg/cli"
 	"github.com/guacsec/guac/pkg/version"
 	"github.com/spf13/cobra"
@@ -34,35 +35,6 @@ var flags = struct {
 	tlsKeyFile  string
 	debug       bool
 	tracegql    bool
-
-	// Needed only if using neo4j backend
-	nAddr  string
-	nUser  string
-	nPass  string
-	nRealm string
-
-	// Needed only if using ent backend
-	dbAddress  string
-	dbDriver   string
-	dbDebug    bool
-	dbMigrate  bool
-	dbConnTime string
-
-	// Needed only if using arangodb backend
-	arangoAddr string
-	arangoUser string
-	arangoPass string
-
-	// Needed only if using neptune backend
-	neptuneEndpoint string
-	neptunePort     int
-	neptuneRegion   string
-	neptuneUser     string
-	neptuneRealm    string
-
-	kvStore string
-	kvRedis string
-	kvTiKV  string
 }{}
 
 var rootCmd = &cobra.Command{
@@ -77,51 +49,32 @@ var rootCmd = &cobra.Command{
 		flags.debug = viper.GetBool("gql-debug")
 		flags.tracegql = viper.GetBool("gql-trace")
 
-		flags.nUser = viper.GetString("neo4j-user")
-		flags.nPass = viper.GetString("neo4j-pass")
-		flags.nAddr = viper.GetString("neo4j-addr")
-		flags.nRealm = viper.GetString("neo4j-realm")
-
-		// Needed only if using ent backend
-		flags.dbAddress = viper.GetString("db-address")
-		flags.dbDriver = viper.GetString("db-driver")
-		flags.dbDebug = viper.GetBool("db-debug")
-		flags.dbMigrate = viper.GetBool("db-migrate")
-		flags.dbConnTime = viper.GetString("db-conn-time")
-
-		flags.arangoUser = viper.GetString("arango-user")
-		flags.arangoPass = viper.GetString("arango-pass")
-		flags.arangoAddr = viper.GetString("arango-addr")
-
-		flags.neptuneEndpoint = viper.GetString("neptune-endpoint")
-		flags.neptunePort = viper.GetInt("neptune-port")
-		flags.neptuneRegion = viper.GetString("neptune-region")
-		flags.neptuneUser = viper.GetString("neptune-user")
-		flags.neptuneRealm = viper.GetString("neptune-realm")
-
-		flags.kvStore = viper.GetString("kv-store")
-		flags.kvRedis = viper.GetString("kv-redis")
-		flags.kvTiKV = viper.GetString("kv-tikv")
 		startServer(cmd)
 	},
 }
 
 func init() {
 	cobra.OnInitialize(cli.InitConfig)
 
+	// Register common flags
 	set, err := cli.BuildFlags([]string{
-		"arango-addr", "arango-user", "arango-pass",
-		"neo4j-addr", "neo4j-user", "neo4j-pass", "neo4j-realm",
-		"neptune-endpoint", "neptune-port", "neptune-region", "neptune-user", "neptune-realm",
-		"gql-listen-port", "gql-tls-cert-file", "gql-tls-key-file", "gql-debug", "gql-backend", "gql-trace",
-		"db-address", "db-driver", "db-debug", "db-migrate", "db-conn-time",
-		"kv-store", "kv-redis", "kv-tikv", "enable-prometheus",
+		"gql-listen-port",
+		"gql-tls-cert-file",
+		"gql-tls-key-file",
+		"gql-debug",
+		"gql-backend",
+		"gql-trace",
+		"enable-prometheus",
 	})
 	if err != nil {
 		fmt.Fprintf(os.Stderr, "failed to setup flag: %v", err)
 		os.Exit(1)
 	}
 	rootCmd.Flags().AddFlagSet(set)
+
+	// Register backend-specific flags
+	backends.RegisterFlags(rootCmd)
+
 	if err := viper.BindPFlags(rootCmd.Flags()); err != nil {
 		fmt.Fprintf(os.Stderr, "failed to bind flags: %v", err)
 		os.Exit(1)

cmd/guacgql/cmd/server.go

@@ -26,47 +26,18 @@ import (
 	"syscall"
 	"time"
 
-	"github.com/guacsec/guac/pkg/version"
-
 	"github.com/99designs/gqlgen/graphql/handler/debug"
 	"github.com/99designs/gqlgen/graphql/playground"
 	"github.com/guacsec/guac/pkg/assembler/backends"
-	"github.com/guacsec/guac/pkg/assembler/backends/arangodb"
 	_ "github.com/guacsec/guac/pkg/assembler/backends/keyvalue"
-	"github.com/guacsec/guac/pkg/assembler/backends/neo4j"
-	"github.com/guacsec/guac/pkg/assembler/backends/neptune"
-	"github.com/guacsec/guac/pkg/assembler/kv"
-	"github.com/guacsec/guac/pkg/assembler/kv/redis"
 	"github.com/guacsec/guac/pkg/assembler/server"
 	"github.com/guacsec/guac/pkg/logging"
 	"github.com/guacsec/guac/pkg/metrics"
+	"github.com/guacsec/guac/pkg/version"
 	"github.com/spf13/cobra"
 	"github.com/spf13/viper"
-	"golang.org/x/exp/maps"
-)
-
-const (
-	arango   = "arango"
-	neo4js   = "neo4j"
-	ent      = "ent"
-	neptunes = "neptune"
-	keyvalue = "keyvalue"
 )
 
-type optsFunc func(context.Context) backends.BackendArgs
-
-var getOpts map[string]optsFunc
-
-func init() {
-	if getOpts == nil {
-		getOpts = make(map[string]optsFunc)
-	}
-	getOpts[arango] = getArango
-	getOpts[neo4js] = getNeo4j
-	getOpts[neptunes] = getNeptune
-	getOpts[keyvalue] = getKeyValue
-}
-
 func startServer(cmd *cobra.Command) {
 	var srvHandler http.Handler
 	ctx := logging.WithLogger(context.Background())
@@ -78,9 +49,15 @@ func startServer(cmd *cobra.Command) {
 		os.Exit(1)
 	}
 
-	backend, err := backends.Get(flags.backend, ctx, getOpts[flags.backend](ctx))
+	backendArgs, err := backends.GetBackendArgs(ctx, flags.backend)
+	if err != nil {
+		logger.Errorf("failed to parse backend flags with error: %v", err)
+		os.Exit(1)
+	}
+
+	backend, err := backends.Get(flags.backend, ctx, backendArgs)
 	if err != nil {
-		logger.Errorf("error creating %v backend: %w", flags.backend, err)
+		logger.Errorf("Error creating %v backend: %v", flags.backend, err)
 		os.Exit(1)
 	}
 
@@ -161,15 +138,9 @@ func setupPrometheus(ctx context.Context, name string) (metrics.MetricCollector,
 }
 
 func validateFlags() error {
-	if !slices.Contains(maps.Keys(getOpts), flags.backend) {
-		return fmt.Errorf("invalid graphql backend specified: %v", flags.backend)
-	}
 	if !slices.Contains(backends.List(), flags.backend) {
 		return fmt.Errorf("invalid graphql backend specified: %v", flags.backend)
 	}
-	if !slices.Contains([]string{"memmap", "redis", "tikv"}, flags.kvStore) {
-		return fmt.Errorf("invalid kv store specified: %v", flags.kvStore)
-	}
 	return nil
 }
 
@@ -183,57 +154,3 @@ func versionHandler(w http.ResponseWriter, r *http.Request) {
 	w.WriteHeader(http.StatusOK)
 	_, _ = fmt.Fprint(w, version.Version)
 }
-
-func getArango(_ context.Context) backends.BackendArgs {
-	return &arangodb.ArangoConfig{
-		User:   flags.arangoUser,
-		Pass:   flags.arangoPass,
-		DBAddr: flags.arangoAddr,
-	}
-}
-
-func getNeo4j(_ context.Context) backends.BackendArgs {
-	return &neo4j.Neo4jConfig{
-		User:   flags.nUser,
-		Pass:   flags.nPass,
-		Realm:  flags.nRealm,
-		DBAddr: flags.nAddr,
-	}
-}
-
-var tikvGS func(context.Context, string) (kv.Store, error)
-
-func getKeyValue(ctx context.Context) backends.BackendArgs {
-	logger := logging.FromContext(ctx)
-	switch flags.kvStore {
-	case "memmap":
-		// default is memmap
-		return nil
-	case "redis":
-		s, err := redis.GetStore(flags.kvRedis)
-		if err != nil {
-			logger.Fatalf("error with Redis: %v", err)
-		}
-		return s
-	case "tikv":
-		if tikvGS == nil {
-			logger.Fatal("TiKV not supported on 32-bit")
-		}
-		s, err := tikvGS(ctx, flags.kvTiKV)
-		if err != nil {
-			logger.Fatalf("error with TiKV: %v", err)
-		}
-		return s
-	}
-	return nil
-}
-
-func getNeptune(_ context.Context) backends.BackendArgs {
-	return &neptune.NeptuneConfig{
-		Endpoint: flags.neptuneEndpoint,
-		Port:     flags.neptunePort,
-		Region:   flags.neptuneRegion,
-		User:     flags.neptuneUser,
-		Realm:    flags.neptuneRealm,
-	}
-}

pkg/assembler/backends/arangodb/backend.go

@@ -23,6 +23,8 @@ import (
 	"time"
 
 	"github.com/99designs/gqlgen/graphql"
+	"github.com/spf13/cobra"
+	"github.com/spf13/viper"
 
 	jsoniter "github.com/json-iterator/go"
 
@@ -42,6 +44,13 @@ type ArangoConfig struct {
 	TestData bool
 }
 
+// flags holds the command-line flags for ArangoDB configuration
+var flags = struct {
+	addr string
+	user string
+	pass string
+}{}
+
 type arangoQueryBuilder struct {
 	query strings.Builder
 }
@@ -59,7 +68,7 @@ type index struct {
 }
 
 func init() {
-	backends.Register("arango", getBackend)
+	backends.Register("arango", getBackend, registerFlags, parseFlags)
 }
 
 func initIndex(name string, fields []string, unique bool) index {
@@ -117,6 +126,28 @@ func DeleteDatabase(ctx context.Context, args backends.BackendArgs) error {
 	return nil
 }
 
+// registerFlags registers ArangoDB-specific command line flags
+func registerFlags(cmd *cobra.Command) {
+	flagSet := cmd.Flags()
+	flagSet.StringVar(&flags.addr, "arango-addr", "", "ArangoDB database address")
+	flagSet.StringVar(&flags.user, "arango-user", "", "ArangoDB username")
+	flagSet.StringVar(&flags.pass, "arango-pass", "", "ArangoDB password")
+
+	// Bind to viper
+	viper.BindPFlag("arango-addr", flagSet.Lookup("arango-addr"))
+	viper.BindPFlag("arango-user", flagSet.Lookup("arango-user"))
+	viper.BindPFlag("arango-pass", flagSet.Lookup("arango-pass"))
+}
+
+// parseFlags returns the ArangoDB configuration from parsed flags
+func parseFlags(ctx context.Context) (backends.BackendArgs, error) {
+	return &ArangoConfig{
+		DBAddr: viper.GetString("arango-addr"),
+		User:   viper.GetString("arango-user"),
+		Pass:   viper.GetString("arango-pass"),
+	}, nil
+}
+
 func getBackend(ctx context.Context, args backends.BackendArgs) (backends.Backend, error) {
 	config, ok := args.(*ArangoConfig)
 	if !ok {