From 4d18ae5c652733ca88b7835325c42ee54b8bc607 Mon Sep 17 00:00:00 2001
From: Pradyumn-cloud <pradyumnprasad.567@gmail.com>
Date: Wed, 3 Dec 2025 03:53:57 +0530
Subject: [PATCH] feat: Add @with_logging decorator to automate logging
 setup/teardown. This solves the issue where manual logging setup/teardown was
 error-prone and easy to forget for new public functions.

---
 datashuttle/utils/decorators.py | 89 +++++++++++++++++++++++++++++++++
 1 file changed, 89 insertions(+)

diff --git a/datashuttle/utils/decorators.py b/datashuttle/utils/decorators.py
index 173758844..a0fdbbc9a 100644
--- a/datashuttle/utils/decorators.py
+++ b/datashuttle/utils/decorators.py
@@ -1,4 +1,5 @@
 from functools import wraps
+from typing import Optional
 
 from datashuttle.utils.custom_exceptions import ConfigError
 from datashuttle.utils.utils import log_and_raise_error
@@ -88,3 +89,91 @@ def wrapper(*args, **kwargs):
             return func(*args, **kwargs)
 
     return wrapper
+
+
+def with_logging(
+    command_name: Optional[str] = None,
+    store_in_temp_folder: bool = False,
+    conditional_param: Optional[str] = None,
+):
+    """Automatically handle logging for DataShuttle methods.
+
+    This decorator:
+    1. Starts logging at the beginning of the function
+    2. Captures local variables for logging
+    3. Ensures logging is closed even if an exception occurs
+
+    Parameters
+    ----------
+    command_name
+        Name of the command for logging. If None, uses the function name
+        with underscores replaced by hyphens.
+    store_in_temp_folder
+        If True, store logs in temp folder instead of project logging path.
+    conditional_param
+        Name of parameter that controls whether logging occurs (e.g., "log").
+        If specified and that parameter is False, logging is skipped.
+
+    Examples
+    --------
+    @check_configs_set
+    @with_logging()
+    def upload_rawdata(self, ...):
+        ...
+
+    @with_logging(conditional_param="log")
+    def create_folders(self, ..., log: bool = True):
+        ...
+
+    """
+
+    def decorator(func):
+        @wraps(func)
+        def wrapper(*args, **kwargs):
+            import inspect
+
+            from datashuttle.utils import ds_logger
+
+            # Get the DataShuttle instance (first argument)
+            self = args[0]
+
+            # Check if logging should be skipped based on conditional parameter
+            if conditional_param:
+                sig = inspect.signature(func)
+                bound_args = sig.bind(*args, **kwargs)
+                bound_args.apply_defaults()
+                if not bound_args.arguments.get(conditional_param, True):
+                    # Skip logging - just run the function
+                    return func(*args, **kwargs)
+
+            # Determine command name
+            log_command_name = (
+                command_name
+                if command_name
+                else func.__name__.replace("_", "-")
+            )
+
+            # Capture local variables for logging
+            sig = inspect.signature(func)
+            bound_args = sig.bind(*args, **kwargs)
+            bound_args.apply_defaults()
+            local_vars = dict(bound_args.arguments)
+
+            # Start logging
+            self._start_log(
+                log_command_name,
+                local_vars=local_vars,
+                store_in_temp_folder=store_in_temp_folder,
+            )
+
+            try:
+                # Execute the function
+                result = func(*args, **kwargs)
+                return result
+            finally:
+                # Always close logging, even if exception occurs
+                ds_logger.close_log_filehandler()
+
+        return wrapper
+
+    return decorator