From ae54808df93a936b0f06ab7cfa16f2a3a7bfc572 Mon Sep 17 00:00:00 2001
From: Dimitri Yatsenko <dimitri.yatsenko@gmail.com>
Date: Thu, 8 Jan 2026 23:29:08 -0600
Subject: [PATCH] feat: add context manager support to Connection class

Add __enter__ and __exit__ methods to Connection for use with Python's
`with` statement. This enables automatic connection cleanup, particularly
useful for serverless environments (AWS Lambda, Cloud Functions).

Usage:
    with dj.Connection(host, user, password) as conn:
        schema = dj.schema('my_schema', connection=conn)
        # perform operations
    # connection automatically closed

Closes #1081

Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>
---
 src/datajoint/connection.py          | 39 ++++++++++++++++++++++++++++
 tests/integration/test_connection.py | 30 +++++++++++++++++++++
 2 files changed, 69 insertions(+)

diff --git a/src/datajoint/connection.py b/src/datajoint/connection.py
index 394952886..257034fec 100644
--- a/src/datajoint/connection.py
+++ b/src/datajoint/connection.py
@@ -283,6 +283,45 @@ def close(self) -> None:
         """Close the database connection."""
         self._conn.close()
 
+    def __enter__(self) -> "Connection":
+        """
+        Enter context manager.
+
+        Returns
+        -------
+        Connection
+            This connection object.
+
+        Examples
+        --------
+        >>> with dj.Connection(host, user, password) as conn:
+        ...     schema = dj.schema('my_schema', connection=conn)
+        ...     # perform operations
+        ... # connection automatically closed
+        """
+        return self
+
+    def __exit__(self, exc_type, exc_val, exc_tb) -> bool:
+        """
+        Exit context manager and close connection.
+
+        Parameters
+        ----------
+        exc_type : type or None
+            Exception type if an exception was raised.
+        exc_val : Exception or None
+            Exception instance if an exception was raised.
+        exc_tb : traceback or None
+            Traceback if an exception was raised.
+
+        Returns
+        -------
+        bool
+            False to propagate exceptions.
+        """
+        self.close()
+        return False
+
     def register(self, schema) -> None:
         """
         Register a schema with this connection.
diff --git a/tests/integration/test_connection.py b/tests/integration/test_connection.py
index 8a30d4a46..ff3940587 100644
--- a/tests/integration/test_connection.py
+++ b/tests/integration/test_connection.py
@@ -46,6 +46,36 @@ def test_dj_connection_class(connection_test):
     assert connection_test.is_connected
 
 
+def test_connection_context_manager(db_creds_test):
+    """
+    Connection should support context manager protocol for automatic cleanup.
+    """
+    # Test basic context manager usage
+    with dj.Connection(**db_creds_test) as conn:
+        assert conn.is_connected
+        # Verify we can use the connection
+        result = conn.query("SELECT 1").fetchone()
+        assert result[0] == 1
+
+    # Connection should be closed after exiting context
+    assert not conn.is_connected
+
+
+def test_connection_context_manager_exception(db_creds_test):
+    """
+    Connection should close even when exception is raised inside context.
+    """
+    conn = None
+    with pytest.raises(ValueError):
+        with dj.Connection(**db_creds_test) as conn:
+            assert conn.is_connected
+            raise ValueError("Test exception")
+
+    # Connection should still be closed after exception
+    assert conn is not None
+    assert not conn.is_connected
+
+
 def test_persistent_dj_conn(db_creds_root):
     """
     conn() method should provide persistent connection across calls.