feat: add data model for client side metrics

google/cloud/bigtable/data/_metrics/data_model.py

@@ -64,7 +64,23 @@ class OperationType(Enum):
 
 
 class OperationState(Enum):
-    """Enum for the state of the active operation."""
+    """Enum for the state of the active operation.
+
+       ┌───────────┐
+       │  CREATED  │
+       └─────┬─────┘
+             │
+             ▼
+      ┌▶ ACTIVE_ATTEMPT ───┐
+      │      │             │
+      │      ▼             │
+      └─ BETWEEN_ATTEMPTS  │
+             │             │
+             ▼             ▼
+       ┌───────────┐       │
+       │ COMPLETED │ ◀─────┘
+       └───────────┘
+    """
 
     CREATED = 0
     ACTIVE_ATTEMPT = 1
@@ -138,6 +154,7 @@ class ActiveOperationMetric:
 
     op_type: OperationType
     uuid: str = field(default_factory=lambda: str(uuid.uuid4()))
+    state: OperationState = OperationState.CREATED
     # create a default backoff generator, initialized with standard default backoff values
     backoff_generator: TrackedBackoffGenerator = field(
         default_factory=lambda: TrackedBackoffGenerator(
@@ -151,7 +168,6 @@ class ActiveOperationMetric:
     zone: str | None = None
     completed_attempts: list[CompletedAttemptMetric] = field(default_factory=list)
     is_streaming: bool = False  # only True for read_rows operations
-    was_completed: bool = False
     handlers: list[MetricsHandler] = field(default_factory=list)
     # the time it takes to recieve the first response from the server, in nanoseconds
     # attached by interceptor
@@ -186,18 +202,6 @@ def from_context(cls) -> ActiveOperationMetric | None:
             return None
         return op
 
-    @property
-    def state(self) -> OperationState:
-        if self.was_completed:
-            return OperationState.COMPLETED
-        elif self.active_attempt is None:
-            if self.completed_attempts:
-                return OperationState.BETWEEN_ATTEMPTS
-            else:
-                return OperationState.CREATED
-        else:
-            return OperationState.ACTIVE_ATTEMPT
-
     def __post_init__(self):
         """
         Save new instances to contextvars on init
@@ -209,7 +213,8 @@ def start(self) -> None:
         Optionally called to mark the start of the operation. If not called,
         the operation will be started at initialization.
 
-        Assumes operation is in CREATED state.
+        StartState: CREATED
+        EndState: CREATED
         """
         if self.state != OperationState.CREATED:
             return self._handle_error(INVALID_STATE_ERROR.format("start", self.state))
@@ -221,7 +226,8 @@ def start_attempt(self) -> ActiveAttemptMetric | None:
         """
         Called to initiate a new attempt for the operation.
 
-        Assumes operation is in either CREATED or BETWEEN_ATTEMPTS states
+        StartState: CREATED | BETWEEN_ATTEMPTS
+        EndState: ACTIVE_ATTEMPT
         """
         if (
             self.state != OperationState.BETWEEN_ATTEMPTS
@@ -244,6 +250,7 @@ def start_attempt(self) -> ActiveAttemptMetric | None:
             backoff_ns = 0
 
         self.active_attempt = ActiveAttemptMetric(backoff_before_attempt_ns=backoff_ns)
+        self.state = OperationState.ACTIVE_ATTEMPT
         return self.active_attempt
 
     def add_response_metadata(self, metadata: dict[str, bytes | str]) -> None:
@@ -252,7 +259,8 @@ def add_response_metadata(self, metadata: dict[str, bytes | str]) -> None:
 
         If not called, default values for the metadata will be used.
 
-        Assumes operation is in ACTIVE_ATTEMPT state.
+        StartState: ACTIVE_ATTEMPT
+        EndState: ACTIVE_ATTEMPT
 
         Args:
           - metadata: the metadata as extracted from the grpc call
@@ -312,7 +320,8 @@ def end_attempt_with_status(self, status: StatusCode | BaseException) -> None:
         be attempted, `end_with_status` or `end_with_success` should be used
         to finalize the operation along with the attempt.
 
-        Assumes operation is in ACTIVE_ATTEMPT state.
+        StartState: ACTIVE_ATTEMPT
+        EndState: BETWEEN_ATTEMPTS
 
         Args:
           - status: The status of the attempt.
@@ -332,6 +341,7 @@ def end_attempt_with_status(self, status: StatusCode | BaseException) -> None:
         )
         self.completed_attempts.append(complete_attempt)
         self.active_attempt = None
+        self.state = OperationState.BETWEEN_ATTEMPTS
         for handler in self.handlers:
             handler.on_attempt_complete(complete_attempt, self)
 
@@ -340,7 +350,8 @@ def end_with_status(self, status: StatusCode | BaseException) -> None:
         Called to mark the end of the operation. If there is an active attempt,
         end_attempt_with_status will be called with the same status.
 
-        Assumes operation is not already in COMPLETED state.
+        StartState: CREATED | ACTIVE_ATTEMPT | BETWEEN_ATTEMPTS
+        EndState: COMPLETED
 
         Causes on_operation_completed to be called for each registered handler.
 
@@ -356,7 +367,7 @@ def end_with_status(self, status: StatusCode | BaseException) -> None:
         )
         if self.state == OperationState.ACTIVE_ATTEMPT:
             self.end_attempt_with_status(final_status)
-        self.was_completed = True
+        self.state = OperationState.COMPLETED
         finalized = CompletedOperationMetric(
             op_type=self.op_type,
             uuid=self.uuid,
@@ -376,7 +387,8 @@ def end_with_success(self):
         """
         Called to mark the end of the operation with a successful status.
 
-        Assumes operation is not already in COMPLETED state.
+        StartState: CREATED | ACTIVE_ATTEMPT | BETWEEN_ATTEMPTS
+        EndState: COMPLETED
 
         Causes on_operation_completed to be called for each registered handler.
         """

tests/unit/data/_metrics/test_data_model.py

@@ -40,10 +40,10 @@ def test_ctor_defaults(self, mock_monotonic_ns):
         assert metric.cluster_id is None
         assert metric.zone is None
         assert len(metric.completed_attempts) == 0
-        assert metric.was_completed is False
         assert len(metric.handlers) == 0
         assert metric.is_streaming is False
         assert metric.flow_throttling_time_ns == 0
+        assert metric.state == State.CREATED
 
     def test_ctor_explicit(self):
         """
@@ -55,7 +55,7 @@ def test_ctor_explicit(self):
         expected_cluster_id = "cluster"
         expected_zone = "zone"
         expected_completed_attempts = [mock.Mock()]
-        expected_was_completed = True
+        expected_state = State.COMPLETED
         expected_handlers = [mock.Mock()]
         expected_is_streaming = True
         expected_flow_throttling = 12
@@ -65,8 +65,8 @@ def test_ctor_explicit(self):
             active_attempt=expected_active_attempt,
             cluster_id=expected_cluster_id,
             zone=expected_zone,
+            state=expected_state,
             completed_attempts=expected_completed_attempts,
-            was_completed=expected_was_completed,
             handlers=expected_handlers,
             is_streaming=expected_is_streaming,
             flow_throttling_time_ns=expected_flow_throttling,
@@ -77,7 +77,7 @@ def test_ctor_explicit(self):
         assert metric.cluster_id == expected_cluster_id
         assert metric.zone == expected_zone
         assert metric.completed_attempts == expected_completed_attempts
-        assert metric.was_completed == expected_was_completed
+        assert metric.state == expected_state
         assert metric.handlers == expected_handlers
         assert metric.is_streaming == expected_is_streaming
         assert metric.flow_throttling_time_ns == expected_flow_throttling
@@ -99,27 +99,18 @@ def test_state_machine_w_methods(self):
         metric.end_with_success()
         assert metric.state == State.COMPLETED
 
-    def test_state_machine_w_state(self):
+    def test_state_machine(self):
         """
-        Exercise state machine by directly manupulating state variables
-
-        relevant variables are: active_attempt, completed_attempts, was_completed
+        Exercise state machine by moving through states
         """
         metric = self._make_one(mock.Mock())
-        for was_completed_value in [False, True]:
-            metric.was_completed = was_completed_value
-            for active_operation_value in [None, mock.Mock()]:
-                metric.active_attempt = active_operation_value
-                for completed_attempts_value in [[], [mock.Mock()]]:
-                    metric.completed_attempts = completed_attempts_value
-                    if was_completed_value:
-                        assert metric.state == State.COMPLETED
-                    elif active_operation_value is not None:
-                        assert metric.state == State.ACTIVE_ATTEMPT
-                    elif completed_attempts_value:
-                        assert metric.state == State.BETWEEN_ATTEMPTS
-                    else:
-                        assert metric.state == State.CREATED
+        assert metric.state == State.CREATED
+        metric.start_attempt()
+        assert metric.state == State.ACTIVE_ATTEMPT
+        metric.end_attempt_with_status(0)
+        assert metric.state == State.BETWEEN_ATTEMPTS
+        metric.end_with_success()
+        assert metric.state == State.COMPLETED
 
     @pytest.mark.parametrize(
         "method,args,valid_states,error_method_name",
@@ -161,13 +152,7 @@ def test_error_invalid_states(self, method, args, valid_states, error_method_nam
         for state in invalid_states:
             with mock.patch.object(cls, "_handle_error") as mock_handle_error:
                 mock_handle_error.return_value = None
-                metric = self._make_one(mock.Mock())
-                if state == State.ACTIVE_ATTEMPT:
-                    metric.active_attempt = mock.Mock()
-                elif state == State.BETWEEN_ATTEMPTS:
-                    metric.completed_attempts.append(mock.Mock())
-                elif state == State.COMPLETED:
-                    metric.was_completed = True
+                metric = self._make_one(mock.Mock(), state=state)
                 return_obj = getattr(metric, method)(*args)
                 assert return_obj is None
                 assert mock_handle_error.call_count == 1
@@ -286,7 +271,7 @@ def test_add_response_metadata_cbt_header(
         cls = type(self._make_one(mock.Mock()))
         with mock.patch.object(cls, "_handle_error") as mock_handle_error:
             metric = self._make_one(
-                mock.Mock(), cluster_id=start_cluster, zone=start_zone
+                mock.Mock(), cluster_id=start_cluster, zone=start_zone, state=State.ACTIVE_ATTEMPT
             )
             metric.active_attempt = mock.Mock()
             metric.active_attempt.gfe_latency_ns = None
@@ -323,7 +308,7 @@ def test_add_response_metadata_cbt_header_w_error(self, metadata_field):
 
         cls = type(self._make_one(mock.Mock()))
         with mock.patch.object(cls, "_handle_error") as mock_handle_error:
-            metric = self._make_one(mock.Mock())
+            metric = self._make_one(mock.Mock(), state=State.ACTIVE_ATTEMPT)
             metric.cluster_id = None
             metric.zone = None
             metric.active_attempt = mock.Mock()
@@ -371,7 +356,7 @@ def test_add_response_metadata_server_timing_header(
 
         cls = type(self._make_one(mock.Mock()))
         with mock.patch.object(cls, "_handle_error") as mock_handle_error:
-            metric = self._make_one(mock.Mock())
+            metric = self._make_one(mock.Mock(), state=State.ACTIVE_ATTEMPT)
             metric.active_attempt = mock.Mock()
             metric.active_attempt.gfe_latency_ns = None
             metadata = grpc.aio.Metadata()
@@ -477,7 +462,7 @@ def test_end_with_status(self, mock_monotonic_ns):
 
         handlers = [mock.Mock(), mock.Mock()]
         metric = self._make_one(
-            expected_type, handlers=handlers, start_time_ns=expected_start_time
+            expected_type, handlers=handlers, start_time_ns=expected_start_time, state=State.ACTIVE_ATTEMPT
         )
         metric.cluster_id = expected_cluster
         metric.zone = expected_zone
@@ -492,7 +477,6 @@ def test_end_with_status(self, mock_monotonic_ns):
         metric.end_with_status(expected_status)
         # test that ActiveOperation was updated to terminal state
         assert metric.state == State.COMPLETED
-        assert metric.was_completed is True
         assert metric.active_attempt is None
         assert len(metric.completed_attempts) == 1
         # check that finalized operation was passed to handlers
@@ -588,7 +572,6 @@ def test_end_on_empty_operation(self):
         metric = self._make_one(mock.Mock(), handlers=handlers)
         metric.end_with_success()
         assert metric.state == State.COMPLETED
-        assert metric.was_completed is True
         final_op = handlers[0].on_operation_complete.call_args[0][0]
         assert final_op.final_status == StatusCode.OK
         assert final_op.completed_attempts == []