feat: add udf for returning a namespace ID from a namespace name (#1712)

Author: caphrim007

src/carnot/funcs/metadata/metadata_ops.cc

@@ -131,6 +131,7 @@ void RegisterMetadataOpsOrDie(px::carnot::udf::Registry* registry) {
   registry->RegisterOrDie<VizierNameUDF>("vizier_name");
   registry->RegisterOrDie<VizierNamespaceUDF>("vizier_namespace");
   registry->RegisterOrDie<GetClusterCIDRRangeUDF>("get_cidrs");
+  registry->RegisterOrDie<NamespaceNameToNamespaceIDUDF>("namespace_name_to_namespace_id");
 
   /*****************************************
    * Aggregate UDFs.

src/carnot/funcs/metadata/metadata_ops.h

@@ -3262,6 +3262,24 @@ class GetClusterCIDRRangeUDF : public udf::ScalarUDF {
   std::string cidrs_str_;
 };
 
+class NamespaceNameToNamespaceIDUDF : public ScalarUDF {
+ public:
+  StringValue Exec(FunctionContext* ctx, StringValue namespace_name) {
+    auto md = GetMetadataState(ctx);
+    auto namespace_id =
+        md->k8s_metadata_state().NamespaceIDByName(std::make_pair(namespace_name, namespace_name));
+    return namespace_id;
+  }
+
+  static udf::ScalarUDFDocBuilder Doc() {
+    return udf::ScalarUDFDocBuilder("Get the Kubernetes UID of the given namespace name.")
+        .Details("Get the Kubernetes UID of the given namespace name.")
+        .Example("df.kube_system_namespace_uid = px.namespace_name_to_namespace_id('kube-system')")
+        .Arg("arg1", "The name of the namespace to get the UID of.")
+        .Returns("The Kubernetes UID of the given namespace name");
+  }
+};
+
 void RegisterMetadataOpsOrDie(px::carnot::udf::Registry* registry);
 
 }  // namespace metadata

src/carnot/funcs/metadata/metadata_ops_test.cc

@@ -67,11 +67,13 @@ class MetadataOpsTest : public ::testing::Test {
     updates_->enqueue(px::metadatapb::testutils::CreateRunningServiceIPUpdatePB());
     updates_->enqueue(px::metadatapb::testutils::CreateRunningReplicaSetUpdatePB());
     updates_->enqueue(px::metadatapb::testutils::CreateRunningDeploymentUpdatePB());
+    updates_->enqueue(px::metadatapb::testutils::CreateRunningNamespaceUpdatePB());
     updates_->enqueue(px::metadatapb::testutils::CreateTerminatingContainerUpdatePB());
     updates_->enqueue(px::metadatapb::testutils::CreateTerminatingPodUpdatePB());
     updates_->enqueue(px::metadatapb::testutils::CreateTerminatingServiceUpdatePB());
     updates_->enqueue(px::metadatapb::testutils::CreateTerminatingReplicaSetUpdatePB());
     updates_->enqueue(px::metadatapb::testutils::CreateTerminatingDeploymentUpdatePB());
+    updates_->enqueue(px::metadatapb::testutils::CreateTerminatingNamespaceUpdatePB());
 
     auto s = px::md::ApplyK8sUpdates(10, metadata_state_.get(), &md_filter_, updates_.get());
 
@@ -1347,6 +1349,19 @@ TEST_F(MetadataOpsTest, get_cidrs) {
       absl::Substitute(R"(["$0","$1","$2"])", "10.0.0.1/32", pod_cidr_str, service_cidr_str));
 }
 
+TEST_F(MetadataOpsTest, namespace_name_to_namespace_id_test) {
+  auto function_ctx = std::make_unique<FunctionContext>(metadata_state_, nullptr);
+  auto udf_tester =
+      px::carnot::udf::UDFTester<NamespaceNameToNamespaceIDUDF>(std::move(function_ctx));
+  udf_tester.ForInput("namespace1").Expect("namespace_uid");
+  udf_tester.ForInput("terminating_namespace1").Expect("terminating_namespace_uid");
+  udf_tester.ForInput("badformat").Expect("");
+  updates_->enqueue(px::metadatapb::testutils::CreateTerminatedNamespaceUpdatePB());
+  // keep information about the namespace after termination
+  EXPECT_OK(px::md::ApplyK8sUpdates(11, metadata_state_.get(), &md_filter_, updates_.get()));
+  udf_tester.ForInput("terminating_namespace1").Expect("terminating_namespace_uid");
+}
+
 }  // namespace metadata
 }  // namespace funcs
 }  // namespace carnot

src/shared/k8s/metadatapb/test_proto.h

@@ -362,6 +362,30 @@ conditions: {
 }
 )";
 
+/*
+ * Templates for namespace updates.
+ */
+const char* kRunningNamespaceUpdatePbTxt = R"(
+uid: "namespace_uid"
+name: "namespace1"
+start_timestamp_ns: 101
+stop_timestamp_ns: 0
+)";
+
+const char* kTerminatingNamespaceUpdatePbTxt = R"(
+uid: "terminating_namespace_uid"
+name: "terminating_namespace1"
+start_timestamp_ns: 123
+stop_timestamp_ns: 0
+)";
+
+const char* kTerminatedNamespaceUpdatePbTxt = R"(
+uid: "terminating_namespace_uid"
+name: "terminating_namespace1"
+start_timestamp_ns: 123
+stop_timestamp_ns: 150
+)";
+
 std::unique_ptr<px::shared::k8s::metadatapb::ResourceUpdate> CreateRunningPodUpdatePB() {
   auto update = std::make_unique<px::shared::k8s::metadatapb::ResourceUpdate>();
   auto update_proto = absl::Substitute(kResourceUpdateTmpl, "pod_update", kRunningPodUpdatePbTxt);
@@ -530,6 +554,33 @@ std::unique_ptr<px::shared::k8s::metadatapb::ResourceUpdate> CreateTerminatedDep
   return update;
 }
 
