From 65f4df6a309cfdd0c6cac2efab1875132743c9c6 Mon Sep 17 00:00:00 2001 From: Andrew Moore Date: Tue, 13 May 2025 15:04:22 +0100 Subject: [PATCH] feat: add pool_maxsize parameter to connection managers Introduces the `pool_maxsize` parameter to `KeycloakOpenID` and `KeycloakAdmin` classes, allowing control over the underlying connection pool size in the `ConnectionManager`. Adds corresponding tests to verify the parameter is correctly passed and stored for both synchronous and asynchronous clients. --- docs/source/modules/admin.rst | 4 +++- docs/source/modules/openid_client.rst | 3 ++- src/keycloak/connection.py | 25 ++++++++++++++++++++++++- src/keycloak/keycloak_admin.py | 6 ++++++ src/keycloak/keycloak_openid.py | 7 ++++++- src/keycloak/openid_connection.py | 4 ++++ tests/conftest.py | 6 ++++++ tests/test_keycloak_admin.py | 16 ++++++++++++++++ tests/test_keycloak_openid.py | 10 ++++++++++ 9 files changed, 77 insertions(+), 4 deletions(-) diff --git a/docs/source/modules/admin.rst b/docs/source/modules/admin.rst index 3c88c96..7f20522 100644 --- a/docs/source/modules/admin.rst +++ b/docs/source/modules/admin.rst @@ -15,7 +15,8 @@ Configure admin client username='example-admin', password='secret', realm_name="master", - user_realm_name="only_if_other_realm_than_master") + user_realm_name="only_if_other_realm_than_master", + pool_maxsize=20) Configure admin client with connection @@ -34,6 +35,7 @@ Configure admin client with connection user_realm_name="only_if_other_realm_than_master", client_id="my_client", client_secret_key="client-secret", + pool_maxsize=25, verify=True) keycloak_admin = KeycloakAdmin(connection=keycloak_connection) diff --git a/docs/source/modules/openid_client.rst b/docs/source/modules/openid_client.rst index c3c0c90..a7ae009 100644 --- a/docs/source/modules/openid_client.rst +++ b/docs/source/modules/openid_client.rst @@ -16,7 +16,8 @@ Configure client OpenID keycloak_openid = KeycloakOpenID(server_url="http://localhost:8080/", client_id="example_client", realm_name="example_realm", - client_secret_key="secret") + client_secret_key="secret", + pool_maxsize=15) # Example: Set connection pool size Get .well_know diff --git a/src/keycloak/connection.py b/src/keycloak/connection.py index 400ee23..987c6f8 100644 --- a/src/keycloak/connection.py +++ b/src/keycloak/connection.py @@ -59,6 +59,8 @@ class ConnectionManager: :type cert: Union[str,Tuple[str,str]] :param max_retries: The total number of times to retry HTTP requests. :type max_retries: int + :param pool_maxsize: The maximum number of connections to save in the pool. + :type pool_maxsize: int """ def __init__( @@ -70,6 +72,7 @@ class ConnectionManager: proxies: dict | None = None, cert: str | tuple | None = None, max_retries: int = 1, + pool_maxsize: int | None = None, ) -> None: """ Init method. @@ -91,19 +94,25 @@ class ConnectionManager: :type cert: Union[str,Tuple[str,str]] :param max_retries: The total number of times to retry HTTP requests. :type max_retries: int + :param pool_maxsize: The maximum number of connections to save in the pool. + :type pool_maxsize: int """ self.base_url = base_url self.headers = headers self.timeout = timeout self.verify = verify self.cert = cert + self.pool_maxsize = pool_maxsize self._s = requests.Session() self._s.auth = lambda x: x # don't let requests add auth headers # retry once to reset connection with Keycloak after tomcat's ConnectionTimeout # see https://github.com/marcospereirampj/python-keycloak/issues/36 for protocol in ("https://", "http://"): - adapter = HTTPAdapter(max_retries=max_retries) + adapter_kwargs = {"max_retries": max_retries} + if pool_maxsize is not None: + adapter_kwargs["pool_maxsize"] = pool_maxsize + adapter = HTTPAdapter(**adapter_kwargs) # adds POST to retry whitelist allowed_methods = set(adapter.max_retries.allowed_methods) allowed_methods.add("POST") @@ -184,6 +193,20 @@ class ConnectionManager: def cert(self, value: str | tuple) -> None: self._cert = value + @property + def pool_maxsize(self) -> int | None: + """ + Return the maximum number of connections to save in the pool. + + :returns: Pool maxsize + :rtype: int or None + """ + return self._pool_maxsize + + @pool_maxsize.setter + def pool_maxsize(self, value: int | None) -> None: + self._pool_maxsize = value + @property def headers(self) -> dict: """ diff --git a/src/keycloak/keycloak_admin.py b/src/keycloak/keycloak_admin.py index 3af3319..f73ba31 100644 --- a/src/keycloak/keycloak_admin.py +++ b/src/keycloak/keycloak_admin.py @@ -87,6 +87,8 @@ class KeycloakAdmin: :type max_retries: int :param connection: A KeycloakOpenIDConnection as an alternative to individual params. :type connection: KeycloakOpenIDConnection + :param pool_maxsize: The maximum number of connections to save in the pool. + :type pool_maxsize: int """ PAGE_SIZE = 100 @@ -109,6 +111,7 @@ class KeycloakAdmin: cert: str | tuple | None = None, max_retries: int = 1, connection: KeycloakOpenIDConnection | None = None, + pool_maxsize: int | None = None, ) -> None: """ Init method. @@ -148,6 +151,8 @@ class KeycloakAdmin: :type max_retries: int :param connection: An OpenID Connection as an alternative to individual params. :type connection: KeycloakOpenIDConnection + :param pool_maxsize: The maximum number of connections to save in the pool. + :type pool_maxsize: int """ self.connection = connection or KeycloakOpenIDConnection( server_url=server_url, @@ -165,6 +170,7 @@ class KeycloakAdmin: timeout=timeout, cert=cert, max_retries=max_retries, + pool_maxsize=pool_maxsize, ) @property diff --git a/src/keycloak/keycloak_openid.py b/src/keycloak/keycloak_openid.py index 60f85e6..dddabbf 100644 --- a/src/keycloak/keycloak_openid.py +++ b/src/keycloak/keycloak_openid.py @@ -87,7 +87,8 @@ class KeycloakOpenID: Either a path to an SSL certificate file, or two-tuple of (certificate file, key file). :param max_retries: The total number of times to retry HTTP requests. - :type max_retries: int + :param pool_maxsize: The maximum number of connections to save in the pool. + :type pool_maxsize: int """ def __init__( @@ -102,6 +103,7 @@ class KeycloakOpenID: timeout: int = 60, cert: str | tuple | None = None, max_retries: int = 1, + pool_maxsize: int | None = None, ) -> None: """ Init method. @@ -129,6 +131,8 @@ class KeycloakOpenID: :type cert: Union[str,Tuple[str,str]] :param max_retries: The total number of times to retry HTTP requests. :type max_retries: int + :param pool_maxsize: The maximum number of connections to save in the pool. + :type pool_maxsize: int """ self.client_id = client_id self.client_secret_key = client_secret_key @@ -142,6 +146,7 @@ class KeycloakOpenID: proxies=proxies, cert=cert, max_retries=max_retries, + pool_maxsize=pool_maxsize, ) self.authorization = Authorization() diff --git a/src/keycloak/openid_connection.py b/src/keycloak/openid_connection.py index 3ec0519..c22d376 100644 --- a/src/keycloak/openid_connection.py +++ b/src/keycloak/openid_connection.py @@ -82,6 +82,7 @@ class KeycloakOpenIDConnection(ConnectionManager): timeout: int | None = 60, cert: str | tuple | None = None, max_retries: int = 1, + pool_maxsize: int | None = None, ) -> None: """ Init method. @@ -120,6 +121,8 @@ class KeycloakOpenIDConnection(ConnectionManager): :type cert: Union[str,Tuple[str,str]] :param max_retries: The total number of times to retry HTTP requests. :type max_retries: int + :param pool_maxsize: The maximum number of connections to save in the pool. + :type pool_maxsize: int """ # token is renewed when it hits 90% of its lifetime. This is to account for any possible # clock skew. @@ -154,6 +157,7 @@ class KeycloakOpenIDConnection(ConnectionManager): verify=self.verify, cert=cert, max_retries=max_retries, + pool_maxsize=pool_maxsize, ) @property diff --git a/tests/conftest.py b/tests/conftest.py index b5434c4..7261671 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -161,6 +161,7 @@ def admin(env: KeycloakTestEnv) -> KeycloakAdmin: server_url=f"http://{env.keycloak_host}:{env.keycloak_port}", username=env.keycloak_admin, password=env.keycloak_admin_password, + pool_maxsize=5, ) @@ -179,6 +180,7 @@ def admin_frozen(env: KeycloakTestEnv) -> KeycloakAdmin: server_url=f"http://{env.keycloak_host}:{env.keycloak_port}", username=env.keycloak_admin, password=env.keycloak_admin_password, + pool_maxsize=5, ) @@ -218,6 +220,7 @@ def oid( server_url=f"http://{env.keycloak_host}:{env.keycloak_port}", realm_name=realm, client_id=client, + pool_maxsize=5, ) # Cleanup admin.delete_client(client_id=client_id) @@ -279,6 +282,7 @@ def oid_with_credentials( realm_name=realm, client_id=client, client_secret_key=secret, + pool_maxsize=5, ), username, password, @@ -354,6 +358,7 @@ def oid_with_credentials_authz( realm_name=realm, client_id=client, client_secret_key=secret, + pool_maxsize=5, ), username, password, @@ -421,6 +426,7 @@ def oid_with_credentials_device( realm_name=realm, client_id=client, client_secret_key=secret, + pool_maxsize=5, ), username, password, diff --git a/tests/test_keycloak_admin.py b/tests/test_keycloak_admin.py index 03d7591..7b55296 100644 --- a/tests/test_keycloak_admin.py +++ b/tests/test_keycloak_admin.py @@ -57,6 +57,7 @@ def test_keycloak_admin_init(env: KeycloakTestEnv) -> None: server_url=f"http://{env.keycloak_host}:{env.keycloak_port}", username=env.keycloak_admin, password=env.keycloak_admin_password, + pool_maxsize=5, ) assert admin.connection.server_url == f"http://{env.keycloak_host}:{env.keycloak_port}", ( admin.connection.server_url @@ -72,6 +73,14 @@ def test_keycloak_admin_init(env: KeycloakTestEnv) -> None: assert admin.connection.token is None, admin.connection.token assert admin.connection.user_realm_name is None, admin.connection.user_realm_name assert admin.connection.custom_headers is None, admin.connection.custom_headers + assert admin.connection.pool_maxsize == 5, admin.connection.pool_maxsize + + admin_default = KeycloakAdmin( + server_url=f"http://{env.keycloak_host}:{env.keycloak_port}", + username=env.keycloak_admin, + password=env.keycloak_admin_password, + ) + assert admin_default.connection.pool_maxsize is None admin = KeycloakAdmin( server_url=f"http://{env.keycloak_host}:{env.keycloak_port}", @@ -79,6 +88,7 @@ def test_keycloak_admin_init(env: KeycloakTestEnv) -> None: password=env.keycloak_admin_password, realm_name=None, user_realm_name="master", + pool_maxsize=5, ) assert admin.connection.token is None admin = KeycloakAdmin( @@ -87,6 +97,7 @@ def test_keycloak_admin_init(env: KeycloakTestEnv) -> None: password=env.keycloak_admin_password, realm_name=None, user_realm_name=None, + pool_maxsize=5, ) assert admin.connection.token is None @@ -97,6 +108,7 @@ def test_keycloak_admin_init(env: KeycloakTestEnv) -> None: token=token, realm_name=None, user_realm_name=None, + pool_maxsize=5, ) assert admin.connection.token == token @@ -121,6 +133,7 @@ def test_keycloak_admin_init(env: KeycloakTestEnv) -> None: user_realm_name="authz", client_id="authz-client", client_secret_key=secret["value"], + pool_maxsize=5, ) admin_auth.connection.refresh_token() assert admin_auth.connection.token is not None @@ -133,6 +146,7 @@ def test_keycloak_admin_init(env: KeycloakTestEnv) -> None: password=None, client_secret_key=None, custom_headers={"custom": "header"}, + pool_maxsize=5, ).connection.token is None ) @@ -144,6 +158,7 @@ def test_keycloak_admin_init(env: KeycloakTestEnv) -> None: realm_name="master", client_id="admin-cli", verify=True, + pool_maxsize=5, ) keycloak_admin = KeycloakAdmin(connection=keycloak_connection) keycloak_admin.connection.get_token() @@ -3615,6 +3630,7 @@ async def test_a_realms(admin: KeycloakAdmin) -> None: realms = await admin.a_get_realms() assert len(realms) == 1, realms assert realms[0]["realm"] == "master" + assert admin.connection.pool_maxsize == 5, admin.connection.pool_maxsize # Create a test realm res = await admin.a_create_realm(payload={"realm": "test"}) diff --git a/tests/test_keycloak_openid.py b/tests/test_keycloak_openid.py index 1daf1af..ccb9ecc 100644 --- a/tests/test_keycloak_openid.py +++ b/tests/test_keycloak_openid.py @@ -35,6 +35,7 @@ def test_keycloak_openid_init(env: KeycloakTestEnv) -> None: server_url=f"http://{env.keycloak_host}:{env.keycloak_port}", realm_name="master", client_id="admin-cli", + pool_maxsize=5, ) assert oid.client_id == "admin-cli" @@ -42,6 +43,14 @@ def test_keycloak_openid_init(env: KeycloakTestEnv) -> None: assert oid.realm_name == "master" assert isinstance(oid.connection, ConnectionManager) assert isinstance(oid.authorization, Authorization) + assert oid.connection.pool_maxsize == 5 + + oid_default = KeycloakOpenID( + server_url=f"http://{env.keycloak_host}:{env.keycloak_port}", + realm_name="master", + client_id="admin-cli", + ) + assert oid_default.connection.pool_maxsize is None def test_well_known(oid: KeycloakOpenID) -> None: @@ -576,6 +585,7 @@ async def test_a_well_known(oid: KeycloakOpenID) -> None: res = await oid.a_well_known() assert res is not None assert res != {} + assert oid.connection.pool_maxsize == 5 for key in [ "acr_values_supported", "authorization_encryption_alg_values_supported",