Make Transaction helper response an object (#176)

* Make trx helper response an object

* linting

Author: philkra

tests/integration-tests/helpers_transaction_test.py

@@ -358,11 +358,37 @@ def test_flush_on_error_flag(self):
         before_run = trx.size()
         assert before_run == 3
         response = trx.run(flush_on_error=False)
-        
+
         assert response["status_code"] == 400
         assert response["has_errors"]
         assert trx.size() == before_run
 
         response = trx.run(flush_on_error=True)
         assert response["status_code"] == 400
         assert trx.size() == 0
+
+    def test_object_response_without_errors(self):
+        trx = Transaction(self.client)
+        trx.insert("Posts", self._get_record())
+        trx.insert("Posts", self._get_record())
+        response = trx.run()
+
+        assert response["status_code"] == 200
+        assert response.status_code == response["status_code"]
+        assert response.has_errors == response["has_errors"]
+        assert response.errors == response["errors"]
+        assert response.results == response["results"]
+        assert response.attempts == response["attempts"]
+
+    def test_object_response_with_errors(self):
+        trx = Transaction(self.client)
+        trx.insert("Posts", self._get_record())
+        trx.insert("ChooChoo", self._get_record())
+        response = trx.run()
+
+        assert response["status_code"] == 400
+        assert response.status_code == response["status_code"]
+        assert response.has_errors == response["has_errors"]
+        assert response.errors == response["errors"]
+        assert response.results == response["results"]
+        assert response.attempts == response["attempts"]

xata/helpers.py

@@ -22,6 +22,8 @@
 from datetime import datetime, timezone
 from threading import Lock, Thread
 
+from xata.api_response import ApiResponse
+
 from .client import XataClient
 
 BP_DEFAULT_THREAD_POOL_SIZE = 4
@@ -396,7 +398,7 @@ def get(self, table: str, record_id: str, columns: list[str] = []) -> None:
     def run(self, branch_name: str = None, retry: bool = True, flush_on_error: bool = False) -> dict:
         """
         Commit the transactions. Flushes the operations queue if no error happened.
-        In case of too many connections, hitting rate limits, two extra attempts are taken 
+        In case of too many connections, hitting rate limits, two extra attempts are taken
         with an incremental back off.
 
         :param branch_name: str Override the branch name from the client init
@@ -414,7 +416,9 @@ def run(self, branch_name: str = None, retry: bool = True, flush_on_error: bool
             while attempt < 3 and not r.is_success():
                 wait = attempt * TRX_BACKOFF
                 time.sleep(wait)
-                self.logger.info(f"request {attempt} encountered a 429: too many requests error. will retry in {wait} ms.")
+                self.logger.info(
+                    f"request {attempt} encountered a 429: too many requests error. will retry in {wait} ms."
+                )
                 r = self.client.records().transaction(self.operations, branch_name=branch_name)
                 attempt += 1
 
@@ -423,17 +427,44 @@ def run(self, branch_name: str = None, retry: bool = True, flush_on_error: bool
             self.operations["operations"] = []
 
         # build response
-        return {
-            "status_code": r.status_code,
-            "results": r["results"] if "results" in r else [],
-            "has_errors": True if "errors" in r else False,
-            "errors": r["errors"] if "errors" in r else [],
-            "attempts": attempt
-        }
+        return self.Summary(r, attempt)
 
     def size(self) -> int:
         """
         Get amount of operations in queue
         :returns int
         """
         return len(self.operations["operations"])
+
+    class Summary(dict):
+        """
+        :link https://github.com/xataio/xata-py/issues/170
+        """
+
+        def __init__(self, response: ApiResponse, attempts: int):
+            super()
+            super().__setitem__("status_code", response.status_code)
+            super().__setitem__("results", response.get("results", []))
+            super().__setitem__("errors", response.get("errors", []))
+            super().__setitem__("has_errors", len(response.get("errors", [])) > 0)
+            super().__setitem__("attempts", attempts)
+
+        @property
+        def status_code(self) -> int:
+            return self.__getitem__("status_code")
+
+        @property
+        def attempts(self) -> int:
+            return self.__getitem__("attempts")
+
+        @property
+        def results(self) -> list:
+            return self.__getitem__("results")
+
+        @property
+        def errors(self) -> list:
+            return self.__getitem__("errors")
+
+        @property
+        def has_errors(self) -> bool:
+            return self.__getitem__("has_errors")