diff --git a/cmd/ax/main.go b/cmd/ax/main.go index 9771f464..31a61f62 100644 --- a/cmd/ax/main.go +++ b/cmd/ax/main.go @@ -49,6 +49,7 @@ func main() { cmd string cleanArgs []string atespace = "default" + atespaceFlag *string explicitServer = "" kubeContext = "" axNamespace = "ax-system" @@ -60,10 +61,12 @@ func main() { if arg == "-a" || arg == "--atespace" { if i+1 < len(args) { atespace = args[i+1] + atespaceFlag = &atespace i++ } } else if strings.HasPrefix(arg, "--atespace=") { atespace = strings.TrimPrefix(arg, "--atespace=") + atespaceFlag = &atespace } else if arg == "--server" { if i+1 < len(args) { explicitServer = args[i+1] @@ -132,7 +135,7 @@ func main() { switch cmd { case "apply": - err = runApply(serverURL, cleanArgs) + err = runApply(serverURL, atespaceFlag, cleanArgs) case "get": err = runGet(serverURL, atespace, cleanArgs) case "describe": @@ -204,7 +207,8 @@ func getAXClient(serverURL string) (v1alpha1.AXClient, *grpc.ClientConn, error) return v1alpha1.NewAXClient(conn), conn, nil } -func runApply(serverURL string, args []string) error { +// A nil atespace leaves manifest atespaces unchanged and lets the server default missing ones. +func runApply(serverURL string, atespace *string, args []string) error { data, ok, err := manifestFromArgs(args) if err != nil { return err @@ -244,7 +248,7 @@ func runApply(serverURL string, args []string) error { } } - kind, name, outcome, err := applyDocument(ctx, client, &doc) + kind, name, outcome, err := applyDocument(ctx, client, &doc, atespace) if err != nil { return fmt.Errorf("applying document %d: %w", docIndex, err) } @@ -255,7 +259,7 @@ func runApply(serverURL string, args []string) error { // applyDocument decodes one manifest by its kind and submits it with the matching // Update RPC. It reports the kind, the resource name, and whether the resource was // created, configured (spec changed), or unchanged, in the style of kubectl apply. -func applyDocument(ctx context.Context, client v1alpha1.AXClient, doc *yaml.Node) (kind, name, outcome string, err error) { +func applyDocument(ctx context.Context, client v1alpha1.AXClient, doc *yaml.Node, atespace *string) (kind, name, outcome string, err error) { var head struct { Kind string `yaml:"kind"` } @@ -269,6 +273,9 @@ func applyDocument(ctx context.Context, client v1alpha1.AXClient, doc *yaml.Node if err := doc.Decode(&task); err != nil { return "", "", "", err } + if err := setApplyAtespace(task.Metadata, atespace); err != nil { + return "", "", "", err + } existing, err := client.GetTask(ctx, &v1alpha1.GetTaskRequest{Atespace: task.GetMetadata().GetAtespace(), Name: task.GetMetadata().GetName()}) outcome, err := applyOutcome(err, existing.GetSpec(), task.GetSpec()) if err != nil { @@ -282,6 +289,9 @@ func applyDocument(ctx context.Context, client v1alpha1.AXClient, doc *yaml.Node if err := doc.Decode(&ws); err != nil { return "", "", "", err } + if err := setApplyAtespace(ws.Metadata, atespace); err != nil { + return "", "", "", err + } existing, err := client.GetWorkspace(ctx, &v1alpha1.GetWorkspaceRequest{Atespace: ws.GetMetadata().GetAtespace(), Name: ws.GetMetadata().GetName()}) outcome, err := applyOutcome(err, existing.GetSpec(), ws.GetSpec()) if err != nil { @@ -295,6 +305,9 @@ func applyDocument(ctx context.Context, client v1alpha1.AXClient, doc *yaml.Node if err := doc.Decode(&m); err != nil { return "", "", "", err } + if err := setApplyAtespace(m.Metadata, atespace); err != nil { + return "", "", "", err + } existing, err := client.GetModel(ctx, &v1alpha1.GetModelRequest{Atespace: m.GetMetadata().GetAtespace(), Name: m.GetMetadata().GetName()}) outcome, err := applyOutcome(err, existing.GetSpec(), m.GetSpec()) if err != nil { @@ -310,6 +323,18 @@ func applyDocument(ctx context.Context, client v1alpha1.AXClient, doc *yaml.Node } } +// setApplyAtespace fills a missing atespace or rejects a mismatch with an explicit flag. +func setApplyAtespace(meta *v1alpha1.ObjectMeta, atespace *string) error { + if meta == nil || atespace == nil { + return nil + } + if meta.Atespace != "" && meta.Atespace != *atespace { + return fmt.Errorf("metadata.atespace %q does not match --atespace %q", meta.Atespace, *atespace) + } + meta.Atespace = *atespace + return nil +} + // applyOutcome classifies an apply from the result of looking up the existing // resource: "created" when it did not exist, "unchanged" when its spec already // matches, and "configured" otherwise. Lookup failures other than NotFound are diff --git a/cmd/ax/main_test.go b/cmd/ax/main_test.go index d1c9d031..90bfa21d 100644 --- a/cmd/ax/main_test.go +++ b/cmd/ax/main_test.go @@ -16,13 +16,18 @@ package main import ( "context" + "errors" + "fmt" "net" "os" + "os/exec" "path/filepath" "strings" "testing" + "time" "github.com/google/ax/internal/server" + "github.com/google/ax/internal/store" "github.com/google/ax/internal/store/memory" "github.com/google/ax/pkg/apis/v1alpha1" "google.golang.org/grpc/codes" @@ -30,6 +35,95 @@ import ( "gopkg.in/yaml.v3" ) +func TestApplyAtespace(t *testing.T) { + // Exercise flag parsing and dispatch as well as the resource RPCs. + binary := filepath.Join(t.TempDir(), "ax") + if output, err := exec.Command("go", "build", "-o", binary, ".").CombinedOutput(); err != nil { + t.Fatalf("building CLI: %v\n%s", err, output) + } + for _, kind := range []string{"Task", "Workspace", "Model"} { + for _, tc := range []struct { + name string + manifestAtespace string + flags []string + wantAtespace string + wantErr string + }{ + {name: "default", wantAtespace: "default"}, + {name: "flag only", flags: []string{"-a", "team-a"}, wantAtespace: "team-a"}, + {name: "manifest only", manifestAtespace: "team-a", wantAtespace: "team-a"}, + {name: "matching", manifestAtespace: "team-a", flags: []string{"--atespace", "team-a"}, wantAtespace: "team-a"}, + {name: "conflicting", manifestAtespace: "team-a", flags: []string{"--atespace=team-b"}, wantErr: `metadata.atespace "team-a" does not match --atespace "team-b"`}, + {name: "explicit default", manifestAtespace: "team-a", flags: []string{"-a", "default"}, wantErr: `metadata.atespace "team-a" does not match --atespace "default"`}, + } { + t.Run(kind+"/"+tc.name, func(t *testing.T) { + s := memory.NewStore() + listener, err := net.Listen("tcp", "127.0.0.1:0") + if err != nil { + t.Fatal(err) + } + srv := server.NewServer(s).GRPCServer() + t.Cleanup(srv.Stop) + go func() { _ = srv.Serve(listener) }() + + path := filepath.Join(t.TempDir(), "resource.yaml") + manifest := fmt.Sprintf("apiVersion: ax.io/v1alpha1\nkind: %s\nmetadata:\n name: example\n", kind) + if tc.manifestAtespace != "" { + manifest += fmt.Sprintf(" atespace: %q\n", tc.manifestAtespace) + } + manifest += "spec: {}\n" + if err := os.WriteFile(path, []byte(manifest), 0600); err != nil { + t.Fatal(err) + } + ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second) + defer cancel() + args := append([]string{"--server", listener.Addr().String(), "apply", "-f", path}, tc.flags...) + output, err := exec.CommandContext(ctx, binary, args...).CombinedOutput() + if tc.wantErr != "" { + var exitErr *exec.ExitError + if !errors.As(err, &exitErr) || exitErr.ExitCode() != 1 || !strings.Contains(string(output), tc.wantErr) { + t.Fatalf("expected exit 1 with %q, got %v\n%s", tc.wantErr, err, output) + } + } else { + wantOutput := strings.ToLower(kind) + ".ax.io/example created\n" + if err != nil || string(output) != wantOutput { + t.Fatalf("expected %q, got %v\n%s", wantOutput, err, output) + } + // Reapplying must look up the resource in the same atespace. + output, err = exec.CommandContext(ctx, binary, args...).CombinedOutput() + wantOutput = strings.ToLower(kind) + ".ax.io/example unchanged\n" + if err != nil || string(output) != wantOutput { + t.Fatalf("expected %q on reapply, got %v\n%s", wantOutput, err, output) + } + } + + for _, atespace := range []string{"default", "team-a", "team-b"} { + var meta *v1alpha1.ObjectMeta + var getErr error + switch kind { + case "Task": + resource, err := s.GetTask(ctx, atespace, "example") + meta, getErr = resource.GetMetadata(), err + case "Workspace": + resource, err := s.GetWorkspace(ctx, atespace, "example") + meta, getErr = resource.GetMetadata(), err + case "Model": + resource, err := s.GetModel(ctx, atespace, "example") + meta, getErr = resource.GetMetadata(), err + } + if atespace == tc.wantAtespace { + if getErr != nil || meta.GetAtespace() != atespace { + t.Errorf("expected resource in %q, got metadata %v, error %v", atespace, meta, getErr) + } + } else if !errors.Is(getErr, store.ErrNotFound) { + t.Errorf("expected no resource in %q, got metadata %v, error %v", atespace, meta, getErr) + } + } + }) + } + } +} + func TestRunApplyEmptyDocuments(t *testing.T) { const first = "apiVersion: ax.io/v1alpha1\nkind: Workspace\nmetadata:\n name: first\nspec: {}\n" second := strings.Replace(first, "name: first", "name: second", 1) @@ -77,7 +171,7 @@ func TestRunApplyEmptyDocuments(t *testing.T) { _ = output.Close() }) - runErr := runApply(listener.Addr().String(), []string{"-f", path}) + runErr := runApply(listener.Addr().String(), nil, []string{"-f", path}) if tc.wantErr == "" { if runErr != nil { t.Fatal(runErr) diff --git a/docs/manifests.md b/docs/manifests.md index 744248fa..772c1330 100644 --- a/docs/manifests.md +++ b/docs/manifests.md @@ -4,6 +4,8 @@ All four kinds can live in one multi-document YAML file. See [`examples/task.yam `metadata.name` and `metadata.atespace` become Substrate resource names, so they must be lowercase RFC 1123 labels: at most 63 lowercase alphanumeric characters or `-`, starting and ending with an alphanumeric character. `ax apply` rejects anything else up front rather than letting the task fail later with `ActorCreationFailed`. +When `metadata.atespace` is omitted, `ax apply -a team-a -f manifest.yaml` uses `team-a`. Without `-a` or `--atespace`, the manifest's atespace is preserved, or defaults to `default` if omitted. If an explicit flag disagrees with `metadata.atespace`, that resource is rejected before it is submitted. + ## Task ```yaml