fixed the connection closing scenario

Author: pavanjava

README.md

@@ -178,7 +178,7 @@ Tests do not require a running Qdrant instance — the Qdrant client is mocked.
 pytest tests/ -v
 ```
 
-Expected: **500 tests passing**.
+Expected: **549 tests passing**.
 
 ---
 

src/qql/__init__.py

@@ -36,11 +36,15 @@ def run_query(
 ) -> ExecutionResult:
     """One-shot convenience function kept for backward compatibility.
 
-    Creates a :class:`Connection`, runs one query, and closes the connection.
+    Creates a :class:`Connection`, runs one query, closes the connection, and
+    returns the result.  The underlying ``QdrantClient`` is always released —
+    even if the query raises — so repeated calls do not leak resources.
+
     For workloads that issue multiple queries, prefer :class:`Connection`
     directly — it reuses a single client across all calls::
 
         with Connection(url, secret=secret) as conn:
             result = conn.run_query(query)
     """
-    return Connection(url=url, secret=secret, default_model=default_model).run_query(query)
+    with Connection(url=url, secret=secret, default_model=default_model) as conn:
+        return conn.run_query(query)

tests/test_connection.py

@@ -150,13 +150,43 @@ def test_run_query_delegates_to_connection(self, mocker):
         conn_instance.run_query.return_value = ExecutionResult(
             success=True, message="ok", data=[]
         )
+        # Context-manager protocol: __enter__ returns the mock, __exit__ is a no-op
+        conn_instance.__enter__ = mocker.MagicMock(return_value=conn_instance)
+        conn_instance.__exit__ = mocker.MagicMock(return_value=False)
         conn_cls = mocker.patch("qql.Connection", return_value=conn_instance)
         run_query("SHOW COLLECTIONS", url="http://localhost:6333")
         conn_cls.assert_called_once_with(
             url="http://localhost:6333", secret=None, default_model=None
         )
         conn_instance.run_query.assert_called_once_with("SHOW COLLECTIONS")
 
+    def test_run_query_closes_connection_after_query(self, mocker):
+        """run_query() must call close() — it must not leak the QdrantClient."""
+        mock_client = mocker.MagicMock()
+        mocker.patch("qdrant_client.QdrantClient", return_value=mock_client)
+        mock_executor = mocker.MagicMock()
+        mock_executor.execute.return_value = ExecutionResult(
+            success=True, message="ok", data=[]
+        )
+        mocker.patch("qql.connection.Executor", return_value=mock_executor)
+        run_query("SHOW COLLECTIONS", url="http://localhost:6333")
+        # close() must have been called exactly once
+        mock_client.close.assert_called_once()
+
+    def test_run_query_closes_connection_even_when_query_raises(self, mocker):
+        """run_query() must call close() even if the query throws."""
+        mock_client = mocker.MagicMock()
+        mocker.patch("qdrant_client.QdrantClient", return_value=mock_client)
+        # Make the query raise a runtime error (e.g. collection not found)
+        from qql.exceptions import QQLRuntimeError
+        mock_executor = mocker.MagicMock()
+        mock_executor.execute.side_effect = QQLRuntimeError("collection 'x' does not exist")
+        mocker.patch("qql.connection.Executor", return_value=mock_executor)
+        with pytest.raises(QQLRuntimeError):
+            run_query("SEARCH x SIMILAR TO 'q' LIMIT 5", url="http://localhost:6333")
+        # close() still called
+        mock_client.close.assert_called_once()
+
     def test_run_query_invalid_syntax_still_raises(self, mocker):
         mocker.patch("qdrant_client.QdrantClient")
         with pytest.raises(QQLSyntaxError):