diff --git a/granica/__init__.py b/granica/__init__.py index 8d7b807..15afbe4 100755 --- a/granica/__init__.py +++ b/granica/__init__.py @@ -22,7 +22,7 @@ # Override Session Class class Session(Boto3Session): - def __init__(self): + def __init__(self, single_endpoint_mode=False): super(Session, self).__init__() custom_domain = environ.get("GRANICA_CUSTOM_DOMAIN") @@ -90,6 +90,20 @@ def __init__(self): def client(self, *args, **kwargs): if kwargs.get("service_name") == "s3" or "s3" in args: kwargs["config"] = self._merge_bolt_config(kwargs.get("config")) + single_endpoint_mode = kwargs.get("single_endpoint_mode", False) + self.bolt_router.single_endpoint_mode = single_endpoint_mode + + if single_endpoint_mode: + if not kwargs.get("crunch_endpoint", False): + raise ValueError( + "single_endpoint_mode is True but crunch_endpoint is not provided" + ) + else: + self.bolt_router.crunch_endpoint = kwargs.get("crunch_endpoint") + + kwargs.pop("single_endpoint_mode", None) + kwargs.pop("crunch_endpoint", None) + return self._session.create_client(*args, **kwargs) else: return self._session.create_client(*args, **kwargs) diff --git a/granica/bolt_router.py b/granica/bolt_router.py index 6cb964b..c5ed12d 100755 --- a/granica/bolt_router.py +++ b/granica/bolt_router.py @@ -265,6 +265,9 @@ def __init__( for _ in range(4) ) + self.single_endpoint_mode = False + self.crunch_endpoint = None + if update_interval > 0: @async_function @@ -354,6 +357,9 @@ def _get_endpoints(self): raise e def _select_endpoint(self, method): + if self.single_endpoint_mode: + return self.crunch_endpoint + preferred_order = ( self.PREFERRED_READ_ENDPOINT_ORDER if method in {"GET", "HEAD"}