Make sunset warning handling configurable

src/kognic/auth/__init__.py

@@ -3,8 +3,12 @@
 from logging import NullHandler
 from pathlib import Path
 
+from kognic.auth._sunset import SunsetHandler, default_sunset_handler
+
 logging.getLogger(__name__).addHandler(NullHandler())
 
+__all__ = ["SunsetHandler", "default_sunset_handler"]
+
 try:
     from ._version import __version__
 except ImportError:

src/kognic/auth/_sunset.py

@@ -2,40 +2,69 @@
 
 import logging
 from datetime import datetime, timezone
-from typing import TYPE_CHECKING, Optional, Union
+from typing import TYPE_CHECKING, Callable, Optional, Union
 
 if TYPE_CHECKING:
     from ._protocols import Response, Url
 
 SUNSET_HEADER = "sunset-date"
-SUNSET_DIFF_THRESHOLD = 14 * 60 * 60 * 24  # two weeks
 
 # Expected formats of sunset date
 DATETIME_FMT = "%Y-%m-%dT%H:%M:%S.%fZ"
 DATETIME_FMT_NO_MICRO = "%Y-%m-%dT%H:%M:%SZ"
 
 logger = logging.getLogger(__name__)
 
+SunsetHandler = Callable[[datetime, str, str], None]
 
-def handle_sunset(response: "Response") -> None:
-    """Check for Sunset header and log warnings/errors.
+
+def default_sunset_handler(threshold_days: int = 14) -> SunsetHandler:
+    """Return a sunset handler that logs a warning or error based on time until sunset.
+
+    Args:
+        threshold_days: Days remaining before the sunset date at which the log level
+            escalates from warning to error. Defaults to 14.
+
+    Returns:
+        A callable ``(sunset_date, method, url) -> None`` that logs the deprecation notice.
+
+    Example::
+
+        from kognic.auth import default_sunset_handler
+        client = BaseApiClient(sunset_handler=default_sunset_handler(threshold_days=30))
+    """
+    threshold_seconds = threshold_days * 60 * 60 * 24
+
+    def handler(sunset_date: datetime, method: str, url: str) -> None:
+        now = datetime.now(tz=timezone.utc)
+        diff = sunset_date - now
+        log_method = logger.warning if diff.total_seconds() > threshold_seconds else logger.error
+        log_method(
+            f"Endpoint has been deprecated and will be removed at {sunset_date}. Please update your client. "
+            f"Endpoint: {method} {url}"
+        )
+
+    return handler
+
+
+_default_handler: SunsetHandler = default_sunset_handler()
+
+
+def handle_sunset(response: "Response", handler: Optional[SunsetHandler] = _default_handler) -> None:
+    """Check for Sunset header and invoke the handler if present.
 
     Args:
         response: The HTTP response object (requests.Response or httpx.Response)
+        handler: Callable invoked with ``(sunset_date, method, url)``. Pass ``None`` to disable.
     """
+    if handler is None:
+        return
     sunset_string = response.headers.get(SUNSET_HEADER)
     sunset_date = _parse_date(sunset_string) if sunset_string else None
     if not sunset_date:
-        return None
-
-    now = datetime.now(tz=timezone.utc)
-    diff = sunset_date - now
+        return
 
-    log_method = logger.warning if diff.total_seconds() > SUNSET_DIFF_THRESHOLD else logger.error
-    log_method(
-        f"Endpoint has been deprecated and will be removed at {sunset_date}. Please update your client. "
-        f"Endpoint: {response.request.method} {_parse_url(response.request.url)}"
-    )
+    handler(sunset_date, response.request.method, _parse_url(response.request.url))
 
 
 def _parse_date(date: str) -> Optional[datetime]:

src/kognic/auth/httpx/base_client.py

@@ -13,14 +13,16 @@
 import httpx
 
 from kognic.auth import DEFAULT_HOST, DEFAULT_TOKEN_ENDPOINT_RELPATH
-from kognic.auth._sunset import handle_sunset
+from kognic.auth._sunset import SunsetHandler, default_sunset_handler, handle_sunset
 from kognic.auth._user_agent import get_user_agent
 from kognic.auth.env_config import DEFAULT_ENV_CONFIG_FILE_PATH, load_kognic_env_config
 from kognic.auth.httpx.async_client import HttpxAuthAsyncClient
 from kognic.auth.serde import serialize_body
 
 logger = logging.getLogger(__name__)
 
+_DEFAULT_SUNSET_HANDLER: SunsetHandler = default_sunset_handler()
+
 
 def _handle_http_error(resp: httpx.Response):
     """Try to get the error message from the response and raise with that message."""
@@ -66,6 +68,7 @@ def __init__(
         auth_token_endpoint: str = DEFAULT_TOKEN_ENDPOINT_RELPATH,
         client_name: Optional[str] = "auto",
         json_serializer: Callable[[Any], Any] = serialize_body,
+        sunset_handler: Optional[SunsetHandler] = _DEFAULT_SUNSET_HANDLER,
         **kwargs,
     ):
         """Initialize the async API client.
