Add retry mechanism

Author: Peng Ren

README.md

@@ -49,6 +49,9 @@ PyMongoSQL implements the DB API 2.0 interfaces to provide SQL-like access to Mo
 - **JMESPath** (JSON/Dict Path Query)
   - jmespath >= 1.0.0
 
+- **Tenacity** (Transient Failure Retry)
+    - tenacity >= 9.0.0
+
 ### Optional Dependencies
 
 - **SQLAlchemy** (for ORM/Core support)
@@ -206,6 +209,22 @@ cursor.execute(
 
 Parameters are substituted into the MongoDB filter during execution, providing protection against injection attacks.
 
+### Retry on Transient System Errors
+
+PyMongoSQL supports retrying transient, system-level MongoDB failures (for example connection timeout and reconnect errors) using Tenacity.
+
+```python
+connection = connect(
+    host="mongodb://localhost:27017/database",
+    retry_enabled=True,     # default: True
+    retry_attempts=3,       # default: 3
+    retry_wait_min=0.1,     # default: 0.1 seconds
+    retry_wait_max=1.0,     # default: 1.0 seconds
+)
+```
+
+These options apply to connection ping checks, query/DML command execution, and paginated `getMore` fetches.
+
 ## Supported SQL Features
 
 ### SELECT Statements

pymongosql/connection.py

@@ -13,6 +13,7 @@
 from .cursor import Cursor
 from .error import DatabaseError, OperationalError
 from .helper import ConnectionHelper
+from .retry import RetryConfig, execute_with_retry
 
 _logger = logging.getLogger(__name__)
 
@@ -55,6 +56,10 @@ def __init__(
         if not self._mode and mode:
             self._mode = mode
 
+        # Retry behavior for transient system-level PyMongo failures.
+        # These kwargs are consumed by PyMongoSQL and are not passed to MongoClient.
+        self._retry_config = RetryConfig.from_kwargs(kwargs)
+
         # Extract commonly used parameters for backward compatibility
         self._host = host or "localhost"
         self._port = port or 27017
@@ -109,7 +114,11 @@ def _connect(self) -> None:
             self._client = MongoClient(**self._pymongo_params)
 
             # Test connection
-            self._client.admin.command("ping")
+            execute_with_retry(
+                lambda: self._client.admin.command("ping"),
+                self._retry_config,
+                "initial MongoDB ping",
+            )
 
             # Initialize the database according to explicit parameter or client's default
             # This may raise OperationalError if no database could be determined; allow it to bubble up
@@ -179,6 +188,11 @@ def mode(self) -> str:
         """Get the specified mode"""
         return self._mode
 
+    @property
+    def retry_config(self) -> RetryConfig:
+        """Get retry configuration used for transient system-level errors."""
+        return self._retry_config
+
     def use_database(self, database_name: str) -> None:
         """Switch to a different database"""
         if self._client is None:
@@ -332,15 +346,23 @@ def _start_session(self, **kwargs) -> ClientSession:
         if self._client is None:
             raise OperationalError("No active connection")
 
-        session = self._client.start_session(**kwargs)
+        session = execute_with_retry(
+            lambda: self._client.start_session(**kwargs),
+            self._retry_config,
+            "start MongoDB session",
+        )
         self._session = session
         _logger.info("Started new MongoDB session")
         return session
 
     def _end_session(self) -> None:
         """End the current session (internal method)"""
         if self._session is not None:
-            self._session.end_session()
+            execute_with_retry(
+                lambda: self._session.end_session(),
+                self._retry_config,
+                "end MongoDB session",
+            )
             self._session = None
             _logger.info("Ended MongoDB session")
 
@@ -357,7 +379,11 @@ def _start_transaction(self, **kwargs) -> None:
         if self._session is None:
             raise OperationalError("No active session")
 
-        self._session.start_transaction(**kwargs)
+        execute_with_retry(
+            lambda: self._session.start_transaction(**kwargs),
+            self._retry_config,
+            "start MongoDB transaction",
+        )
         self._in_transaction = True
         self._autocommit = False
         _logger.info("Started MongoDB transaction")
@@ -370,7 +396,11 @@ def _commit_transaction(self) -> None:
         if not self._session.in_transaction:
             raise OperationalError("No active transaction to commit")
 
-        self._session.commit_transaction()
+        execute_with_retry(
+            lambda: self._session.commit_transaction(),
+            self._retry_config,
+            "commit MongoDB transaction",
+        )
         self._in_transaction = False
         self._autocommit = True
         _logger.info("Committed MongoDB transaction")
@@ -383,7 +413,11 @@ def _abort_transaction(self) -> None:
         if not self._session.in_transaction:
             raise OperationalError("No active transaction to abort")
 
-        self._session.abort_transaction()
+        execute_with_retry(
+            lambda: self._session.abort_transaction(),
+            self._retry_config,
+            "abort MongoDB transaction",
+        )
         self._in_transaction = False
         self._autocommit = True
         _logger.info("Aborted MongoDB transaction")
@@ -460,7 +494,11 @@ def test_connection(self) -> bool:
         """Test if the connection is alive"""
         try:
             if self._client:
-                self._client.admin.command("ping")
+                execute_with_retry(
+                    lambda: self._client.admin.command("ping"),
+                    self._retry_config,
+                    "connection health check ping",
+                )
                 return True
             return False
         except Exception as e:

pymongosql/cursor.py

@@ -110,6 +110,7 @@ def execute(self: _T, operation: str, parameters: Optional[Any] = None) -> _T:
                     command_result=result,
                     execution_plan=execution_plan_for_rs,
                     database=self.connection.database,
+                    retry_config=self.connection.retry_config,
                     **self._kwargs,
                 )
             else:
@@ -125,6 +126,7 @@ def execute(self: _T, operation: str, parameters: Optional[Any] = None) -> _T:
                     },
                     execution_plan=stub_plan,
                     database=self.connection.database,
