diff --git a/src/carnot/funcs/metadata/metadata_ops.cc b/src/carnot/funcs/metadata/metadata_ops.cc index 9ce87dea2f2..76220e0bf6b 100644 --- a/src/carnot/funcs/metadata/metadata_ops.cc +++ b/src/carnot/funcs/metadata/metadata_ops.cc @@ -131,6 +131,7 @@ void RegisterMetadataOpsOrDie(px::carnot::udf::Registry* registry) { registry->RegisterOrDie("vizier_name"); registry->RegisterOrDie("vizier_namespace"); registry->RegisterOrDie("get_cidrs"); + registry->RegisterOrDie("namespace_name_to_namespace_id"); /***************************************** * Aggregate UDFs. diff --git a/src/carnot/funcs/metadata/metadata_ops.h b/src/carnot/funcs/metadata/metadata_ops.h index bca7437298c..6bd3db98ea8 100644 --- a/src/carnot/funcs/metadata/metadata_ops.h +++ b/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 diff --git a/src/carnot/funcs/metadata/metadata_ops_test.cc b/src/carnot/funcs/metadata/metadata_ops_test.cc index a405be9a45e..b9d225326bb 100644 --- a/src/carnot/funcs/metadata/metadata_ops_test.cc +++ b/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(metadata_state_, nullptr); + auto udf_tester = + px::carnot::udf::UDFTester(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 diff --git a/src/shared/k8s/metadatapb/test_proto.h b/src/shared/k8s/metadatapb/test_proto.h index 9bd6def7bf6..ab05aa34f9c 100644 --- a/src/shared/k8s/metadatapb/test_proto.h +++ b/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 CreateRunningPodUpdatePB() { auto update = std::make_unique(); auto update_proto = absl::Substitute(kResourceUpdateTmpl, "pod_update", kRunningPodUpdatePbTxt); @@ -530,6 +554,33 @@ std::unique_ptr CreateTerminatedDep return update; } +std::unique_ptr CreateRunningNamespaceUpdatePB() { + auto update = std::make_unique(); + 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 CreateTerminatingNamespaceUpdatePB() { + auto update = std::make_unique(); + 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 CreateTerminatedNamespaceUpdatePB() { + auto update = std::make_unique(); + 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