[util] Implement TryGetEnv and UnsetEnv functions

The TryGetEnv function can be used to distinguish the case when a variable is not defined from the case when it has an empty value.
4aab7a2a39f4b9107b0b1d1ef639dc2c41c04f42

3 files changed, 89 insertions, 5 deletions

diff --git a/util/system/env.cpp b/util/system/env.cpp
index ead9b566a5..ea81c3a1d0 100644
--- a/util/system/env.cpp
+++ b/util/system/env.cpp
@@ -1,5 +1,6 @@
 #include "env.h"
 
+#include <util/generic/maybe.h>
 #include <util/generic/string.h>
 #include <util/generic/yexception.h>
 
@@ -21,13 +22,13 @@
  *  - https://a.yandex-team.ru/review/108892/details
  */
 
-TString GetEnv(const TString& key, const TString& def) {
+TMaybe<TString> TryGetEnv(const TString& key) {
 #ifdef _win_
     size_t len = GetEnvironmentVariableA(key.data(), nullptr, 0);
 
     if (len == 0) {
         if (GetLastError() == ERROR_ENVVAR_NOT_FOUND) {
-            return def;
+            return Nothing();
         }
         return TString{};
     }
@@ -45,10 +46,21 @@ TString GetEnv(const TString& key, const TString& def) {
     return TString(buffer.data(), len);
 #else
     const char* env = getenv(key.data());
-    return env ? TString(env) : def;
+    if (!env) {
+        return Nothing();
+    }
+    return TString(env);
 #endif
 }
 
+TString GetEnv(const TString& key, const TString& def) {
+    TMaybe<TString> value = TryGetEnv(key);
+    if (value.Defined()) {
+        return *std::move(value);
+    }
+    return def;
+}
+
 void SetEnv(const TString& key, const TString& value) {
     bool isOk = false;
     int errorCode = 0;
@@ -65,3 +77,17 @@ void SetEnv(const TString& key, const TString& value) {
 #endif
     Y_ENSURE_EX(isOk, TSystemError() << "failed to SetEnv with error-code " << errorCode);
 }
+
+void UnsetEnv(const TString& key) {
+    bool notFound = false;
+#ifdef _win_
+    bool ok = SetEnvironmentVariable(key.c_str(), NULL);
+    notFound = !ok && (GetLastError() == ERROR_ENVVAR_NOT_FOUND);
+#else
+    bool ok = (0 == unsetenv(key.c_str()));
+    #if defined(_darwin_)
+    notFound = !ok && (errno == EINVAL);
+    #endif
+#endif
+    Y_ENSURE_EX(ok || notFound, TSystemError() << "failed to unset environment variable " << key.Quote());
+}
diff --git a/util/system/env.h b/util/system/env.h
index e2ccdd1e95..55f43bd9b8 100644
--- a/util/system/env.h
+++ b/util/system/env.h
@@ -1,5 +1,6 @@
 #pragma once
 
+#include <util/generic/fwd.h>
 #include <util/generic/string.h>
 
 /**
@@ -19,10 +20,27 @@
 TString GetEnv(const TString& key, const TString& def = TString());
 
 /**
+ * Search the environment list provided by the host environment for associated variable.
+ *
+ * @param key               String identifying the name of the environmental variable to look for
+ *
+ * @return                  String that is associated with the matched environment
+ *                          variable or empty optional value if such variable is missing.
+ *
+ * @throws TSystemError     If name of the variable has invalid format
+ *
+ * @note                    Use it only in pair with `SetEnv` as there may be inconsistency
+ *                          in their behaviour otherwise.
+ * @note                    Calls to `TryGetEnv` and `SetEnv` from different threads must be synchronized.
+ * @see                     SetEnv
+ */
+TMaybe<TString> TryGetEnv(const TString& key);
+
+/**
  * Add or change environment variable provided by the host environment.
  *
- * @key         String identifying the name of the environment variable to set or change
- * @value       Value to assign
+ * @param key   String identifying the name of the environment variable to set or change
+ * @param value Value to assign
 
  * @note        Use it only in pair with `GetEnv` as there may be inconsistency in their behaviour
  *              otherwise.
@@ -30,3 +48,15 @@ TString GetEnv(const TString& key, const TString& def = TString());
  * @see         GetEnv
  */
 void SetEnv(const TString& key, const TString& value);
+
+/**
+ * Remove environment variable from the host environment.
+ *
+ * @param key   String identifying the name of the environment variable to remove
+ *
+ * @note        If key does not exist in the environment, then the environment is unchanged,
+ *              and the function returns normally.
+ * @note        Calls to `GetEnv` and `SetEnv` from different threads must be synchronized.
+ * @see         GetEnv
+ */
+void UnsetEnv(const TString& key);
diff --git a/util/system/env_ut.cpp b/util/system/env_ut.cpp
index e03cc01658..6023ce2afd 100644
--- a/util/system/env_ut.cpp
+++ b/util/system/env_ut.cpp
@@ -28,4 +28,32 @@ Y_UNIT_TEST_SUITE(EnvTest) {
         SetEnv(longKey, TString());
         UNIT_ASSERT_VALUES_EQUAL(GetEnv(longKey), TString());
     }
+
+    Y_UNIT_TEST(TryGetEnv) {
+        TString key = "util_TryGetEnv_var";
+        UNIT_ASSERT_NO_EXCEPTION(TryGetEnv(key));
+        SetEnv(key, "value");
+        UNIT_ASSERT(TryGetEnv(key).Defined());
+        UNIT_ASSERT_VALUES_EQUAL(*TryGetEnv(key), "value");
+        UnsetEnv(key);
+        UNIT_ASSERT(TryGetEnv(key).Empty());
+    }
+
+    Y_UNIT_TEST(UnsetEnv) {
+        TString key = "util_UnsetEnv_var";
+        SetEnv(key, "value");
+        UnsetEnv(key);
+        UNIT_ASSERT_VALUES_EQUAL(GetEnv(key, "default_value"), "default_value");
+    }
+
+    Y_UNIT_TEST(UnsetNonexistingEnv) {
+        TString key = "util_UnsetNonexistingEnv_var";
+        UNIT_ASSERT_NO_EXCEPTION(UnsetEnv(key));
+        UNIT_ASSERT_NO_EXCEPTION(UnsetEnv(key));
+    }
+
+    Y_UNIT_TEST(SetEnvInvalidName) {
+        UNIT_ASSERT_EXCEPTION(SetEnv("", "value"), yexception);
+        UNIT_ASSERT_EXCEPTION(SetEnv("A=B", "C=D"), yexception);
+    }
 }