From e056af27eb574374a9ecfa369e930ac50f0c015a Mon Sep 17 00:00:00 2001
From: Aleksandr Omyshev <a.omyshev@gmail.com>
Date: Tue, 11 Jul 2023 14:40:10 +0300
Subject: [PATCH] feat: Allow change request timeout and limits (#59)

---
 CHANGES.md                | 12 ++++++++++++
 async_firebase/_config.py | 39 +++++++++++++++++++++++++++++++++++++++
 async_firebase/base.py    | 20 +++++++++++++++++++-
 async_firebase/client.py  |  6 +++++-
 pyproject.toml            |  2 +-
 5 files changed, 76 insertions(+), 3 deletions(-)
 create mode 100644 async_firebase/_config.py

diff --git a/CHANGES.md b/CHANGES.md
index 52c1085..217776d 100644
--- a/CHANGES.md
+++ b/CHANGES.md
@@ -1,5 +1,17 @@
 # Changelog
 
+## 3.2.0
+* ``AsyncFirebaseClient`` empower with advanced features to configure request behaviour such as timeout, or connection pooling.
+Example:
+```python
+
+from async_firebase.client import AsyncFirebaseClient, RequestTimeout
+
+# This will disable timeout
+client = AsyncFirebaseClient(..., request_timeout=RequestTimeout(None))
+client.send(...)
+```
+
 ## 3.1.1
 * [FIX] The push notification could not be sent to topic because ``messages.Message.token`` is declared as required attribute though it should be optional.
 ``messages.Message.token`` turned into Optional attribute.
diff --git a/async_firebase/_config.py b/async_firebase/_config.py
new file mode 100644
index 0000000..f41cb47
--- /dev/null
+++ b/async_firebase/_config.py
@@ -0,0 +1,39 @@
+import typing as t
+from dataclasses import dataclass
+
+
+@dataclass()
+class RequestTimeout:
+    """
+    Request timeout configuration.
+
+    Arguments:
+    timeout: Timeout on all operations eg, read, write, connect.
+
+    Examples:
+        RequestTimeout(None)               # No timeouts.
+        RequestTimeout(5.0)                # 5s timeout on all operations.
+    """
+
+    timeout: t.Optional[float] = None
+
+
+@dataclass()
+class RequestLimits:
+    """
+    Configuration for request limits.
+
+    Attributes:
+    max_connections: The maximum number of concurrent connections that may be established.
+    max_keepalive_connections: Allow the connection pool to maintain keep-alive connections
+        below this point. Should be less than or equal to `max_connections`.
+    keepalive_expiry: Time limit on idle keep-alive connections in seconds.
+    """
+
+    max_connections: t.Optional[int] = None
+    max_keepalive_connections: t.Optional[int] = None
+    keepalive_expiry: t.Optional[int] = None
+
+
+DEFAULT_REQUEST_TIMEOUT = RequestTimeout(timeout=5.0)
+DEFAULT_REQUEST_LIMITS = RequestLimits(max_connections=100, max_keepalive_connections=20, keepalive_expiry=5)
diff --git a/async_firebase/base.py b/async_firebase/base.py
index d4aeea9..e5d9520 100644
--- a/async_firebase/base.py
+++ b/async_firebase/base.py
@@ -16,6 +16,12 @@
 from google.oauth2 import service_account  # type: ignore
 
 import pkg_resources  # type: ignore
+from async_firebase._config import (
+    DEFAULT_REQUEST_LIMITS,
+    DEFAULT_REQUEST_TIMEOUT,
+    RequestLimits,
+    RequestTimeout,
+)
 from async_firebase.messages import FCMBatchResponse, FCMResponse
 from async_firebase.utils import (
     FCMBatchResponseHandler,
@@ -40,6 +46,9 @@ def __init__(
         self,
         credentials: t.Optional[service_account.Credentials] = None,
         scopes: t.Optional[t.List[str]] = None,
+        *,
+        request_timeout: RequestTimeout = DEFAULT_REQUEST_TIMEOUT,
+        request_limits: RequestLimits = DEFAULT_REQUEST_LIMITS,
     ) -> None:
         """
         :param credentials: instance of ``google.oauth2.service_account.Credentials``.
@@ -54,10 +63,15 @@ def __init__(
                 self.creds_from_service_account_info(service_account_info)
 
         :param scopes: user-defined scopes to request during the authorization grant.
+        :param request_timeout: advanced feature that allows to change request timeout.
+        :param request_limits: advanced feature that allows to control the connection pool size.
         """
         self._credentials: service_account.Credentials = credentials
         self.scopes: t.List[str] = scopes or self.SCOPES
 
+        self._request_timeout = request_timeout
+        self._request_limits = request_limits
+
     def creds_from_service_account_info(self, service_account_info: t.Dict[str, str]) -> None:
         """
         Creates a Credentials instance from parsed service account info.
@@ -166,7 +180,11 @@ async def send_request(
         :param content: request content
         :return: HTTP response
         """
-        async with httpx.AsyncClient(base_url=self.BASE_URL) as client:
+        async with httpx.AsyncClient(
+            base_url=self.BASE_URL,
+            timeout=httpx.Timeout(**self._request_timeout.__dict__),
+            limits=httpx.Limits(**self._request_limits.__dict__),
+        ) as client:
             logging.debug(
                 "Requesting POST %s, payload: %s, content: %s, headers: %s",
                 urljoin(self.BASE_URL, self.FCM_ENDPOINT.format(project_id=self._credentials.project_id)),
diff --git a/async_firebase/client.py b/async_firebase/client.py
index 053a820..a3437ee 100644
--- a/async_firebase/client.py
+++ b/async_firebase/client.py
@@ -15,7 +15,11 @@
 
 import httpx
 
-from async_firebase.base import AsyncClientBase
+from async_firebase.base import (  # noqa: F401
+    AsyncClientBase,
+    RequestLimits,
+    RequestTimeout,
+)
 from async_firebase.encoders import aps_encoder
 from async_firebase.messages import (
     AndroidConfig,
diff --git a/pyproject.toml b/pyproject.toml
index e959ee1..f50f1fe 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "async-firebase"
-version = "3.1.1"
+version = "3.2.0"
 description = "Async Firebase Client - a Python asyncio client to interact with Firebase Cloud Messaging in an easy way."
 license = "MIT"
 authors = [