schemas/graph/cosmo.go

package graph

import (
	"context"
	"fmt"
	"os"
	"os/exec"
	"path/filepath"
	"time"

	"golang.org/x/sync/semaphore"
	"gopkg.in/yaml.v3"

	"gitea.unbound.se/unboundsoftware/schemas/graph/model"
)

// CommandExecutor is an interface for executing external commands
// This allows for mocking in tests
type CommandExecutor interface {
	Execute(name string, args ...string) ([]byte, error)
}

// DefaultCommandExecutor implements CommandExecutor using os/exec
type DefaultCommandExecutor struct{}

// Execute runs a command and returns its combined output
func (e *DefaultCommandExecutor) Execute(name string, args ...string) ([]byte, error) {
	cmd := exec.Command(name, args...)
	return cmd.CombinedOutput()
}

// GenerateCosmoRouterConfig generates a Cosmo Router execution config from subgraphs
// using the official wgc CLI tool via npx
func GenerateCosmoRouterConfig(subGraphs []*model.SubGraph) (string, error) {
	return GenerateCosmoRouterConfigWithExecutor(subGraphs, &DefaultCommandExecutor{})
}

// GenerateCosmoRouterConfigWithExecutor generates a Cosmo Router execution config from subgraphs
// using the provided command executor (useful for testing)
func GenerateCosmoRouterConfigWithExecutor(subGraphs []*model.SubGraph, executor CommandExecutor) (string, error) {
	if len(subGraphs) == 0 {
		return "", fmt.Errorf("no subgraphs provided")
	}

	// Create a temporary directory for composition
	tmpDir, err := os.MkdirTemp("", "cosmo-compose-*")
	if err != nil {
		return "", fmt.Errorf("create temp dir: %w", err)
	}
	defer os.RemoveAll(tmpDir)

	// Write each subgraph SDL to a file
	type SubgraphConfig struct {
		Name         string                 `yaml:"name"`
		RoutingURL   string                 `yaml:"routing_url,omitempty"`
		Schema       map[string]string      `yaml:"schema"`
		Subscription map[string]interface{} `yaml:"subscription,omitempty"`
	}

	type InputConfig struct {
		Version   int              `yaml:"version"`
		Subgraphs []SubgraphConfig `yaml:"subgraphs"`
	}

	inputConfig := InputConfig{
		Version:   1,
		Subgraphs: make([]SubgraphConfig, 0, len(subGraphs)),
	}

	for _, sg := range subGraphs {
		// Write SDL to a temp file
		schemaFile := filepath.Join(tmpDir, fmt.Sprintf("%s.graphql", sg.Service))
		if err := os.WriteFile(schemaFile, []byte(sg.Sdl), 0o644); err != nil {
			return "", fmt.Errorf("write schema file for %s: %w", sg.Service, err)
		}

		subgraphCfg := SubgraphConfig{
			Name: sg.Service,
			Schema: map[string]string{
				"file": schemaFile,
			},
		}

		if sg.URL != nil {
			subgraphCfg.RoutingURL = *sg.URL
		}

		if sg.WsURL != nil {
			subgraphCfg.Subscription = map[string]interface{}{
				"url":                   *sg.WsURL,
				"protocol":              "ws",
				"websocket_subprotocol": "graphql-ws",
			}
		}

		inputConfig.Subgraphs = append(inputConfig.Subgraphs, subgraphCfg)
	}

	// Write input config YAML
	inputFile := filepath.Join(tmpDir, "input.yaml")
	inputYAML, err := yaml.Marshal(inputConfig)
	if err != nil {
		return "", fmt.Errorf("marshal input config: %w", err)
	}
	if err := os.WriteFile(inputFile, inputYAML, 0o644); err != nil {
		return "", fmt.Errorf("write input config: %w", err)
	}

	// Execute wgc router compose
	// wgc is installed globally in the Docker image
	outputFile := filepath.Join(tmpDir, "config.json")
	output, err := executor.Execute("wgc", "router", "compose",
		"--input", inputFile,
		"--out", outputFile,
		"--suppress-warnings",
	)
	if err != nil {
		return "", fmt.Errorf("wgc router compose failed: %w\nOutput: %s", err, string(output))
	}

	// Read the generated config
	configJSON, err := os.ReadFile(outputFile)
	if err != nil {
		return "", fmt.Errorf("read output config: %w", err)
	}

	return string(configJSON), nil
}

// CosmoGenerator wraps config generation with a concurrency limit and timeout
// to prevent unbounded wgc process spawning under rapid schema updates.
type CosmoGenerator struct {
	sem      *semaphore.Weighted
	executor CommandExecutor
	timeout  time.Duration
}

// NewCosmoGenerator creates a CosmoGenerator that allows at most one concurrent
// wgc process and applies the given timeout to each generation attempt.
func NewCosmoGenerator(executor CommandExecutor, timeout time.Duration) *CosmoGenerator {
	return &CosmoGenerator{
		sem:      semaphore.NewWeighted(1),
		executor: executor,
		timeout:  timeout,
	}
}

// Generate produces a Cosmo Router config, blocking if another generation is
// already in progress. The provided context (plus the configured timeout)
// controls cancellation.
func (g *CosmoGenerator) Generate(ctx context.Context, subGraphs []*model.SubGraph) (string, error) {
	ctx, cancel := context.WithTimeout(ctx, g.timeout)
	defer cancel()

	if err := g.sem.Acquire(ctx, 1); err != nil {
		return "", fmt.Errorf("acquire cosmo generator: %w", err)
	}
	defer g.sem.Release(1)

	return GenerateCosmoRouterConfigWithExecutor(subGraphs, g.executor)
}