fix: context close on Create/Update/Delete operations (#2038)

* fix: context close on crud operations

* Use the existing error handler decorator to close the context

* Remaining cleanup

---------

Co-authored-by: Iaroslav Zeigerman <zeigerman.ia@gmail.com>

Author: eakmanrq

sqlmesh/cli/__init__.py

@@ -5,6 +5,7 @@
 import click
 from sqlglot.errors import SqlglotError
 
+from sqlmesh.core.context import Context
 from sqlmesh.utils import debug_mode_enabled
 from sqlmesh.utils.concurrency import NodeExecutionFailedError
 from sqlmesh.utils.errors import SQLMeshError
@@ -19,15 +20,21 @@ def error_handler(
     func: t.Callable[..., DECORATOR_RETURN_TYPE]
 ) -> t.Callable[..., DECORATOR_RETURN_TYPE]:
     @wraps(func)
-    def wrapper(*args: t.Any, **kwargs: t.Any) -> DECORATOR_RETURN_TYPE:
+    def wrapper(*args: t.List[t.Any], **kwargs: t.Any) -> DECORATOR_RETURN_TYPE:
+        context_or_obj = args[0]
+        sqlmesh_context = (
+            context_or_obj.obj if isinstance(context_or_obj, click.Context) else context_or_obj
+        )
+        if not isinstance(sqlmesh_context, Context):
+            sqlmesh_context = None
         handler = _debug_exception_handler if debug_mode_enabled() else _default_exception_handler
-        return handler(lambda: func(*args, **kwargs))
+        return handler(sqlmesh_context, lambda: func(*args, **kwargs))
 
     return wrapper
 
 
 def _default_exception_handler(
-    func: t.Callable[[], DECORATOR_RETURN_TYPE]
+    context: t.Optional[Context], func: t.Callable[[], DECORATOR_RETURN_TYPE]
 ) -> DECORATOR_RETURN_TYPE:
     try:
         return func()
@@ -36,11 +43,19 @@ def _default_exception_handler(
         raise click.ClickException(f"Failed processing {ex.node}. {cause}")
     except (SQLMeshError, SqlglotError, ValueError) as ex:
         raise click.ClickException(str(ex))
+    finally:
+        if context:
+            context.close()
 
 
-def _debug_exception_handler(func: t.Callable[[], DECORATOR_RETURN_TYPE]) -> DECORATOR_RETURN_TYPE:
+def _debug_exception_handler(
+    context: t.Optional[Context], func: t.Callable[[], DECORATOR_RETURN_TYPE]
+) -> DECORATOR_RETURN_TYPE:
     try:
         return func()
     except Exception:
         logger.exception("Unhandled exception")
         raise
+    finally:
+        if context:
+            context.close()