pyetp 0.0.46__py3-none-any.whl → 0.0.47__py3-none-any.whl

pyetp/client.py

@@ -6,7 +6,7 @@ import typing as T
 import uuid
 import warnings
 from collections import defaultdict
-from collections.abc import AsyncGenerator
+from collections.abc import AsyncGenerator, Generator
 from types import TracebackType
 
 import numpy as np
@@ -16,138 +16,89 @@ import websockets.client
 from etpproto.connection import CommunicationProtocol, ConnectionType, ETPConnection
 from etpproto.messages import Message, MessageFlags
 from etptypes import ETPModel
-from etptypes.energistics.etp.v12.datatypes.any_array_type import AnyArrayType
-from etptypes.energistics.etp.v12.datatypes.any_logical_array_type import (
+from pydantic import SecretStr
+from xtgeo import RegularSurface
+
+import resqml_objects.v201 as ro
+from energistics.etp.v12.datatypes import (
+    AnyArrayType,
     AnyLogicalArrayType,
+    ArrayOfString,
+    DataValue,
+    ErrorInfo,
+    SupportedDataObject,
+    SupportedProtocol,
+    Uuid,
+    Version,
 )
-from etptypes.energistics.etp.v12.datatypes.array_of_string import ArrayOfString
-from etptypes.energistics.etp.v12.datatypes.data_array_types.data_array_identifier import (
+from energistics.etp.v12.datatypes.data_array_types import (
     DataArrayIdentifier,
-)
-from etptypes.energistics.etp.v12.datatypes.data_array_types.data_array_metadata import (
     DataArrayMetadata,
-)
-from etptypes.energistics.etp.v12.datatypes.data_array_types.get_data_subarrays_type import (
     GetDataSubarraysType,
-)
-from etptypes.energistics.etp.v12.datatypes.data_array_types.put_data_arrays_type import (
     PutDataArraysType,
-)
-from etptypes.energistics.etp.v12.datatypes.data_array_types.put_data_subarrays_type import (
     PutDataSubarraysType,
-)
-from etptypes.energistics.etp.v12.datatypes.data_array_types.put_uninitialized_data_array_type import (
     PutUninitializedDataArrayType,
 )
-from etptypes.energistics.etp.v12.datatypes.data_value import DataValue
-from etptypes.energistics.etp.v12.datatypes.error_info import ErrorInfo
-from etptypes.energistics.etp.v12.datatypes.object.context_info import ContextInfo
-from etptypes.energistics.etp.v12.datatypes.object.context_scope_kind import (
+from energistics.etp.v12.datatypes.object import (
+    ContextInfo,
     ContextScopeKind,
-)
-from etptypes.energistics.etp.v12.datatypes.object.data_object import DataObject
-from etptypes.energistics.etp.v12.datatypes.object.dataspace import Dataspace
-from etptypes.energistics.etp.v12.datatypes.object.relationship_kind import (
+    DataObject,
+    Dataspace,
     RelationshipKind,
+    Resource,
 )
-from etptypes.energistics.etp.v12.datatypes.object.resource import Resource
-from etptypes.energistics.etp.v12.datatypes.supported_data_object import (
-    SupportedDataObject,
-)
-from etptypes.energistics.etp.v12.datatypes.supported_protocol import SupportedProtocol
-from etptypes.energistics.etp.v12.datatypes.uuid import Uuid
-from etptypes.energistics.etp.v12.datatypes.version import Version
-from etptypes.energistics.etp.v12.protocol.core.authorize import Authorize
-from etptypes.energistics.etp.v12.protocol.core.authorize_response import (
+from energistics.etp.v12.protocol.core import (
+    Authorize,
     AuthorizeResponse,
-)
-from etptypes.energistics.etp.v12.protocol.core.close_session import CloseSession
-from etptypes.energistics.etp.v12.protocol.core.open_session import OpenSession
-from etptypes.energistics.etp.v12.protocol.core.protocol_exception import (
+    CloseSession,
+    OpenSession,
     ProtocolException,
+    RequestSession,
 )
-from etptypes.energistics.etp.v12.protocol.core.request_session import RequestSession
-from etptypes.energistics.etp.v12.protocol.data_array.get_data_array_metadata import (
+from energistics.etp.v12.protocol.data_array import (
     GetDataArrayMetadata,
-)
-from etptypes.energistics.etp.v12.protocol.data_array.get_data_array_metadata_response import (
     GetDataArrayMetadataResponse,
-)
-from etptypes.energistics.etp.v12.protocol.data_array.get_data_arrays import (
     GetDataArrays,
-)
-from etptypes.energistics.etp.v12.protocol.data_array.get_data_arrays_response import (
     GetDataArraysResponse,
-)
-from etptypes.energistics.etp.v12.protocol.data_array.get_data_subarrays import (
     GetDataSubarrays,
-)
-from etptypes.energistics.etp.v12.protocol.data_array.get_data_subarrays_response import (
     GetDataSubarraysResponse,
-)
-from etptypes.energistics.etp.v12.protocol.data_array.put_data_arrays import (
     PutDataArrays,
-)
-from etptypes.energistics.etp.v12.protocol.data_array.put_data_arrays_response import (
     PutDataArraysResponse,
-)
-from etptypes.energistics.etp.v12.protocol.data_array.put_data_subarrays import (
     PutDataSubarrays,
-)
-from etptypes.energistics.etp.v12.protocol.data_array.put_data_subarrays_response import (
     PutDataSubarraysResponse,
-)
-from etptypes.energistics.etp.v12.protocol.data_array.put_uninitialized_data_arrays import (
     PutUninitializedDataArrays,
-)
-from etptypes.energistics.etp.v12.protocol.data_array.put_uninitialized_data_arrays_response import (
     PutUninitializedDataArraysResponse,
 )
-from etptypes.energistics.etp.v12.protocol.dataspace.delete_dataspaces import (
+from energistics.etp.v12.protocol.dataspace import (
     DeleteDataspaces,
-)
-from etptypes.energistics.etp.v12.protocol.dataspace.delete_dataspaces_response import (
     DeleteDataspacesResponse,
-)
-from etptypes.energistics.etp.v12.protocol.dataspace.get_dataspaces import GetDataspaces
-from etptypes.energistics.etp.v12.protocol.dataspace.get_dataspaces_response import (
+    GetDataspaces,
     GetDataspacesResponse,
-)
-from etptypes.energistics.etp.v12.protocol.dataspace.put_dataspaces import PutDataspaces
-from etptypes.energistics.etp.v12.protocol.dataspace.put_dataspaces_response import (
+    PutDataspaces,
     PutDataspacesResponse,
 )
-from etptypes.energistics.etp.v12.protocol.discovery.get_resources import GetResources
-from etptypes.energistics.etp.v12.protocol.store.delete_data_objects import (
-    DeleteDataObjects,
+from energistics.etp.v12.protocol.discovery import (
+    GetResources,
+    GetResourcesResponse,
 )
-from etptypes.energistics.etp.v12.protocol.store.delete_data_objects_response import (
+from energistics.etp.v12.protocol.store import (
+    DeleteDataObjects,
     DeleteDataObjectsResponse,
-)
-from etptypes.energistics.etp.v12.protocol.store.get_data_objects import GetDataObjects
-from etptypes.energistics.etp.v12.protocol.store.get_data_objects_response import (
+    GetDataObjects,
     GetDataObjectsResponse,
-)
-from etptypes.energistics.etp.v12.protocol.store.put_data_objects import PutDataObjects
-from etptypes.energistics.etp.v12.protocol.store.put_data_objects_response import (
+    PutDataObjects,
     PutDataObjectsResponse,
 )
-from etptypes.energistics.etp.v12.protocol.transaction.commit_transaction import (
+from energistics.etp.v12.protocol.transaction import (
     CommitTransaction,
-)
-from etptypes.energistics.etp.v12.protocol.transaction.rollback_transaction import (
     RollbackTransaction,
-)
-from etptypes.energistics.etp.v12.protocol.transaction.start_transaction import (
     StartTransaction,
+    StartTransactionResponse,
 )
-from pydantic import SecretStr
-from xtgeo import RegularSurface
-
-import resqml_objects.v201 as ro
 from pyetp import utils_arrays, utils_xml
 from pyetp._version import version
 from pyetp.config import SETTINGS
+from pyetp.errors import ETPTransactionFailure
 from pyetp.uri import DataObjectURI, DataspaceURI
 from resqml_objects import parse_resqml_v201_object, serialize_resqml_v201_object
 
@@ -223,11 +174,11 @@ class ETPClient(ETPConnection):
     # client
     #
 
-    async def send(self, body: ETPModel):
+    async def send(self, body: ETPModel) -> list[ETPModel]:
         correlation_id = await self._send(body)
         return await self._recv(correlation_id)
 
-    async def _send(self, body: ETPModel):
+    async def _send(self, body: ETPModel) -> int:
         msg = Message.get_object_message(body, message_flags=MessageFlags.FINALPART)
         if msg is None:
             raise TypeError(f"{type(body)} not valid etp protocol")
@@ -246,37 +197,11 @@ class ETPClient(ETPConnection):
 
         return msg.header.message_id
 
-    async def _recv(self, correlation_id: int) -> ETPModel:
+    async def _recv(self, correlation_id: int) -> list[ETPModel]:
         assert correlation_id in self._recv_events, (
             "Trying to receive a response on non-existing message"
         )
 
-        def timeout_intervals(etp_timeout):
-            # Local function generating progressively longer timeout intervals.
-
-            # Use the timeout-interval generator from the Python websockets
-            # library.
-            backoff_generator = websockets.client.backoff(
-                initial_delay=5.0, min_delay=5.0, max_delay=20.0
-            )
-
-            # Check if we should never time out.
-            if etp_timeout is None:
-                # This is an infinite generator, so it should never exit.
-                yield from backoff_generator
-                return
-
-            # Generate timeout intervals until we have reached the
-            # `etp_timeout`-threshold.
-            csum = 0.0
-            for d in backoff_generator:
-                yield d
-
-                csum += d
-
-                if csum >= etp_timeout:
-                    break
-
         for ti in timeout_intervals(self.etp_timeout):
             try:
                 # Wait for an event for `ti` seconds.
@@ -327,11 +252,7 @@ class ETPClient(ETPConnection):
                 "Server responded with ETPErrors:", ETPError.from_protos(errors)
             )
 
-        if len(bodies) > 1:
-            logger.warning(f"Recived {len(bodies)} messages, but only expected one")
-
-        # ok
-        return bodies[0]
+        return bodies
 
     @staticmethod
     def _parse_error_info(bodies: list[ETPModel]) -> list[ErrorInfo]:
@@ -416,7 +337,7 @@ class ETPClient(ETPConnection):
 
         logger.debug("Client closed")
 
-    async def close(self):
+    async def close(self) -> None:
         """Closing method that tears down the ETP-connection via the
         `ETPClient.__aexit__`-method, and closes the websockets connection.
         This method should _only_ be used if the user has set up a connection
@@ -473,7 +394,7 @@ class ETPClient(ETPConnection):
 
             return "store"
 
-        msg = await self.send(
+        msgs = await self.send(
             RequestSession(
                 applicationName=self.application_name,
                 applicationVersion=self.application_version,
@@ -497,6 +418,10 @@ class ETPClient(ETPConnection):
                 ),
             )
         )
+
+        assert len(msgs) == 1
+        msg = msgs[0]
+
         assert msg and isinstance(msg, OpenSession)
 
         self.is_connected = True
@@ -510,12 +435,15 @@ class ETPClient(ETPConnection):
     async def authorize(
         self, authorization: str, supplemental_authorization: T.Mapping[str, str] = {}
     ):
-        msg = await self.send(
+        msgs = await self.send(
             Authorize(
                 authorization=authorization,
                 supplementalAuthorization=supplemental_authorization,
             )
         )
+        assert len(msgs) == 1
+        msg = msgs[0]
+
         assert msg and isinstance(msg, AuthorizeResponse)
 
         return msg
@@ -547,8 +475,10 @@ class ETPClient(ETPConnection):
             raise Exception("Max one / in dataspace name")
         return DataspaceURI.from_name(ds)
 
-    def list_objects(self, dataspace_uri: DataspaceURI | str, depth: int = 1) -> list:
-        return self.send(
+    async def list_objects(
+        self, dataspace_uri: DataspaceURI | str, depth: int = 1
+    ) -> GetResourcesResponse:
+        responses = await self.send(
             GetResources(
                 scope=ContextScopeKind.TARGETS_OR_SELF,
                 context=ContextInfo(
@@ -559,6 +489,8 @@ class ETPClient(ETPConnection):
                 ),
             )
         )
+        assert len(responses) == 1
+        return responses[0]
 
     #
     # dataspace
@@ -567,10 +499,17 @@ class ETPClient(ETPConnection):
     async def get_dataspaces(
         self, store_last_write_filter: int = None
     ) -> GetDataspacesResponse:
-        return await self.send(
+        responses = await self.send(
             GetDataspaces(store_last_write_filter=store_last_write_filter)
         )
 
+        assert all(
+            [isinstance(response, GetDataspacesResponse) for response in responses]
+        ), "Expected GetDataspacesResponse"
+        assert len(responses) == 1
+
+        return responses[0]
+
     async def put_dataspaces(
         self,
         legaltags: list[str],
@@ -578,13 +517,13 @@ class ETPClient(ETPConnection):
         owners: list[str],
         viewers: list[str],
         *dataspace_uris: DataspaceURI,
-    ):
+    ) -> dict[str, str]:
         _uris = list(map(DataspaceURI.from_any, dataspace_uris))
         for i in _uris:
             if i.raw_uri.count("/") > 4:  # includes the 3 eml
                 raise Exception("Max one / in dataspace name")
         time = self.timestamp
-        response = await self.send(
+        responses = await self.send(
             PutDataspaces(
                 dataspaces={
                     d.raw_uri: Dataspace(
@@ -607,15 +546,19 @@ class ETPClient(ETPConnection):
                 }
             )
         )
-        assert isinstance(response, PutDataspacesResponse), (
-            "Expected PutDataspacesResponse"
-        )
+        assert all(
+            [isinstance(response, PutDataspacesResponse) for response in responses]
+        ), "Expected PutDataspacesResponse"
 
-        assert len(response.success) == len(dataspace_uris), (
-            f"expected {len(dataspace_uris)} success's"
+        successes = {}
+        for response in responses:
+            successes = {**successes, **response.success}
+
+        assert len(successes) == len(dataspace_uris), (
+            f"expected {len(dataspace_uris)} successes"
         )
 
-        return response.success
+        return successes
 
     async def put_dataspaces_no_raise(
         self,
@@ -624,22 +567,30 @@ class ETPClient(ETPConnection):
         owners: list[str],
         viewers: list[str],
         *dataspace_uris: DataspaceURI,
-    ):
+    ) -> dict[str, str]:
         try:
             return await self.put_dataspaces(
                 legaltags, otherRelevantDataCountries, owners, viewers, *dataspace_uris
             )
         except ETPError:
-            pass
+            return {}
 
-    async def delete_dataspaces(self, *dataspace_uris: DataspaceURI):
+    async def delete_dataspaces(self, *dataspace_uris: DataspaceURI) -> dict[str, str]:
         _uris = list(map(str, dataspace_uris))
 
-        response = await self.send(DeleteDataspaces(uris=dict(zip(_uris, _uris))))
-        assert isinstance(response, DeleteDataspacesResponse), (
-            "Expected DeleteDataspacesResponse"
+        responses = await self.send(DeleteDataspaces(uris=dict(zip(_uris, _uris))))
+        assert all(
+            [isinstance(response, DeleteDataspacesResponse) for response in responses]
+        ), "Expected DeleteDataspacesResponse"
+
+        successes = {}
+        for response in responses:
+            successes = {**successes, **response.success}
+
+        assert len(successes) == len(dataspace_uris), (
+            f"expected {len(dataspace_uris)} successes"
         )
-        return response.success
+        return successes
 
     async def get_data_objects(self, *uris: T.Union[DataObjectURI, str]):
         tasks = []
@@ -647,7 +598,8 @@ class ETPClient(ETPConnection):
             task = self.send(GetDataObjects(uris={str(uri): str(uri)}))
             tasks.append(task)
 
-        responses = await asyncio.gather(*tasks)
+        task_responses = await asyncio.gather(*tasks)
+        responses = [r for tr in task_responses for r in tr]
         assert len(responses) == len(uris)
 
         data_objects = []
@@ -680,7 +632,8 @@ class ETPClient(ETPConnection):
             )
             tasks.append(task)
 
-        responses = await asyncio.gather(*tasks)
+        task_responses = await asyncio.gather(*tasks)
+        responses = [r for tr in task_responses for r in tr]
 
         errors = []
         for response in responses:
@@ -742,12 +695,16 @@ class ETPClient(ETPConnection):
     ):
         _uris = list(map(str, uris))
 
-        response = await self.send(
+        responses = await self.send(
             DeleteDataObjects(
                 uris=dict(zip(_uris, _uris)),
                 prune_contained_objects=prune_contained_objects,
             )
         )
+
+        assert len(responses) == 1
+        response = responses[0]
+
         assert isinstance(response, DeleteDataObjectsResponse), (
             "Expected DeleteDataObjectsResponse"
         )
@@ -755,23 +712,34 @@ class ETPClient(ETPConnection):
         return response.deleted_uris
 
     async def start_transaction(
-        self, dataspace_uri: DataspaceURI, read_only: bool = True
+        self, dataspace_uri: DataspaceURI | str, read_only: bool = True
     ) -> Uuid:
-        trans_id = await self.send(
-            StartTransaction(
-                read_only=read_only, dataspace_uris=[dataspace_uri.raw_uri]
-            )
+        dataspace_uri = str(DataspaceURI.from_any(dataspace_uri))
+        responses = await self.send(
+            StartTransaction(read_only=read_only, dataspace_uris=[dataspace_uri])
         )
-        if trans_id.successful is False:
-            raise Exception(f"Failed starting transaction {dataspace_uri.raw_uri}")
-        # uuid.UUID(bytes=trans_id.transaction_uuid)
-        return Uuid(trans_id.transaction_uuid)
+        assert all(
+            [isinstance(response, StartTransactionResponse) for response in responses]
+        ), "Expected StartTransactionResponse"
+
+        assert len(responses) == 1
+        response = responses[0]
+
+        if not response.successful:
+            raise ETPTransactionFailure(f"Failed starting transaction {dataspace_uri}")
+
+        return response.transaction_uuid
 
     async def commit_transaction(self, transaction_uuid: Uuid):
-        r = await self.send(CommitTransaction(transaction_uuid=transaction_uuid))
-        if r.successful is False:
-            raise Exception(r.failure_reason)
-        return r
+        responses = await self.send(
+            CommitTransaction(transaction_uuid=transaction_uuid)
+        )
+        assert len(responses) == 1
+        response = responses[0]
+
+        if response.successful is False:
+            raise ETPTransactionFailure(response.failure_reason)
+        return response
 
     async def rollback_transaction(self, transaction_id: Uuid):
         return await self.send(RollbackTransaction(transactionUuid=transaction_id))
@@ -877,9 +845,12 @@ class ETPClient(ETPConnection):
     #
 
     async def get_array_metadata(self, *uids: DataArrayIdentifier):
-        response = await self.send(
+        responses = await self.send(
             GetDataArrayMetadata(dataArrays={i.path_in_resource: i for i in uids})
         )
+        assert len(responses) == 1
+        response = responses[0]
+
         assert isinstance(response, GetDataArrayMetadataResponse)
 
         if len(response.array_metadata) != len(uids):
@@ -899,9 +870,13 @@ class ETPClient(ETPConnection):
         ):
             return await self._get_array_chunked(uid)
 
-        response = await self.send(
+        responses = await self.send(
             GetDataArrays(dataArrays={uid.path_in_resource: uid})
         )
+
+        assert len(responses) == 1
+        response = responses[0]
+
         assert isinstance(response, GetDataArraysResponse), (
             "Expected GetDataArraysResponse"
         )
@@ -920,10 +895,13 @@ class ETPClient(ETPConnection):
             path_in_resource=path_in_resource,
         )
 
-        response = await self.send(
+        responses = await self.send(
             GetDataArrayMetadata(data_arrays={dai.path_in_resource: dai}),
         )
 
+        assert len(responses) == 1
+        response = responses[0]
+
         self.assert_response(response, GetDataArrayMetadataResponse)
         assert (
             len(response.array_metadata) == 1
@@ -986,7 +964,12 @@ class ETPClient(ETPConnection):
                 )
                 tasks.append(task)
 
-            responses = await asyncio.gather(*tasks)
+            task_responses = await asyncio.gather(*tasks)
+            responses = [
+                response
+                for task_response in task_responses
+                for response in task_response
+            ]
 
             data_blocks = []
             for i, response in enumerate(responses):
@@ -1021,10 +1004,13 @@ class ETPClient(ETPConnection):
             return data
 
         # Download the full array in one go.
-        response = await self.send(
+        responses = await self.send(
             GetDataArrays(data_arrays={dai.path_in_resource: dai}),
         )
 
+        assert len(responses) == 1
+        response = responses[0]
+
         self.assert_response(response, GetDataArraysResponse)
         assert (
             len(response.data_arrays) == 1
@@ -1056,7 +1042,7 @@ class ETPClient(ETPConnection):
         now = self.timestamp
 
         # Allocate space on server for the array.
-        response = await self.send(
+        responses = await self.send(
             PutUninitializedDataArrays(
                 data_arrays={
                     dai.path_in_resource: PutUninitializedDataArrayType(
@@ -1073,6 +1059,9 @@ class ETPClient(ETPConnection):
             ),
         )
 
+        assert len(responses) == 1
+        response = responses[0]
+
         self.assert_response(response, PutUninitializedDataArraysResponse)
         assert len(response.success) == 1 and dai.path_in_resource in response.success
 
@@ -1122,7 +1111,14 @@ class ETPClient(ETPConnection):
                 tasks.append(task)
 
             # Upload all blocks.
-            responses = await asyncio.gather(*tasks)
+            task_responses = await asyncio.gather(*tasks)
+
+            # Flatten list of responses.
+            responses = [
+                response
+                for task_response in task_responses
+                for response in task_response
+            ]
 
             # Check for successful responses.
             for response in responses:
@@ -1139,7 +1135,7 @@ class ETPClient(ETPConnection):
         etp_array_data = utils_arrays.get_etp_data_array_from_numpy(data)
 
         # Pass entire array in one message.
-        response = await self.send(
+        responses = await self.send(
             PutDataArrays(
                 data_arrays={
                     dai.path_in_resource: PutDataArraysType(
@@ -1150,6 +1146,9 @@ class ETPClient(ETPConnection):
             )
         )
 
+        assert len(responses) == 1
+        response = responses[0]
+
         self.assert_response(response, PutDataArraysResponse)
         assert len(response.success) == 1 and dai.path_in_resource in response.success
 
@@ -1171,7 +1170,7 @@ class ETPClient(ETPConnection):
         if data.nbytes > self.max_array_size:
             return await self._put_array_chunked(uid, data)
 
-        response = await self.send(
+        responses = await self.send(
             PutDataArrays(
                 data_arrays={
                     uid.path_in_resource: PutDataArraysType(
@@ -1182,6 +1181,9 @@ class ETPClient(ETPConnection):
             )
         )
 
+        assert len(responses) == 1
+        response = responses[0]
+
         assert isinstance(response, PutDataArraysResponse), (
             "Expected PutDataArraysResponse"
         )
@@ -1205,9 +1207,13 @@ class ETPClient(ETPConnection):
             starts=starts.tolist(),
             counts=counts.tolist(),
         )
-        response = await self.send(
+        responses = await self.send(
             GetDataSubarrays(dataSubarrays={uid.path_in_resource: payload})
         )
+
+        assert len(responses) == 1
+        response = responses[0]
+
         assert isinstance(response, GetDataSubarraysResponse), (
             "Expected GetDataSubarraysResponse"
         )
@@ -1246,9 +1252,13 @@ class ETPClient(ETPConnection):
             f"{dataarray.data.item.__class__.__name__}"
         )
 
-        response = await self.send(
+        responses = await self.send(
             PutDataSubarrays(dataSubarrays={uid.path_in_resource: payload})
         )
+
+        assert len(responses) == 1
+        response = responses[0]
+
         assert isinstance(response, PutDataSubarraysResponse), (
             "Expected PutDataSubarraysResponse"
         )
@@ -1370,9 +1380,13 @@ class ETPClient(ETPConnection):
                 )
             ),
         )
-        response = await self.send(
+        responses = await self.send(
             PutUninitializedDataArrays(dataArrays={uid.path_in_resource: payload})
         )
+
+        assert len(responses) == 1
+        response = responses[0]
+
         assert isinstance(response, PutUninitializedDataArraysResponse), (
             "Expected PutUninitializedDataArraysResponse"
         )
@@ -1521,7 +1535,12 @@ class etp_connect:
     ) -> None:
         self.uri = uri
         self.data_partition_id = data_partition_id
-        self.authorization = SecretStr(authorization)
+
+        if isinstance(authorization, SecretStr):
+            self.authorization = authorization
+        else:
+            self.authorization = SecretStr(authorization)
+
         self.etp_timeout = etp_timeout
         self.max_message_size = max_message_size
         self.subprotocols = ["etp12.energistics.org"]
@@ -1586,3 +1605,30 @@ class etp_connect:
                 max_message_size=self.max_message_size,
             ) as etp_client:
                 yield etp_client
+
+
+def timeout_intervals(total_timeout: float) -> Generator[float]:
+    # Local function generating progressively longer timeout intervals.
+
+    # Use the timeout-interval generator from the Python websockets
+    # library.
+    backoff_generator = websockets.client.backoff(
+        initial_delay=5.0, min_delay=5.0, max_delay=20.0
+    )
+
+    # Check if we should never time out.
+    if total_timeout is None:
+        # This is an infinite generator, so it should never exit.
+        yield from backoff_generator
+        return
+
+    # Generate timeout intervals until we have reached the
+    # `total_timeout`-threshold.
+    csum = 0.0
+    for d in backoff_generator:
+        yield d
+
+        csum += d
+
+        if csum >= total_timeout:
+            break