Add external_job_uri as a first-class proto field

The external job URI is already treated as special in Armada (dedicated
DB column, dedicated query API endpoint), but clients have to know the
magic annotation key "armadaproject.io/externalJobUri" to use it. This
makes the field undiscoverable and means clients get no help from type
checking or generated docs.

This adds external_job_uri as a proper proto field on JobSubmitRequestItem
and SubmitJob, so clients can set it directly. The server resolves the
value from the proto field first and falls back to the annotation for
backward compatibility. The value is mirrored into the annotation on the
outgoing event so older ingesters that only read the annotation continue
to work during rolling deploys.

The Airflow operator is updated to use the proto field when available
(guarded by hasattr for older client library versions) while continuing
to set the annotation for backward compat with older servers.

Signed-off-by: Dejan Zele Pejchev <pejcev.dejan@gmail.com>

Author: dejanzele

client/rust/src/builder.rs

@@ -35,6 +35,7 @@ pub struct JobRequestItemBuilder<S> {
     labels: HashMap<String, String>,
     annotations: HashMap<String, String>,
     scheduler: String,
+    external_job_uri: String,
     ingress: Vec<IngressConfig>,
     services: Vec<ServiceConfig>,
     pod_specs: Vec<PodSpec>,
@@ -56,6 +57,7 @@ impl JobRequestItemBuilder<NoPodSpec> {
             labels: HashMap::new(),
             annotations: HashMap::new(),
             scheduler: String::new(),
+            external_job_uri: String::new(),
             ingress: Vec::new(),
             services: Vec::new(),
             pod_specs: Vec::new(),
@@ -137,6 +139,13 @@ impl<S> JobRequestItemBuilder<S> {
         self
     }
 
+    /// Set a URI identifying this job in an external system (e.g. Airflow).
+    #[must_use]
+    pub fn external_job_uri(mut self, uri: impl Into<String>) -> Self {
+        self.external_job_uri = uri.into();
+        self
+    }
+
     /// Replace the entire ingress config list.
     #[must_use]
     pub fn ingress(mut self, i: Vec<IngressConfig>) -> Self {
@@ -189,6 +198,7 @@ impl<S> JobRequestItemBuilder<S> {
             labels: self.labels,
             annotations: self.annotations,
             scheduler: self.scheduler,
+            external_job_uri: self.external_job_uri,
             ingress: self.ingress,
             services: self.services,
             pod_specs: specs,
@@ -219,6 +229,7 @@ impl JobRequestItemBuilder<HasPodSpec> {
             ingress: self.ingress,
             services: self.services,
             scheduler: self.scheduler,
+            external_job_uri: self.external_job_uri,
             // Deprecated singular fields — zeroed, not used by this builder
             pod_spec: None,
             required_node_labels: HashMap::new(),

client/rust/src/gen/api.rs

@@ -84,6 +84,10 @@ pub struct JobSubmitRequestItem {
     /// If empty, the default scheduler is used.
     #[prost(string, tag = "11")]
     pub scheduler: ::prost::alloc::string::String,
+    /// URI identifying this job in an external system (e.g. Airflow).
+    /// If not set, the server falls back to the "armadaproject.io/externalJobUri" annotation.
+    #[prost(string, tag = "13")]
+    pub external_job_uri: ::prost::alloc::string::String,
 }
 #[derive(Clone, PartialEq, ::prost::Message)]
 pub struct IngressConfig {

internal/common/constants/constants.go

@@ -33,6 +33,9 @@ const (
 	FailFastAnnotation  = "armadaproject.io/failFast"
 	PoolAnnotation      = "armadaproject.io/pool"
 	ReservationTaintKey = "armadaproject.io/reservation"
+	// ExternalJobUriAnnotation is the legacy annotation key for setting an external job URI.
+	// Prefer the ExternalJobUri proto field on JobSubmitRequestItem / SubmitJob instead.
+	ExternalJobUriAnnotation = "armadaproject.io/externalJobUri"
 )
 
 var schedulingAnnotations = map[string]bool{

internal/lookoutingester/instructions/instructions.go

@@ -10,6 +10,7 @@ import (
 
 	"github.com/armadaproject/armada/internal/common/armadacontext"
 	"github.com/armadaproject/armada/internal/common/compress"
+	"github.com/armadaproject/armada/internal/common/constants"
 	"github.com/armadaproject/armada/internal/common/database/lookout"
 	"github.com/armadaproject/armada/internal/common/eventutil"
 	"github.com/armadaproject/armada/internal/common/ingest/metrics"
@@ -176,7 +177,12 @@ func (c *InstructionConverter) handleSubmitJob(
 
 	annotations := event.GetObjectMeta().GetAnnotations()
 	userAnnotations := extractUserAnnotations(c.userAnnotationPrefix, c.blocklistAnnotations, annotations)
-	externalJobUri := util.Truncate(annotations["armadaproject.io/externalJobUri"], maxAnnotationValLen)
+	// Prefer the proto field; fall back to annotation for events produced before the field existed.
+	externalJobUri := event.ExternalJobUri
+	if externalJobUri == "" {
+		externalJobUri = annotations[constants.ExternalJobUriAnnotation]
+	}
+	externalJobUri = util.Truncate(externalJobUri, maxAnnotationValLen)
 
 	job := model.CreateJobInstruction{
 		JobId:                     event.JobId,

internal/lookoutingester/instructions/instructions_test.go

@@ -14,6 +14,7 @@ import (
 
 	"github.com/armadaproject/armada/internal/common/armadacontext"
 	"github.com/armadaproject/armada/internal/common/compress"
+	"github.com/armadaproject/armada/internal/common/constants"
 	"github.com/armadaproject/armada/internal/common/database/lookout"
 	"github.com/armadaproject/armada/internal/common/eventutil"
 	"github.com/armadaproject/armada/internal/common/ingest/testfixtures"
@@ -223,9 +224,9 @@ func TestConvert(t *testing.T) {
 	}
 	submit.GetSubmitJob().GetMainObject().GetPodSpec().GetPodSpec().PriorityClassName = priorityClass
 	submit.GetSubmitJob().GetObjectMeta().Annotations = map[string]string{
-		userAnnotationPrefix + "a":        "0",
-		"b":                               "1",
-		"armadaproject.io/externalJobUri": "external-job-uri",
+		userAnnotationPrefix + "a":         "0",
+		"b":                                "1",
+		constants.ExternalJobUriAnnotation: "external-job-uri",
 	}
 	job, err := eventutil.ApiJobFromLogSubmitJob(testfixtures.UserId, []string{}, testfixtures.Queue, testfixtures.JobsetName, testfixtures.BaseTime, submit.GetSubmitJob())
 	assert.NoError(t, err)
@@ -249,9 +250,9 @@ func TestConvert(t *testing.T) {
 		JobProto:                  jobProto,
 		PriorityClass:             pointer.String(priorityClass),
 		Annotations: map[string]string{
-			"a":                               "0",
-			"b":                               "1",
-			"armadaproject.io/externalJobUri": "external-job-uri",
+			"a":                                "0",
+			"b":                                "1",
+			constants.ExternalJobUriAnnotation: "external-job-uri",
 		},
 		ExternalJobUri: "external-job-uri",
 	}
@@ -538,14 +539,63 @@ func TestConvert(t *testing.T) {
 	}
 }
 
+func TestExternalJobUriProtoFieldPreferred(t *testing.T) {
+	tests := map[string]struct {
+		protoField  string
+		annotations map[string]string
+		expected    string
+	}{
+		"Proto field takes priority over annotation": {
+			protoField: "airflow://dag/task/run/0",
+			annotations: map[string]string{
+				constants.ExternalJobUriAnnotation: "old-value",
+			},
+			expected: "airflow://dag/task/run/0",
+		},
+		"Falls back to annotation when proto field empty": {
+			protoField: "",
+			annotations: map[string]string{
+				constants.ExternalJobUriAnnotation: "airflow://dag/task/run/0",
+			},
+			expected: "airflow://dag/task/run/0",
+		},
+		"Empty when neither set": {
+			protoField:  "",
+			annotations: map[string]string{},
+			expected:    "",
+		},
+	}
+
+	for name, tc := range tests {
+		t.Run(name, func(t *testing.T) {
+			submit, err := testfixtures.DeepCopy(testfixtures.Submit)
+			require.NoError(t, err)
+			submit.GetSubmitJob().ExternalJobUri = tc.protoField
+			submit.GetSubmitJob().GetObjectMeta().Annotations = tc.annotations
+
+			events := &utils.EventsWithIds[*armadaevents.EventSequence]{
+				Events: []*armadaevents.EventSequence{
+					testfixtures.NewEventSequence(submit),
+				},
+				MessageIds: []pulsar.MessageID{pulsarutils.NewMessageId(1)},
+			}
+
+			converter := NewInstructionConverter(metrics.Get().Metrics, userAnnotationPrefix, []string{}, &compress.NoOpCompressor{})
+			instructionSet := converter.Convert(armadacontext.TODO(), events)
+			require.Len(t, instructionSet.JobsToCreate, 1)
+			assert.Equal(t, tc.expected, instructionSet.JobsToCreate[0].ExternalJobUri)
+		})
+	}
+}
+
 func TestTruncatesStringsThatAreTooLong(t *testing.T) {
 	longString := strings.Repeat("x", 4000)
 
 	submit, err := testfixtures.DeepCopy(testfixtures.Submit)
 	assert.NoError(t, err)
 	submit.GetSubmitJob().GetMainObject().GetPodSpec().GetPodSpec().PriorityClassName = longString
 	submit.GetSubmitJob().GetObjectMeta().Annotations = map[string]string{
-		"armadaproject.io/externalJobUri": longString,
+		constants.ExternalJobUriAnnotation: longString,
 	}
 
 	leased, err := testfixtures.DeepCopy(testfixtures.Leased)

internal/server/submit/conversion/conversions.go

@@ -8,6 +8,7 @@ import (
 	networking "k8s.io/api/networking/v1"
 
 	"github.com/armadaproject/armada/internal/common"
+	"github.com/armadaproject/armada/internal/common/constants"
 	log "github.com/armadaproject/armada/internal/common/logging"
 	armadaslices "github.com/armadaproject/armada/internal/common/slices"
 	"github.com/armadaproject/armada/internal/common/util"
@@ -29,13 +30,27 @@ func SubmitJobFromApiRequest(
 	priority := PriorityAsInt32(jobReq.GetPriority())
 	ingressesAndServices := convertIngressesAndServices(config, jobReq, jobId, jobSetId, queue, owner)
 
+	// Resolve externalJobUri: prefer the proto field, fall back to annotation.
+	annotations := jobReq.GetAnnotations()
+	externalJobUri := jobReq.GetExternalJobUri()
+	if externalJobUri == "" {
+		externalJobUri = annotations[constants.ExternalJobUriAnnotation]
+	}
+	// Mirror into annotation so old ingesters (that only read the annotation) still work during rolling deploys.
+	if externalJobUri != "" {
+		if annotations == nil {
+			annotations = make(map[string]string)
+		}
+		annotations[constants.ExternalJobUriAnnotation] = externalJobUri
+	}
+
 	msg := &armadaevents.SubmitJob{
 		JobId:           jobId,
 		DeduplicationId: jobReq.GetClientId(),
 		Priority:        priority,
 		ObjectMeta: &armadaevents.ObjectMeta{
 			Namespace:   jobReq.GetNamespace(),
-			Annotations: jobReq.GetAnnotations(),
+			Annotations: annotations,
 			Labels:      jobReq.GetLabels(),
 		},
 		MainObject: &armadaevents.KubernetesMainObject{
@@ -45,8 +60,9 @@ func SubmitJobFromApiRequest(
 				},
 			},
 		},
-		Objects:   ingressesAndServices,
-		Scheduler: jobReq.Scheduler,
+		Objects:        ingressesAndServices,
+		Scheduler:      jobReq.Scheduler,
+		ExternalJobUri: externalJobUri,
 	}
 
 	postProcess(msg, config)

internal/server/submit/conversion/conversions_test.go

@@ -5,10 +5,12 @@ import (
 	"testing"
 
 	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
 	v1 "k8s.io/api/core/v1"
 	networking "k8s.io/api/networking/v1"
 	"k8s.io/apimachinery/pkg/api/resource"
 
+	"github.com/armadaproject/armada/internal/common/constants"
 	"github.com/armadaproject/armada/internal/server/configuration"
 	"github.com/armadaproject/armada/internal/server/submit/testfixtures"
 	"github.com/armadaproject/armada/pkg/api"
@@ -907,6 +909,62 @@ func jobSubmitRequestItemWithClassicInitContainer() *api.JobSubmitRequestItem {
 	return req
 }
 
+func TestExternalJobUri(t *testing.T) {
+	tests := map[string]struct {
+		protoField     string
+		annotations    map[string]string
+		expectedUri    string
+		expectMirrored bool
+	}{
+		"Proto field takes priority over annotation": {
+			protoField:     "airflow://dag/task/run/0",
+			annotations:    map[string]string{constants.ExternalJobUriAnnotation: "old-annotation-value"},
+			expectedUri:    "airflow://dag/task/run/0",
+			expectMirrored: true,
+		},
+		"Falls back to annotation when proto field empty": {
+			protoField:     "",
+			annotations:    map[string]string{constants.ExternalJobUriAnnotation: "airflow://dag/task/run/0"},
+			expectedUri:    "airflow://dag/task/run/0",
+			expectMirrored: true,
+		},
+		"Empty when neither set": {
+			protoField:     "",
+			annotations:    map[string]string{},
+			expectedUri:    "",
+			expectMirrored: false,
+		},
+		"Proto field mirrored into annotation": {
+			protoField:     "airflow://dag/task/run/0",
+			annotations:    nil,
+			expectedUri:    "airflow://dag/task/run/0",
+			expectMirrored: true,
+		},
+	}
+
+	for name, tc := range tests {
+		t.Run(name, func(t *testing.T) {
+			jobReq := testfixtures.JobSubmitRequestItem(1)
+			jobReq.Annotations = tc.annotations
+			jobReq.ExternalJobUri = tc.protoField
+
+			msg := SubmitJobFromApiRequest(
+				jobReq,
+				testfixtures.DefaultSubmissionConfig(),
+				testfixtures.DefaultJobset,
+				testfixtures.DefaultQueue.Name,
+				testfixtures.DefaultOwner,
+				testfixtures.TestUlidGenerator(),
+			)
+
+			require.Equal(t, tc.expectedUri, msg.ExternalJobUri)
+			if tc.expectMirrored {
+				assert.Equal(t, tc.expectedUri, msg.GetObjectMeta().GetAnnotations()[constants.ExternalJobUriAnnotation])
+			}
+		})
+	}
+}
+
 func SubmitJobMsgWithK8sObjects(objects []*armadaevents.KubernetesObject, initContainers ...v1.Container) *armadaevents.SubmitJob {
 	submitMsg := testfixtures.SubmitJob(1)
 	submitMsg.Objects = objects

pkg/api/api.swagger.go

@@ -1881,6 +1881,10 @@ func SwaggerJsonTemplate() string {
 		"        \"clientId\": {\n" +
 		"          \"type\": \"string\"\n" +
 		"        },\n" +
+		"        \"externalJobUri\": {\n" +
+		"          \"description\": \"URI identifying this job in an external system (e.g. Airflow).\\nIf not set, the server falls back to the \\\"armadaproject.io/externalJobUri\\\" annotation.\",\n" +
+		"          \"type\": \"string\"\n" +
+		"        },\n" +
 		"        \"ingress\": {\n" +
 		"          \"type\": \"array\",\n" +
 		"          \"items\": {\n" +

pkg/api/api.swagger.json

@@ -1870,6 +1870,10 @@
         "clientId": {
           "type": "string"
         },
+        "externalJobUri": {
+          "description": "URI identifying this job in an external system (e.g. Airflow).\nIf not set, the server falls back to the \"armadaproject.io/externalJobUri\" annotation.",
+          "type": "string"
+        },
         "ingress": {
           "type": "array",
           "items": {