+std::unique_ptr<px::shared::k8s::metadatapb::ResourceUpdate> CreateRunningNamespaceUpdatePB() {
+  auto update = std::make_unique<px::shared::k8s::metadatapb::ResourceUpdate>();
+  auto update_proto =
+      absl::Substitute(kResourceUpdateTmpl, "namespace_update", kRunningNamespaceUpdatePbTxt);
+  CHECK(google::protobuf::TextFormat::MergeFromString(update_proto, update.get()))
+      << "Failed to parse proto";
+  return update;
+}
+
+std::unique_ptr<px::shared::k8s::metadatapb::ResourceUpdate> CreateTerminatingNamespaceUpdatePB() {
+  auto update = std::make_unique<px::shared::k8s::metadatapb::ResourceUpdate>();
+  auto update_proto =
+      absl::Substitute(kResourceUpdateTmpl, "namespace_update", kTerminatingNamespaceUpdatePbTxt);
+  CHECK(google::protobuf::TextFormat::MergeFromString(update_proto, update.get()))
+      << "Failed to parse proto";
+  return update;
+}
+
+std::unique_ptr<px::shared::k8s::metadatapb::ResourceUpdate> CreateTerminatedNamespaceUpdatePB() {
+  auto update = std::make_unique<px::shared::k8s::metadatapb::ResourceUpdate>();
+  auto update_proto =
+      absl::Substitute(kResourceUpdateTmpl, "namespace_update", kTerminatedNamespaceUpdatePbTxt);
+  CHECK(google::protobuf::TextFormat::MergeFromString(update_proto, update.get()))
+      << "Failed to parse proto";
+  return update;
+}
+
 }  // namespace testutils
 }  // namespace metadatapb
 }  // namespace px