+                    retry_config=self.connection.retry_config,
                     **self._kwargs,
                 )
                 # Store the actual insert result for reference

pymongosql/executor.py

@@ -8,6 +8,7 @@
 
 from .error import DatabaseError, OperationalError, ProgrammingError, SqlSyntaxError
 from .helper import SQLHelper
+from .retry import execute_with_retry
 from .sql.delete_builder import DeleteExecutionPlan
 from .sql.insert_builder import InsertExecutionPlan
 from .sql.parser import SQLParser
@@ -17,6 +18,24 @@
 _logger = logging.getLogger(__name__)
 
 
+def _run_db_command(db: Any, command: Dict[str, Any], connection: Any, operation_name: str) -> Dict[str, Any]:
+    """Run a MongoDB command with optional transaction session and retry policy."""
+    retry_config = getattr(connection, "retry_config", None)
+
+    if connection and connection.session and connection.session.in_transaction:
+        return execute_with_retry(
+            lambda: db.command(command, session=connection.session),
+            retry_config,
+            operation_name,
+        )
+
+    return execute_with_retry(
+        lambda: db.command(command),
+        retry_config,
+        operation_name,
+    )
+
+
 @dataclass
 class ExecutionContext:
     """Manages execution context for a single query"""
@@ -156,11 +175,8 @@ def _execute_find_plan(
 
             _logger.debug(f"Executing MongoDB command: {find_command}")
 
-            # Execute find command with session if in transaction
-            if connection and connection.session and connection.session.in_transaction:
-                result = db.command(find_command, session=connection.session)
-            else:
-                result = db.command(find_command)
+            # Execute find command with retry for transient system-level errors
+            result = _run_db_command(db, find_command, connection, "find command")
 
             # Create command result
             return result
@@ -214,11 +230,13 @@ def _execute_aggregate_plan(
             # Get collection and call aggregate()
             collection = db[execution_plan.collection]
 
-            # Execute aggregate with options
-            cursor = collection.aggregate(pipeline, **options)
-
-            # Convert cursor to list
-            results = list(cursor)
+            # Execute aggregate with retry for transient system-level errors
+            retry_config = getattr(connection, "retry_config", None)
+            results = execute_with_retry(
+                lambda: list(collection.aggregate(pipeline, **options)),
+                retry_config,
+                "aggregate command",
+            )
 
             # Apply additional filters if specified (from WHERE clause)
             if execution_plan.filter_stage:
@@ -420,11 +438,7 @@ def _execute_execution_plan(
 
             _logger.debug(f"Executing MongoDB insert command: {command}")
 
-            # Execute with session if in transaction
-            if connection and connection.session and connection.session.in_transaction:
-                return db.command(command, session=connection.session)
-            else:
-                return db.command(command)
+            return _run_db_command(db, command, connection, "insert command")
         except PyMongoError as e:
             _logger.error(f"MongoDB insert failed: {e}")
             raise DatabaseError(f"Insert execution failed: {e}")
@@ -504,11 +518,7 @@ def _execute_execution_plan(
 
             _logger.debug(f"Executing MongoDB delete command: {command}")
 
-            # Execute with session if in transaction
-            if connection and connection.session and connection.session.in_transaction:
-                return db.command(command, session=connection.session)
-            else:
-                return db.command(command)
+            return _run_db_command(db, command, connection, "delete command")
         except PyMongoError as e:
             _logger.error(f"MongoDB delete failed: {e}")
             raise DatabaseError(f"Delete execution failed: {e}")
@@ -608,11 +618,7 @@ def _execute_execution_plan(
 
             _logger.debug(f"Executing MongoDB update command: {command}")
 
-            # Execute with session if in transaction
-            if connection and connection.session and connection.session.in_transaction:
-                return db.command(command, session=connection.session)
-            else:
-                return db.command(command)
+            return _run_db_command(db, command, connection, "update command")
         except PyMongoError as e:
             _logger.error(f"MongoDB update failed: {e}")
             raise DatabaseError(f"Update execution failed: {e}")

pymongosql/result_set.py

@@ -9,6 +9,7 @@
 from . import STRING
 from .common import CursorIterator
 from .error import DatabaseError, ProgrammingError
+from .retry import RetryConfig, execute_with_retry
 from .sql.query_builder import QueryExecutionPlan
 
 _logger = logging.getLogger(__name__)
@@ -23,6 +24,7 @@ def __init__(
         execution_plan: QueryExecutionPlan = None,
         arraysize: int = None,
         database: Optional[Any] = None,
+        retry_config: Optional[RetryConfig] = None,
         **kwargs,
     ) -> None:
         super().__init__(arraysize=arraysize or self.DEFAULT_FETCH_SIZE, **kwargs)
@@ -39,6 +41,8 @@ def __init__(
         else:
             raise ProgrammingError("command_result must be provided")
 
+        self._retry_config = retry_config
+
         self._execution_plan = execution_plan
         self._is_closed = False
         self._cache_exhausted = False
@@ -108,7 +112,11 @@ def _ensure_results_available(self, count: int = 1) -> None:
                         "getMore": self._cursor_id,
                         "collection": self._execution_plan.collection,
                     }
-                    result = self._database.command(getmore_cmd)
+                    result = execute_with_retry(
+                        lambda: self._database.command(getmore_cmd),
+                        self._retry_config,
+                        "getMore command",
+                    )
 
                     # Extract and process next batch
                     cursor_info = result.get("cursor", {})

pymongosql/retry.py

@@ -0,0 +1,82 @@
+# -*- coding: utf-8 -*-
+import logging
+from dataclasses import dataclass
+from typing import Any, Callable, Optional, Tuple, TypeVar
+
+from pymongo.errors import AutoReconnect, ConnectionFailure, NetworkTimeout, PyMongoError, ServerSelectionTimeoutError
+from tenacity import Retrying, retry_if_exception_type, stop_after_attempt, wait_exponential
+
+_logger = logging.getLogger(__name__)
+_T = TypeVar("_T")
+
+RETRYABLE_SYSTEM_EXCEPTIONS: Tuple[type, ...] = (
+    AutoReconnect,
+    NetworkTimeout,
+    ConnectionFailure,
+    ServerSelectionTimeoutError,
+)
+
+
+@dataclass(frozen=True)
+class RetryConfig:
+    enabled: bool = True
+    attempts: int = 3
+    wait_min: float = 0.1
+    wait_max: float = 1.0
+
+    @classmethod
+    def from_kwargs(cls, kwargs: dict) -> "RetryConfig":
+        enabled = bool(kwargs.pop("retry_enabled", True))
+        attempts = int(kwargs.pop("retry_attempts", 3))
+        wait_min = float(kwargs.pop("retry_wait_min", 0.1))
+        wait_max = float(kwargs.pop("retry_wait_max", 1.0))
+
+        if attempts < 1:
+            attempts = 1
+        if wait_min < 0:
+            wait_min = 0.0
+        if wait_max < wait_min:
+            wait_max = wait_min
+
+        return cls(
+            enabled=enabled,
+            attempts=attempts,
+            wait_min=wait_min,
+            wait_max=wait_max,
+        )
+
+
+def execute_with_retry(
+    operation: Callable[[], _T],
+    retry_config: Optional[RetryConfig],
+    operation_name: str,
+) -> _T:
+    """Execute an operation with retry on transient system-level PyMongo failures."""
+    config = retry_config or RetryConfig(enabled=False, attempts=1, wait_min=0.0, wait_max=0.0)
+
+    if not config.enabled or config.attempts <= 1:
+        return operation()
+
+    def _before_sleep(retry_state: Any) -> None:
+        error = retry_state.outcome.exception() if retry_state.outcome else None
+        _logger.warning(
+            "Retrying %s after transient error (attempt %s/%s): %s",
+            operation_name,
+            retry_state.attempt_number,
+            config.attempts,
+            error,
+        )
+
+    retrying = Retrying(
+        reraise=True,
+        stop=stop_after_attempt(config.attempts),
+        wait=wait_exponential(min=config.wait_min, max=config.wait_max),
+        retry=retry_if_exception_type(RETRYABLE_SYSTEM_EXCEPTIONS),
+        before_sleep=_before_sleep,
+    )
+
+    return retrying(operation)
+
+
+def is_retryable_system_error(error: Exception) -> bool:
+    return isinstance(error, RETRYABLE_SYSTEM_EXCEPTIONS) and isinstance(error, PyMongoError)