Add capability to get pod labels

src/carnot/funcs/metadata/metadata_ops.cc

@@ -45,6 +45,7 @@ void RegisterMetadataOpsOrDie(px::carnot::udf::Registry* registry) {
   registry->RegisterOrDie<HasServiceNameUDF>("has_service_name");
   registry->RegisterOrDie<IPToPodIDUDF>("ip_to_pod_id");
   registry->RegisterOrDie<PodIDToPodNameUDF>("pod_id_to_pod_name");
+  registry->RegisterOrDie<PodIDToPodLabelsUDF>("pod_id_to_pod_labels");
   registry->RegisterOrDie<PodIDToNamespaceUDF>("pod_id_to_namespace");
   registry->RegisterOrDie<PodIDToNodeNameUDF>("pod_id_to_node_name");
   registry->RegisterOrDie<PodIDToPodStartTimeUDF>("pod_id_to_start_time");

src/carnot/funcs/metadata/metadata_ops.h

@@ -122,6 +122,27 @@ class PodIDToPodNameUDF : public ScalarUDF {
   }
 };
 
+class PodIDToPodLabelsUDF : public ScalarUDF {
+ public:
+  StringValue Exec(FunctionContext* ctx, StringValue pod_id) {
+    auto md = GetMetadataState(ctx);
+
+    const auto* pod_info = md->k8s_metadata_state().PodInfoByID(pod_id);
+    if (pod_info != nullptr) {
+      return pod_info->labels();
+    }
+    return "";
+  }
+
+  static udf::ScalarUDFDocBuilder Doc() {
+    return udf::ScalarUDFDocBuilder("Get labels of a pod from its pod ID.")
+        .Details("Gets the kubernetes pod labels for the pod from its pod ID.")
+        .Example("df.labels = px.pod_id_to_pod_labels(df.pod_id)")
+        .Arg("pod_id", "The pod ID of the pod to get the labels for.")
+        .Returns("The k8s pod labels for the pod ID passed in.");
+  }
+};
+
 class PodNameToPodIDUDF : public ScalarUDF {
  public:
   StringValue Exec(FunctionContext* ctx, StringValue pod_name) {

src/carnot/funcs/metadata/metadata_ops_test.cc

@@ -105,6 +105,14 @@ TEST_F(MetadataOpsTest, pod_id_to_pod_name_test) {
   udf_tester.ForInput("missing").Expect("");
 }
 
+TEST_F(MetadataOpsTest, pod_id_to_pod_labels_test) {
+  auto function_ctx = std::make_unique<FunctionContext>(metadata_state_, nullptr);
+  auto udf_tester = px::carnot::udf::UDFTester<PodIDToPodLabelsUDF>(std::move(function_ctx));
+  udf_tester.ForInput("1_uid").Expect("{\"k1\":\"v1\", \"k2\":\"v2\"}");
+  udf_tester.ForInput("2_uid").Expect("{\"k1\":\"v1\"}");
+  udf_tester.ForInput("missing").Expect("");
+}
+
 TEST_F(MetadataOpsTest, pod_name_to_pod_id_test) {
   auto function_ctx = std::make_unique<FunctionContext>(metadata_state_, nullptr);
   auto udf_tester = px::carnot::udf::UDFTester<PodNameToPodIDUDF>(std::move(function_ctx));

src/shared/k8s/metadatapb/metadata.proto

@@ -413,6 +413,8 @@ message PodUpdate {
   string message = 14;
   // A brief CamelCase message indicating details about why the pod is in this state.
   string reason = 15;
+  // A json object containing pod label keys and values
+  string labels = 17;
 }
 
 enum ContainerType {

src/shared/k8s/metadatapb/test_proto.h

@@ -49,6 +49,7 @@ const char* kRunningPodUpdatePbTxt = R"(
 uid: "1_uid"
 name: "running_pod"
 namespace: "pl"
+labels: "{\"k1\":\"v1\", \"k2\":\"v2\"}"
 start_timestamp_ns: 5
 container_ids: "pod1_container_1"
 qos_class: QOS_CLASS_GUARANTEED
@@ -68,6 +69,7 @@ const char* kToBeTerminatedPodUpdatePbTxt = R"(
 uid: "2_uid"
 name: "terminating_pod"
 namespace: "pl"
+labels: "{\"k1\":\"v1\"}"
 start_timestamp_ns: 10
 container_ids: "pod2_container_1"
 qos_class: QOS_CLASS_BEST_EFFORT

src/shared/metadata/k8s_objects.h

@@ -292,9 +292,12 @@ class PodInfo : public K8sMetadataObject {
   void set_node_name(std::string_view node_name) { node_name_ = node_name; }
   void set_hostname(std::string_view hostname) { hostname_ = hostname; }
   void set_pod_ip(std::string_view pod_ip) { pod_ip_ = pod_ip; }
+  void set_pod_labels(std::string labels) { labels_ = labels; }
+
   const std::string& node_name() const { return node_name_; }
   const std::string& hostname() const { return hostname_; }
   const std::string& pod_ip() const { return pod_ip_; }
+  const std::string& labels() const { return labels_; }
 
   const absl::flat_hash_set<std::string>& containers() const { return containers_; }
   const absl::flat_hash_set<std::string>& services() const { return services_; }
@@ -334,6 +337,7 @@ class PodInfo : public K8sMetadataObject {
   std::string node_name_;
   std::string hostname_;
   std::string pod_ip_;
+  std::string labels_;
 };
 
 struct UPIDStartTSCompare {

src/shared/metadata/metadata_state.cc

@@ -195,6 +195,7 @@ Status K8sMetadataState::HandlePodUpdate(const PodUpdate& update) {
   pod_info->set_conditions(ConvertToPodConditions(update.conditions()));
   pod_info->set_phase_message(update.message());
   pod_info->set_phase_reason(update.reason());
+  pod_info->set_pod_labels(update.labels());
 
   pods_by_name_[{ns, name}] = object_uid;
   // Filter out daemonsets which don't have their own, unique podIP.

src/shared/metadata/metadata_state_test.cc

@@ -33,6 +33,7 @@ constexpr char kPod0UpdatePbTxt[] = R"(
   uid: "pod0_uid"
   name: "pod0"
   namespace: "ns0"
+  labels: "{\"k1\":\"v1\", \"k2\":\"v2\"}"
   start_timestamp_ns: 101
   stop_timestamp_ns: 103
   container_ids: "container0_uid"
@@ -57,6 +58,7 @@ constexpr char kPod1UpdatePbTxt[] = R"(
   uid: "pod1_uid"
   name: "pod1"
   namespace: "ns0"
+  labels: "{\"k1\":\"v1\"}"
   start_timestamp_ns: 101
   stop_timestamp_ns: 103
   container_ids: "container0_uid"
@@ -81,6 +83,7 @@ constexpr char kPod2UpdatePbTxt[] = R"(
   uid: "pod2_uid"
   name: "pod2"
   namespace: "ns0"
+  labels: "{\"k1\":\"v1\"}"
   start_timestamp_ns: 101
   stop_timestamp_ns: 0
   container_ids: "container0_uid"
@@ -217,6 +220,7 @@ TEST(K8sMetadataStateTest, HandlePodUpdate) {
   EXPECT_EQ("pod0_uid", pod_info->uid());
   EXPECT_EQ("pod0", pod_info->name());
   EXPECT_EQ("ns0", pod_info->ns());
+  EXPECT_EQ("{\"k1\":\"v1\", \"k2\":\"v2\"}", pod_info->labels());
   EXPECT_EQ(PodQOSClass::kGuaranteed, pod_info->qos_class());
   EXPECT_EQ(PodPhase::kRunning, pod_info->phase());
   EXPECT_EQ(1, pod_info->conditions().size());

src/vizier/services/metadata/controllers/k8smeta/k8s_metadata_handler.go

@@ -19,6 +19,7 @@
 package k8smeta
 
 import (
+	"encoding/json"
 	"fmt"
 	"net"
 	"path"
@@ -938,13 +939,19 @@ func getResourceUpdateFromPod(pod *metadatapb.Pod, uv int64) *metadatapb.Resourc
 		hostname = pod.Spec.Hostname
 	}
 
+	var podLabels []byte
+	if pod.Metadata.Labels != nil {
+		podLabels, _ = json.Marshal(pod.Metadata.Labels)
+	}
+
 	update := &metadatapb.ResourceUpdate{
 		UpdateVersion: uv,
 		Update: &metadatapb.ResourceUpdate_PodUpdate{
 			PodUpdate: &metadatapb.PodUpdate{
 				UID:              pod.Metadata.UID,
 				Name:             pod.Metadata.Name,
 				Namespace:        pod.Metadata.Namespace,
+				Labels:           string(podLabels),
 				StartTimestampNS: pod.Metadata.CreationTimestampNS,
 				StopTimestampNS:  pod.Metadata.DeletionTimestampNS,
 				QOSClass:         pod.Status.QOSClass,

src/vizier/services/metadata/controllers/k8smeta/k8s_metadata_handler_test.go

@@ -981,6 +981,7 @@ func TestPodUpdateProcessor_GetUpdatesToSend(t *testing.T) {
 					UID:              "ijkl",
 					Name:             "object_md",
 					Namespace:        "",
+					Labels:           `{"app":"myApp1","project":"myProj1"}`,
 					StartTimestampNS: 4,
 					StopTimestampNS:  6,
 					QOSClass:         metadatapb.QOS_CLASS_BURSTABLE,

src/vizier/services/metadata/controllers/testutils/test_proto.go

@@ -593,6 +593,7 @@ metadata {
     name: "test"
     uid: "abcd"
   }
+  labels: <key:"app" value:"myApp1"> labels: <key:"project" value:"myProj1">
   creation_timestamp_ns: 4
   deletion_timestamp_ns: 6
 }