dont stomp on original transaction id

whylabs · Nov 21, 2024 · d6150ba · d6150ba
1 parent 520dc0f
commit d6150ba
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 2 deletions.
diff --git a/python/tests/api/writer/test_whylabs_integration.py b/python/tests/api/writer/test_whylabs_integration.py
@@ -649,8 +649,9 @@ def test_transaction_aborted():
     transaction_id = writer.start_transaction()
     status, id = writer.write(result)
     assert status
-    writer._whylabs_client.abort_transaction(transaction_id)
+    writer.abort_transaction()
     status, id = writer.write(result)
+    assert transaction_id == writer._whylabs_client._transaction_id
     with pytest.raises(TransactionAbortedException) as e:
         writer.commit_transaction()
         assert str(e) == "Transaction has been aborted"

diff --git a/python/whylogs/api/writer/whylabs_transaction_writer.py b/python/whylogs/api/writer/whylabs_transaction_writer.py
@@ -67,7 +67,7 @@ def __init__(
         transaction_id: Optional[str] = None,
     ):
         super().__init__(org_id, api_key, dataset_id, api_client, ssl_ca_cert, _timeout_seconds, whylabs_client)
-        transaction_id = transaction_id or self._whylabs_client.get_transaction_id()  # type: ignore
+        transaction_id = transaction_id or self._whylabs_client._transaction_id or self._whylabs_client.get_transaction_id()  # type: ignore
         self._whylabs_client._transaction_id = transaction_id  # type: ignore
         self._aborted: bool = False