add auto_transaction connection arg to control whether a transaction is auto-started when creating a cursor (#76)

Author: amaizr

pydataapi/dbapi.py

@@ -96,6 +96,7 @@ class ConnectArgs(BaseModel):
     client: Optional[Any] = None
     rollback_exception: Optional[Type[Exception]] = None
     rds_client: Optional[Any] = None
+    auto_transaction: Optional[bool] = True
 
 
 class Connection:
@@ -113,6 +114,7 @@ def __init__(self, **kwargs: Any) -> None:
             client=connect_args.client,
             rollback_exception=connect_args.rollback_exception,
             rds_client=connect_args.rds_client,
+            auto_transaction=connect_args.auto_transaction,
         )
 
         self.closed = False
@@ -132,7 +134,7 @@ def rollback(self) -> None:
             self._data_api._transaction_id = None
 
     def cursor(self) -> 'Cursor':
-        if not self._data_api.transaction_id:
+        if not self._data_api.transaction_id and self._data_api.auto_transaction:
             self._data_api.begin()
         cursor = Cursor(self._data_api)
         self.cursors.append(cursor)

pydataapi/pydataapi.py

@@ -365,6 +365,7 @@ def __init__(
         client: Optional[boto3.session.Session.client] = None,
         rollback_exception: Optional[Type[Exception]] = None,
         rds_client: Optional[boto3.session.Session.client] = None,
+        auto_transaction: Optional[bool] = None,
     ) -> None:
         if resource_name:
             if resource_arn:
@@ -389,6 +390,7 @@ def __init__(
         )
         self._transaction_status: Optional[str] = None
         self.rollback_exception: Optional[Type[Exception]] = rollback_exception
+        self._auto_transaction: Optional[bool] = auto_transaction
 
     def __enter__(self) -> "DataAPI":
         self.begin()
@@ -418,6 +420,10 @@ def transaction_id(self) -> Optional[str]:
     def transaction_status(self) -> Optional[str]:
         return self._transaction_status
 
+    @property
+    def auto_transaction(self) -> Optional[bool]:
+        return self._auto_transaction
+
     def begin(
         self, database: Optional[str] = None, schema: Optional[str] = None
     ) -> str:

tests/pydataapi/test_dbaapi.py

@@ -353,3 +353,59 @@ class CustomError(Exception):
             )
             raise Exception('error')
     second_mocked_client.rollback_transaction.assert_not_called()
+
+
+def test_execute_select_w_auto_transaction(mocked_client, mocker) -> None:
+    mocked_client.begin_transaction.return_value = {'transactionId': 'abc'}
+    mocked_client.execute_statement.return_value = {
+        'numberOfRecordsUpdated': 0,
+        'records': [[{'longValue': 1}, {'stringValue': 'cat'}]],
+    }
+    data_api = connect(
+        resource_arn='arn:aws:rds:dummy',
+        secret_arn='dummy',
+        database='test',
+        client=mocked_client,
+    )
+    result = data_api.execute("select * from pets")
+    assert result.rowcount == 1
+    assert result.fetchone() == [1, 'cat']
+    assert result.fetchone() is None
+    assert mocked_client.execute_statement.call_args == mocker.call(
+        continueAfterTimeout=True,
+        includeResultMetadata=True,
+        resourceArn='arn:aws:rds:dummy',
+        secretArn='dummy',
+        sql="select * from pets",
+        database='test',
+        transactionId='abc',
+    )
+    mocked_client.begin_transaction.assert_called_once()
+
+
+def test_execute_select_wo_auto_transaction(mocked_client, mocker) -> None:
+    mocked_client.begin_transaction.return_value = {'transactionId': 'abc'}
+    mocked_client.execute_statement.return_value = {
+        'numberOfRecordsUpdated': 0,
+        'records': [[{'longValue': 1}, {'stringValue': 'cat'}]],
+    }
+    data_api = connect(
+        resource_arn='arn:aws:rds:dummy',
+        secret_arn='dummy',
+        database='test',
+        client=mocked_client,
+        auto_transaction=False,
+    )
+    result = data_api.execute("select * from pets")
+    assert result.rowcount == 1
+    assert result.fetchone() == [1, 'cat']
+    assert result.fetchone() is None
+    assert mocked_client.execute_statement.call_args == mocker.call(
+        continueAfterTimeout=True,
+        database='test',
+        includeResultMetadata=True,
+        resourceArn='arn:aws:rds:dummy',
+        secretArn='dummy',
+        sql='select * from pets',
+    )
+    mocked_client.begin_transaction.assert_not_called()