Eliminate use of status objects when creating balance trades

robcarver17 · Jul 31, 2023 · 4ab763e · 4ab763e
1 parent 44d004d
commit 4ab763e
Showing 1 changed file with 32 additions and 22 deletions.
diff --git a/sysexecution/stack_handler/balance_trades.py b/sysexecution/stack_handler/balance_trades.py
@@ -1,4 +1,4 @@
-from syscore.constants import success, failure
+from sysexecution.order_stacks.order_stack import failureWithRollback
 from sysexecution.orders.named_order_objects import missing_order
 from sysexecution.stack_handler.fills import stackHandlerForFills
 from sysexecution.orders.instrument_orders import instrumentOrder
@@ -23,20 +23,15 @@ def create_balance_trade(self, broker_order: brokerOrder):
 
         log.debug("Putting balancing trades on stacks")
 
-        (
-            result,
-            instrument_order_id,
-            contract_order_id,
-            broker_order_id,
-        ) = self.put_balance_trades_on_stack(
-            instrument_order, contract_order, broker_order
-        )
-
-        if result is failure:
-            log.error("Something went wrong, rolling back")
-            self.rollback_balance_trades(
-                instrument_order_id, contract_order_id, broker_order_id
+        try:
+            (
+                instrument_order_id,
+                contract_order_id,
+                broker_order_id,
+            ) = self.put_balance_trades_on_stack(
+                instrument_order, contract_order, broker_order
             )
+        except failureWithRollback:
             return None
 
         contract_order.order_id = contract_order_id
@@ -55,8 +50,6 @@ def create_balance_trade(self, broker_order: brokerOrder):
             instrument_order_id, treat_inactive_as_complete=True
         )
 
-        return success
-
     def put_balance_trades_on_stack(
         self,
         instrument_order: instrumentOrder,
@@ -76,7 +69,8 @@ def put_balance_trades_on_stack(
             log.error(
                 "Couldn't add balancing instrument trade error condition %s" % str(e)
             )
-            return failure, missing_order, missing_order, missing_order
+            log.error("Nothing to roll back")
+            raise failureWithRollback from e
 
         try:
             contract_order.parent = instrument_order_id
@@ -85,7 +79,11 @@ def put_balance_trades_on_stack(
             log.error(
                 "Couldn't add balancing contract trade error condition %s " % str(e)
             )
-            return failure, instrument_order_id, missing_order, missing_order
+            log.error("Rolling back")
+            self.rollback_balance_trades(
+                instrument_order_id, missing_order, missing_order
+            )
+            raise failureWithRollback from e
 
         try:
             self.instrument_stack.add_children_to_order_without_existing_children(
@@ -94,26 +92,38 @@ def put_balance_trades_on_stack(
         except Exception as e:
 
             log.error("Couldn't add children to instrument order error %s" % str(e))
-            return failure, instrument_order_id, contract_order_id, missing_order
+            log.error("Rolling back")
+            self.rollback_balance_trades(
+                instrument_order_id, contract_order_id, missing_order
+            )
+            raise failureWithRollback from e
 
         broker_order.parent = contract_order_id
         try:
             broker_order_id = self.broker_stack.put_order_on_stack(broker_order)
         except Exception as e:
             log.error("Couldn't add balancing broker trade error condition %s" % str(e))
-            return failure, instrument_order_id, contract_order_id, missing_order
+            log.error("Rolling back")
+            self.rollback_balance_trades(
+                instrument_order_id, contract_order_id, missing_order
+            )
+            raise failureWithRollback from e
 
         try:
             self.contract_stack.add_children_to_order_without_existing_children(
                 contract_order_id, [broker_order_id]
             )
         except Exception as e:
             log.error("Couldn't add children to contract order exception %s" % str(e))
-            return failure, instrument_order_id, contract_order_id, broker_order_id
+            log.error("Rolling back")
+            self.rollback_balance_trades(
+                instrument_order_id, contract_order_id, broker_order_id
+            )
+            raise failureWithRollback from e
 
         log.debug("All balancing trades added to stacks")
 
-        return success, instrument_order_id, contract_order_id, broker_order_id
+        return instrument_order_id, contract_order_id, broker_order_id
 
     def rollback_balance_trades(
         self, instrument_order_id: int, contract_order_id: int, broker_order_id: int