Add OIDC credential providers

CHANGELOG.md

@@ -1,3 +1,4 @@
+* Add static-token, Client Credentials, and Device Authorization OAuth 2.0 providers under `ydb.oidc`, with non-blocking Client Credentials and Device Authorization variants under `ydb.aio.oidc`
 * Deprecated the table client scan query methods — `TableClient.scan_query`, `TableClient.async_scan_query` and the async `ydb.aio.TableClient.scan_query`: they now emit a `DeprecationWarning` and keep working as before, use QueryService (`ydb.QuerySessionPool` / `ydb.aio.QuerySessionPool`) instead
 * Mark the package as typed so type checkers use the SDK's inline annotations
 

docs/driver.rst

@@ -182,6 +182,62 @@ Pass a static IAM token or API key directly
 
     credentials = ydb.AccessTokenCredentials("your-token")
 
+OIDC and OAuth 2.0 Credentials
+^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+
+``ydb.oidc`` provides three OAuth 2.0 credential modes for an external identity
+provider advertised through OIDC discovery
+(`full example <https://github.com/ydb-platform/ydb-python-sdk/tree/main/examples/oidc-credentials>`__).
+
+Use an access token obtained outside the SDK:
+
+.. code-block:: python
+
+    import ydb.oidc
+
+    credentials = ydb.oidc.OAuth2TokenCredentials("your-access-token")
+
+Use the Client Credentials Grant for a service acting on its own behalf:
+
+.. code-block:: python
+
+    credentials = ydb.oidc.OAuth2ClientCredentials(
+        issuer="https://identity.example.com/realms/example",
+        client_id="service-client",
+        client_secret="service-client-secret",
+        audience="ydb",
+        ca_file="/path/to/idp-ca.pem",
+    )
+
+Use the Device Authorization Grant for a CLI or another input-constrained client. The
+callback is invoked before polling begins and must show the verification URI and user
+code to the user:
+
+.. code-block:: python
+
+    def show_device_authorization(info):
+        print(info.verification_uri_complete or info.verification_uri)
+        print(info.user_code)
+
+    credentials = ydb.oidc.OAuth2DeviceCredentials(
+        issuer="https://identity.example.com/realms/example",
+        client_id="public-device-client",
+        device_authorization_callback=show_device_authorization,
+        scope=["offline_access"],
+        ca_file="/path/to/idp-ca.pem",
+    )
+
+Client Credentials obtains a new access token when needed. Device Authorization uses a
+returned refresh token for subsequent refreshes and starts a new user interaction if
+the refresh token is no longer valid. Access tokens are sent to YDB using the
+``Bearer`` authentication scheme. The providers require HTTPS issuer and endpoint URLs,
+require the discovered issuer to match the configured value exactly, and add the
+``openid`` scope when it is absent.
+
+Non-blocking counterparts are available as ``ydb.aio.oidc.OAuth2ClientCredentials``
+and ``ydb.aio.oidc.OAuth2DeviceCredentials``. The asynchronous Device Authorization
+callback may be either a regular callable or an async callable.
+
 StaticCredentials (username/password)
 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
 

examples/oidc-credentials/main.py