@@ -76,6 +79,8 @@ def __init__(
             auth_token_endpoint: Relative path to token endpoint
             client_name: Name added to User-Agent. Use "auto" for class name, None for no name.
             json_serializer: Callable to serialize request bodies. Defaults to serialize_body.
+            sunset_handler: Callable invoked with ``(sunset_date, method, url)`` when a sunset header
+                is detected. Defaults to logging a warning or error. Pass ``None`` to disable.
             **kwargs: Additional arguments passed to the underlying httpx client (e.g. timeout, verify).
         """
         if client_name == "auto":
@@ -120,7 +125,7 @@ async def call_with_simple_retry(attempts):
 
             resp = await call_with_simple_retry(3)
 
-            handle_sunset(resp)
+            handle_sunset(resp, sunset_handler)
             _handle_http_error(resp)
             return resp
 

src/kognic/auth/requests/base_client.py

@@ -17,7 +17,7 @@
 from requests.adapters import HTTPAdapter, Retry
 
 from kognic.auth import DEFAULT_HOST, DEFAULT_TOKEN_ENDPOINT_RELPATH
-from kognic.auth._sunset import handle_sunset
+from kognic.auth._sunset import SunsetHandler, default_sunset_handler, handle_sunset
 from kognic.auth._user_agent import get_user_agent
 from kognic.auth.credentials_parser import ANY_AUTH_TYPE, resolve_credentials
 from kognic.auth.env_config import DEFAULT_ENV_CONFIG_FILE_PATH, load_kognic_env_config
@@ -31,10 +31,12 @@
 
 DEFAULT_RETRY = Retry(total=3, connect=3, read=3, backoff_factor=0.5, status_forcelist=[502, 503, 504])
 
+_DEFAULT_SUNSET_HANDLER: SunsetHandler = default_sunset_handler()
 
-def _check_response(resp: requests.Response):
+
+def _check_response(resp: requests.Response, sunset_handler: Optional[SunsetHandler] = _DEFAULT_SUNSET_HANDLER):
     """Handle sunset headers and raise for status with enhanced error messages."""
-    handle_sunset(resp)
+    handle_sunset(resp, sunset_handler)
     try:
         resp.raise_for_status()
     except requests.HTTPError as e:
@@ -56,7 +58,11 @@ def _set_session_user_agent(session: Session, client_name: Optional[str] = None)
     session.headers["User-Agent"] = get_user_agent(f"requests/{requests.__version__}", client_name)
 
 
-def _monkey_patch_send(session: Session, json_serializer: Callable[[Any], Any]):
+def _monkey_patch_send(
+    session: Session,
+    json_serializer: Callable[[Any], Any],
+    sunset_handler: Optional[SunsetHandler] = _DEFAULT_SUNSET_HANDLER,
+):
     """
     Monkey patch to serialize JSON and validate paths
     :param session:
@@ -82,7 +88,7 @@ def prepare_request(req, *args, **kwargs):
 
     def send_request(req, *args, **kwargs):
         resp = vanilla_send(req, *args, **kwargs)
-        _check_response(resp)
+        _check_response(resp, sunset_handler)
         return resp
 
     session.send = send_request
@@ -98,6 +104,7 @@ def create_session(
     initial_token: Optional[dict] = None,
     on_token_updated: Optional[Callable[[dict], None]] = None,
     token_provider: Optional[RequestsAuthSession] = None,
+    sunset_handler: Optional[SunsetHandler] = _DEFAULT_SUNSET_HANDLER,
 ) -> Session:
     """Create a requests session with enhancements.
 
@@ -117,6 +124,8 @@ def create_session(
         on_token_updated: Callback invoked with the new token dict whenever a fresh token is fetched.
         token_provider: Explicit token provider to use. When given, auth/initial_token/on_token_updated
             are ignored. Multiple sessions sharing one provider share the same token lifecycle.
+        sunset_handler: Callable invoked with ``(sunset_date, method, url)`` when a sunset header is
+            detected. Defaults to logging a warning or error. Pass ``None`` to disable.
 
     Returns:
         Configured requests Session
@@ -133,7 +142,7 @@ def create_session(
     session = requests.Session()
     session.auth = KognicBearerAuth(token_provider)
     _set_session_user_agent(session, client_name)
-    _monkey_patch_send(session, json_serializer)
+    _monkey_patch_send(session, json_serializer, sunset_handler)
     session.mount("http://", HTTPAdapter(max_retries=DEFAULT_RETRY))
     session.mount("https://", HTTPAdapter(max_retries=DEFAULT_RETRY))
     return session
@@ -235,6 +244,7 @@ def __init__(
         json_serializer: Callable[[Any], Any] = serialize_body,
         token_provider: Optional[RequestsAuthSession] = None,
         token_cache: Optional[TokenCache] = None,
+        sunset_handler: Optional[SunsetHandler] = _DEFAULT_SUNSET_HANDLER,
     ):
         """Initialize the API client.
 
@@ -248,6 +258,8 @@ def __init__(
                 provider is looked up (or created) by credentials + auth_host.
             token_cache: Token cache for cross-process token persistence. When given, a valid cached
                 token is injected on startup and new tokens are saved automatically.
+            sunset_handler: Callable invoked with ``(sunset_date, method, url)`` when a sunset header
+                is detected. Defaults to logging a warning or error. Pass ``None`` to disable.
         """
         self._session: Optional[Session] = None
         self._auth = auth
@@ -256,6 +268,7 @@ def __init__(
         self._json_serializer = json_serializer
         self._token_provider = token_provider
         self._token_cache = token_cache
+        self._sunset_handler = sunset_handler
         self._lock = Lock()
 
         if client_name == "auto":
@@ -278,6 +291,7 @@ def session(self) -> Session:
                         token_provider=provider,
                         client_name=self._client_name,
                         json_serializer=self._json_serializer,
+                        sunset_handler=self._sunset_handler,
                     )
         return self._session