fix(cli): honor the atespace flag when applying resources

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

cmd/ax/main_test.go

@@ -16,20 +16,114 @@ 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"
 	"google.golang.org/grpc/status"
 	"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)

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