@@ -0,0 +1,70 @@
+import argparse
+
+import ydb
+import ydb.oidc
+
+
+def parse_args():
+    parser = argparse.ArgumentParser(description="OIDC and OAuth 2.0 credentials example")
+    parser.add_argument("--endpoint", required=True, help="YDB endpoint")
+    parser.add_argument("--database", required=True, help="YDB database")
+    parser.add_argument("--mode", required=True, choices=("token", "client", "device"))
+    parser.add_argument("--access-token", help="Existing OAuth 2.0 access token")
+    parser.add_argument("--issuer", help="OIDC issuer URL")
+    parser.add_argument("--client-id", help="OAuth 2.0 client ID")
+    parser.add_argument("--client-secret", help="OAuth 2.0 client secret")
+    parser.add_argument("--scope", action="append", help="OAuth 2.0 scope; may be repeated")
+    parser.add_argument("--audience", help="OAuth 2.0 token audience")
+    parser.add_argument("--ca-file", help="CA certificate for the identity provider")
+    args = parser.parse_args()
+
+    if args.mode == "token" and not args.access_token:
+        parser.error("--access-token is required in token mode")
+    if args.mode in ("client", "device") and (not args.issuer or not args.client_id):
+        parser.error("--issuer and --client-id are required in client and device modes")
+    if args.mode == "client" and not args.client_secret:
+        parser.error("--client-secret is required in client mode")
+
+    return args
+
+
+def device_authorization_callback(info):
+    print("Open {}".format(info.verification_uri_complete or info.verification_uri))
+    print("Enter code {}".format(info.user_code))
+
+
+def create_credentials(args):
+    if args.mode == "token":
+        return ydb.oidc.OAuth2TokenCredentials(args.access_token)
+
+    common = {
+        "issuer": args.issuer,
+        "client_id": args.client_id,
+        "audience": args.audience,
+        "ca_file": args.ca_file,
+    }
+    if args.scope is not None:
+        common["scope"] = args.scope
+    if args.mode == "client":
+        return ydb.oidc.OAuth2ClientCredentials(client_secret=args.client_secret, **common)
+    return ydb.oidc.OAuth2DeviceCredentials(
+        device_authorization_callback=device_authorization_callback,
+        **common,
+    )
+
+
+def main():
+    args = parse_args()
+    with ydb.Driver(
+        endpoint=args.endpoint,
+        database=args.database,
+        credentials=create_credentials(args),
+    ) as driver:
+        driver.wait(timeout=10, fail_fast=True)
+        with ydb.QuerySessionPool(driver) as pool:
+            result_sets = pool.execute_with_retries("SELECT 42 AS answer")
+    print(result_sets[0].rows[0].answer)
+
+
+if __name__ == "__main__":
+    main()

tests/aio/test_credentials.py

@@ -6,14 +6,17 @@
 import os
 import json
 import asyncio
+import aiohttp
 from unittest.mock import patch, AsyncMock, MagicMock
 
 import tests.auth.test_credentials
 import tests.oauth2_token_exchange
 import tests.oauth2_token_exchange.test_token_exchange
 import ydb.aio.iam
+import ydb.aio.oidc
 import ydb.aio.oauth2_token_exchange
 import ydb.oauth2_token_exchange.token_source
+from ydb import issues
 
 
 class ServiceAccountCredentialsForTest(ydb.aio.iam.ServiceAccountCredentials):
@@ -263,3 +266,235 @@ def mock_submit(callback):
         token3 = await credentials.token()
         assert token3 == "token_v2"
         assert call_count == 2
