From cfb86c2c0a5f141e9999837941effdd5d466e69e Mon Sep 17 00:00:00 2001
From: Michael Oviedo <mikeovi@amazon.com>
Date: Tue, 7 Jan 2025 19:26:27 +0000
Subject: [PATCH] add new ensure_symlink util function to help prevent
 FileExists errors

Signed-off-by: Michael Oviedo <mikeovi@amazon.com>
---
 osbenchmark/paths.py    | 42 +++++++++++++++++++++++++----------------
 osbenchmark/utils/io.py | 28 +++++++++++++++++++++++++++
 2 files changed, 54 insertions(+), 16 deletions(-)

diff --git a/osbenchmark/paths.py b/osbenchmark/paths.py
index 986c9ec25..d40205a55 100644
--- a/osbenchmark/paths.py
+++ b/osbenchmark/paths.py
@@ -21,33 +21,43 @@
 # KIND, either express or implied.  See the License for the
 # specific language governing permissions and limitations
 # under the License.
+import logging
 import os
 import sys
-from osbenchmark.utils.io import ensure_dir
+from osbenchmark.utils.io import ensure_dir, ensure_symlink
 
 def benchmark_confdir():
+    logger = logging.getLogger("opensearch_benchmark")
     default_home = os.path.expanduser("~")
     old_path = os.path.join(default_home, ".benchmark")
     new_path = os.path.join(default_home, ".osb")
 
-    # ensure both directories exist
-    ensure_dir(old_path)
-    ensure_dir(new_path)
+    try:
+        # Ensure .benchmark directory exists
+        ensure_dir(old_path)
+        logger.info(f"Ensured directory exists: {old_path}")
 
-    # Create symlink from .osb to .benchmark if it doesn't exist
-    if not os.path.islink(new_path):
-        try:
-            os.symlink(old_path, new_path, target_is_directory=True)
-        except OSError as e:
-            error_message = (
-                f"OSError: Failed to create symlink from {new_path} to {old_path}\n"
-                f"Error type: {type(e).__name__}\n"
-                f"Error message: {str(e)}\n"
-            )
-            print(error_message, file=sys.stderr)
+        # Ensure symlink from .osb to .benchmark
+        ensure_symlink(old_path, new_path)
 
-    return os.path.join(os.getenv("BENCHMARK_HOME", default_home), ".osb")
+        final_path = os.path.join(os.getenv("BENCHMARK_HOME", default_home), ".osb")
 
+        return final_path
+
+    except Exception as e:
+        error_message = (
+            f"Error in benchmark_confdir:\n"
+            f"Error type: {type(e).__name__}\n"
+            f"Error message: {str(e)}\n"
+            f"Current user: {os.getlogin()}\n"
+            f"Current working directory: {os.getcwd()}\n"
+            f"Python version: {sys.version}\n"
+            f"Operating system: {sys.platform}\n"
+            f"Permissions of {old_path}: {oct(os.stat(old_path).st_mode) if os.path.exists(old_path) else 'N/A'}\n"
+            f"Permissions of parent of {new_path}: {oct(os.stat(os.path.dirname(new_path)).st_mode)}"
+        )
+        print(error_message)
+        raise
 
 def benchmark_root():
     return os.path.dirname(os.path.realpath(__file__))
diff --git a/osbenchmark/utils/io.py b/osbenchmark/utils/io.py
index 2b740f1ec..e1b596709 100644
--- a/osbenchmark/utils/io.py
+++ b/osbenchmark/utils/io.py
@@ -236,6 +236,34 @@ def ensure_dir(directory, mode=0o777):
     if directory:
         os.makedirs(directory, mode, exist_ok=True)
 
+def ensure_symlink(source, link_name):
+    """
+    Ensure that a symlink exists from link_name to source.
+    If link_name already exists, it will be updated or replaced as necessary.
+
+    :param source: The target of the symlink
+    :param link_name: The path where the symlink should be created
+    """
+    logger = logging.getLogger(__name__)
+    if os.path.exists(link_name):
+        if os.path.islink(link_name):
+            if os.readlink(link_name) != source:
+                os.remove(link_name)
+                os.symlink(source, link_name)
+                logger.info("Updated symlink: %s -> %s", link_name, source)
+            else:
+                logger.info("Symlink already correct: %s -> %s", link_name, source)
+        elif os.path.isdir(link_name):
+            shutil.rmtree(link_name)
+            os.symlink(source, link_name)
+            logger.info("Replaced directory with symlink: %s -> %s", link_name, source)
+        else:
+            os.remove(link_name)
+            os.symlink(source, link_name)
+            logger.info("Replaced file with symlink: %s -> %s", link_name, source)
+    else:
+        os.symlink(source, link_name)
+        logger.info("Created symlink: %s -> %s", link_name, source)
 
 def _zipdir(source_directory, archive):
     for root, _, files in os.walk(source_directory):