From 73197e5132e910400941eaeca0d55aba7129ff1f 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    | 20 ++++----------------
 osbenchmark/utils/io.py | 29 +++++++++++++++++++++++++++++
 2 files changed, 33 insertions(+), 16 deletions(-)

diff --git a/osbenchmark/paths.py b/osbenchmark/paths.py
index 986c9ec25..cd4ffd801 100644
--- a/osbenchmark/paths.py
+++ b/osbenchmark/paths.py
@@ -22,33 +22,21 @@
 # specific language governing permissions and limitations
 # under the License.
 import os
-import sys
-from osbenchmark.utils.io import ensure_dir
+from osbenchmark.utils.io import ensure_dir, ensure_symlink
 
 def benchmark_confdir():
     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 .benchmark directory exists
     ensure_dir(old_path)
-    ensure_dir(new_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")
 
-
 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..059c7a534 100644
--- a/osbenchmark/utils/io.py
+++ b/osbenchmark/utils/io.py
@@ -29,6 +29,7 @@
 import mmap
 import shutil
 import subprocess
+import sys
 import tarfile
 import zipfile
 import urllib.error
@@ -236,6 +237,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(f"Updated symlink: {link_name} -> {source}")
+            else:
+                logger.info(f"Symlink already correct: {link_name} -> {source}")
+        elif os.path.isdir(link_name):
+            shutil.rmtree(link_name)
+            os.symlink(source, link_name)
+            logger.info(f"Replaced directory with symlink: {link_name} -> {source}")
+        else:
+            os.remove(link_name)
+            os.symlink(source, link_name)
+            logger.info(f"Replaced file with symlink: {link_name} -> {source}")
+    else:
+        os.symlink(source, link_name)
+        logger.info(f"Created symlink: {link_name} -> {source}")
 
 def _zipdir(source_directory, archive):
     for root, _, files in os.walk(source_directory):