+
+
+@pytest.mark.asyncio
+async def test_oauth2_client_credentials():
+    issuer = "https://issuer.example"
+    credentials = ydb.aio.oidc.OAuth2ClientCredentials(issuer, "client-id", "client-secret")
+    credentials._request_json = AsyncMock(
+        side_effect=[
+            (200, {"issuer": issuer, "token_endpoint": issuer + "/token"}),
+            (200, {"access_token": "access-token", "token_type": "Bearer", "expires_in": 300}),
+        ]
+    )
+
+    assert await credentials.get_auth_token() == "Bearer access-token"
+    assert credentials._request_json.await_count == 2
+    assert credentials._request_json.await_args_list[1].args[1]["scope"] == "openid"
+
+
+@pytest.mark.asyncio
+async def test_oauth2_async_token_lifetime_starts_after_request_and_preserves_short_lived_token():
+    credentials = ydb.aio.oidc.OAuth2ClientCredentials(
+        "https://issuer.example",
+        "client-id",
+        "client-secret",
+    )
+    now = [1000]
+
+    async def make_token_request():
+        now[0] = 1120
+        return {"access_token": "Bearer access-token", "expires_in": 10}
+
+    credentials._make_token_request = make_token_request
+
+    with patch("ydb.aio.credentials.time.time", side_effect=lambda: now[0]):
+        assert await credentials.get_auth_token() == "Bearer access-token"
+
+    assert credentials._expires_in == 1129
+
+
+@pytest.mark.asyncio
+async def test_oauth2_device_credentials():
+    issuer = "https://issuer.example"
+    callback_values = []
+
+    async def callback(info):
+        callback_values.append(info)
+
+    credentials = ydb.aio.oidc.OAuth2DeviceCredentials(issuer, "device-client", callback)
+    credentials._request_json = AsyncMock(
+        side_effect=[
+            (
+                200,
+                {
+                    "issuer": issuer,
+                    "token_endpoint": issuer + "/token",
+                    "device_authorization_endpoint": issuer + "/device",
+                },
+            ),
+            (
+                200,
+                {
+                    "device_code": "device-code",
+                    "user_code": "user-code",
+                    "verification_uri": issuer + "/verify",
+                    "expires_in": 600,
+                    "interval": 1,
+                },
+            ),
+            (400, {"error": "authorization_pending"}),
+            (400, {"error": "slow_down"}),
+            (200, {"access_token": "access-token", "token_type": "Bearer", "expires_in": 300}),
+        ]
+    )
+
+    with patch("ydb.aio.oidc.asyncio.sleep", new=AsyncMock()) as sleep:
+        assert await credentials.get_auth_token() == "Bearer access-token"
+
+    assert callback_values[0].user_code == "user-code"
+    assert sleep.await_count == 3
+    assert credentials._request_json.await_args_list[2].kwargs["request_timeout"] == 10
+
+
+@pytest.mark.asyncio
+async def test_oauth2_async_http_requests_and_discovery_cache():
+    credentials = ydb.aio.oidc.OAuth2ClientCredentials(
+        "https://issuer.example/",
+        "client-id",
+        "client-secret",
+    )
+    response = MagicMock(status=200)
+    response.read = AsyncMock(
+        return_value=b'{"issuer":"https://issuer.example/","token_endpoint":"https://issuer.example/token"}'
+    )
+    request_context = MagicMock()
+    request_context.__aenter__ = AsyncMock(return_value=response)
+    request_context.__aexit__ = AsyncMock(return_value=None)
+    session = MagicMock()
+    session.request.return_value = request_context
+    session_context = MagicMock()
+    session_context.__aenter__ = AsyncMock(return_value=session)
+    session_context.__aexit__ = AsyncMock(return_value=None)
+
+    with patch("ydb.aio.oidc.aiohttp.ClientSession", return_value=session_context) as client_session:
+        first = await credentials._discovery()
+        second = await credentials._discovery()
+        await credentials._request_json("https://issuer.example/token", request_timeout=0.5)
+        await credentials._request_json("HTTPS://issuer.example/token")
+
+    assert first is second
+    assert client_session.call_count == 3
+    assert client_session.call_args_list[1].kwargs["timeout"].total == 0.5
+    assert session.request.call_args_list[0].args[1] == "https://issuer.example/.well-known/openid-configuration"
+    assert session.request.call_args.kwargs["ssl"] is not None
+    assert session.request.call_args.kwargs["allow_redirects"] is False
+
+    with patch(
+        "ydb.aio.oidc.aiohttp.ClientSession",
+        side_effect=aiohttp.ClientError("unavailable"),
+    ):
+        with pytest.raises(issues.Unavailable):
+            await credentials._request_json("http://issuer.example/token", {"key": "value"})
+
+
+@pytest.mark.asyncio
+async def test_oauth2_async_device_refresh_and_error_paths():
+    with pytest.raises(ValueError):
+        ydb.aio.oidc.OAuth2ClientCredentials("https://issuer.example", "client-id", "")
+    with pytest.raises(ValueError):
+        ydb.aio.oidc.OAuth2DeviceCredentials("https://issuer.example", "client-id", None)
+    with pytest.raises(ValueError):
+        ydb.aio.oidc.OAuth2DeviceCredentials(
+            "https://issuer.example",
+            "client-id",
+            lambda info: None,
+            device_flow_timeout=0,
+        )
+
+    credentials = ydb.aio.oidc.OAuth2DeviceCredentials(
+        "https://issuer.example",
+        "client-id",
+        lambda info: None,
+        client_secret="client-secret",
+    )
+    assert credentials._client_headers()["Authorization"].startswith("Basic ")
+    with pytest.raises(issues.Error, match="refresh_token"):
+        credentials._save_token_response(
+            {
+                "access_token": "token",
+                "token_type": "Bearer",
+                "expires_in": 300,
+                "refresh_token": "",
+            }
+        )
+
+    credentials._refresh_token_value = "refresh-token"
+    credentials._request_json = AsyncMock(
+        return_value=(
+            200,
+            {
+                "access_token": "refreshed-token",
+                "token_type": "Bearer",
+                "expires_in": 300,
+                "refresh_token": "new-refresh-token",
+            },
+        )
+    )
+    credentials._discovery_document = {"token_endpoint": "https://issuer.example/token"}
+    assert await credentials._make_token_request() == {
+        "access_token": "Bearer refreshed-token",
+        "expires_in": 300,
+    }
+    assert credentials._refresh_token_value == "new-refresh-token"
+
+    credentials._request_json = AsyncMock(return_value=(400, {"error": "invalid_grant"}))
+    assert await credentials._try_refresh("https://issuer.example/token") is None
+    assert credentials._refresh_token_value is None
+
+    credentials._discovery_document = {"token_endpoint": "https://issuer.example/token"}
+    with pytest.raises(issues.Error, match="device_authorization_endpoint"):
+        await credentials._make_token_request()
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize(
+    "token_response, expected_message, monotonic_values",
+    [
+        ({"error": "expired_token"}, "expired", [0, 0, 0]),
+        ({"error": "access_denied"}, "access_denied", [0, 0, 0]),
+        (None, "timed out", [0, 2]),
+        (None, "timed out", [0, 0, 2]),
+    ],
+)
+async def test_oauth2_async_device_expiration(token_response, expected_message, monotonic_values):
+    callback_values = []
+    loop = MagicMock()
+
+    def callback(info):
+        callback_values.append(info)
+        assert loop.time.call_count == 1
+
+    credentials = ydb.aio.oidc.OAuth2DeviceCredentials(
+        "https://issuer.example",
+        "client-id",
+        callback,
+        device_flow_timeout=1,
+    )
+    credentials._discovery_document = {
+        "token_endpoint": "https://issuer.example/token",
+        "device_authorization_endpoint": "https://issuer.example/device",
+    }
+    device_response = {
+        "device_code": "device-code",
+        "user_code": "user-code",
+        "verification_uri": "https://issuer.example/verify",
+        "expires_in": 600,
+        "interval": 1,
+    }
+    responses = [(200, device_response)]
+    loop.time.side_effect = monotonic_values
+    if token_response is not None:
+        responses.append((400, token_response))
+    credentials._request_json = AsyncMock(side_effect=responses)
+
+    with patch("ydb.aio.oidc.asyncio.get_running_loop", return_value=loop), patch(
+        "ydb.aio.oidc.asyncio.sleep", new=AsyncMock()
+    ):
+        with pytest.raises(issues.Unauthenticated, match=expected_message):
+            await credentials._make_token_request()
+
+    assert callback_values[0].user_code == "user-code"
+    if token_response is not None:
+        assert credentials._request_json.await_args_list[1].kwargs["request_timeout"] == 1