Manual changes: lazy pagination and paginated tests verified

Author: Gayathri Srividya Rajavarapu

pyiceberg/catalog/rest/__init__.py

@@ -437,49 +437,55 @@ def _create_session(self) -> Session:
                 elif ssl_client_cert := ssl_client.get(CERT):
                     session.cert = ssl_client_cert
 
+        auth_type = None
+        source_env_var = None
+        auth_config = None
         if raw_auth := self.properties.get(AUTH):
             # When auth is configured via an environment variable (e.g. PYICEBERG_CATALOG__<NAME>__AUTH),
             # the value arrives as a JSON string rather than a dict. Decode it before processing.
+            source_env_var = f"PYICEBERG_CATALOG__{self.name.upper()}__AUTH"
             if isinstance(raw_auth, str):
                 try:
-                    auth_config: dict[str, Any] = json.loads(raw_auth)
+                    auth_config = json.loads(raw_auth)
                 except json.JSONDecodeError as e:
-                    raise ValueError(f"Failed to parse auth configuration as JSON: {raw_auth!r}") from e
+                    raise ValueError(f"Failed to parse auth configuration as JSON from {source_env_var}: {raw_auth!r}") from e
+                if not isinstance(auth_config, dict):
+                    raise ValueError(
+                        f"Auth configuration loaded from {source_env_var} must be a JSON object (dict), got {type(auth_config).__name__}: {raw_auth!r}"
+                    )
             else:
+                if not isinstance(raw_auth, dict):
+                    raise ValueError(
+                        f"Auth configuration for {source_env_var} must be a dict or JSON string, got {type(raw_auth).__name__}: {raw_auth!r}"
+                    )
                 auth_config = raw_auth
             auth_type = auth_config.get("type")
-            if auth_type is None:
-                raise ValueError("auth.type must be defined")
-            auth_type_config = auth_config.get(auth_type, {})
+            auth_type_config = auth_config.get(auth_type, {}) if auth_type else None
             auth_impl = auth_config.get("impl")
-
-            if auth_type == CUSTOM and not auth_impl:
-                raise ValueError("auth.impl must be specified when using custom auth.type")
-
-            if auth_type != CUSTOM and auth_impl:
-                raise ValueError("auth.impl can only be specified when using custom auth.type")
-
-            self._auth_manager = AuthManagerFactory.create(auth_impl or auth_type, auth_type_config)
-            session.auth = AuthManagerAdapter(self._auth_manager)
         elif auth_type := self.properties.get(f"{AUTH}.type"):
             # Support flattened env-var style configuration:
             #   PYICEBERG_CATALOG__<NAME>__AUTH__TYPE=oauth2
             #   PYICEBERG_CATALOG__<NAME>__AUTH__OAUTH2__CLIENT_ID=id
             # The env-var parser maps these to flat properties like "auth.type" and "auth.oauth2.client-id".
             # Key names are converted from kebab-case to snake_case to match AuthManager constructor parameters.
             auth_impl = self.properties.get(f"{AUTH}.impl")
-
-            if auth_type == CUSTOM and not auth_impl:
-                raise ValueError("auth.impl must be specified when using custom auth.type")
-
-            if auth_type != CUSTOM and auth_impl:
-                raise ValueError("auth.impl can only be specified when using custom auth.type")
-
             type_prefix = f"{AUTH}.{auth_type}."
             auth_type_config = {
                 k[len(type_prefix) :].replace("-", "_"): v for k, v in self.properties.items() if k.startswith(type_prefix)
             }
-
+        if auth_type:
+            if auth_type is None:
+                raise ValueError(
+                    f"auth.type must be defined in auth configuration{f' from {source_env_var}' if source_env_var else ''}"
+                )
+            if auth_type == CUSTOM and not auth_impl:
+                raise ValueError(
+                    f"auth.impl must be specified when using custom auth.type{f' (from {source_env_var})' if source_env_var else ''}"
+                )
+            if auth_type != CUSTOM and auth_impl:
+                raise ValueError(
+                    f"auth.impl can only be specified when using custom auth.type{f' (from {source_env_var})' if source_env_var else ''}"
+                )
             self._auth_manager = AuthManagerFactory.create(auth_impl or auth_type, auth_type_config)
             session.auth = AuthManagerAdapter(self._auth_manager)
         else:

tests/catalog/test_rest.py

@@ -198,10 +198,10 @@ def test_token_200(rest_mock: Mocker) -> None:
         status_code=200,
         request_headers=OAUTH_TEST_HEADERS,
     )
-    assert (
-        RestCatalog("rest", uri=TEST_URI, credential=TEST_CREDENTIALS)._session.headers["Authorization"]  # pylint: disable=W0212
-        == f"Bearer {TEST_TOKEN}"
-    )
+    catalog = RestCatalog("rest", uri=TEST_URI, credential=TEST_CREDENTIALS)
+    req = Request("GET", "http://example.com")
+    prepped = catalog._session.prepare_request(req)
+    assert prepped.headers["Authorization"] == f"Bearer {TEST_TOKEN}"
 
 
 @pytest.mark.filterwarnings(
@@ -218,10 +218,10 @@ def test_token_200_without_optional_fields(rest_mock: Mocker) -> None:
         status_code=200,
         request_headers=OAUTH_TEST_HEADERS,
     )
-    assert (
-        RestCatalog("rest", uri=TEST_URI, credential=TEST_CREDENTIALS)._session.headers["Authorization"]  # pylint: disable=W0212
-        == f"Bearer {TEST_TOKEN}"
-    )
+    catalog = RestCatalog("rest", uri=TEST_URI, credential=TEST_CREDENTIALS)
+    req = Request("GET", "http://example.com")
+    prepped = catalog._session.prepare_request(req)
+    assert prepped.headers["Authorization"] == f"Bearer {TEST_TOKEN}"
 
 
 @pytest.mark.filterwarnings(
@@ -240,12 +240,10 @@ def test_token_with_optional_oauth_params(rest_mock: Mocker) -> None:
         status_code=200,
         request_headers=OAUTH_TEST_HEADERS,
     )
-    assert (
-        RestCatalog(
-            "rest", uri=TEST_URI, credential=TEST_CREDENTIALS, audience=TEST_AUDIENCE, resource=TEST_RESOURCE
-        )._session.headers["Authorization"]
-        == f"Bearer {TEST_TOKEN}"
-    )
+    catalog = RestCatalog("rest", uri=TEST_URI, credential=TEST_CREDENTIALS, audience=TEST_AUDIENCE, resource=TEST_RESOURCE)
+    req = Request("GET", "http://example.com")
+    prepped = catalog._session.prepare_request(req)
+    assert prepped.headers["Authorization"] == f"Bearer {TEST_TOKEN}"
     assert TEST_AUDIENCE in mock_request.last_request.text
     assert TEST_RESOURCE in mock_request.last_request.text
 
@@ -266,10 +264,10 @@ def test_token_with_optional_oauth_params_as_empty(rest_mock: Mocker) -> None:
         status_code=200,
         request_headers=OAUTH_TEST_HEADERS,
     )
-    assert (
-        RestCatalog("rest", uri=TEST_URI, credential=TEST_CREDENTIALS, audience="", resource="")._session.headers["Authorization"]
-        == f"Bearer {TEST_TOKEN}"
-    )
+    catalog = RestCatalog("rest", uri=TEST_URI, credential=TEST_CREDENTIALS, audience="", resource="")
+    req = Request("GET", "http://example.com")
+    prepped = catalog._session.prepare_request(req)
+    assert prepped.headers["Authorization"] == f"Bearer {TEST_TOKEN}"
     assert TEST_AUDIENCE not in mock_request.last_request.text
     assert TEST_RESOURCE not in mock_request.last_request.text
 
@@ -290,9 +288,10 @@ def test_token_with_default_scope(rest_mock: Mocker) -> None:
         status_code=200,
         request_headers=OAUTH_TEST_HEADERS,
     )
-    assert (
-        RestCatalog("rest", uri=TEST_URI, credential=TEST_CREDENTIALS)._session.headers["Authorization"] == f"Bearer {TEST_TOKEN}"
-    )
+    catalog = RestCatalog("rest", uri=TEST_URI, credential=TEST_CREDENTIALS)
+    req = Request("GET", "http://example.com")
+    prepped = catalog._session.prepare_request(req)
+    assert prepped.headers["Authorization"] == f"Bearer {TEST_TOKEN}"
     assert "catalog" in mock_request.last_request.text
 
 
@@ -312,10 +311,10 @@ def test_token_with_custom_scope(rest_mock: Mocker) -> None:
         status_code=200,
         request_headers=OAUTH_TEST_HEADERS,
     )
-    assert (
-        RestCatalog("rest", uri=TEST_URI, credential=TEST_CREDENTIALS, scope=TEST_SCOPE)._session.headers["Authorization"]
-        == f"Bearer {TEST_TOKEN}"
-    )
+    catalog = RestCatalog("rest", uri=TEST_URI, credential=TEST_CREDENTIALS, scope=TEST_SCOPE)
+    req = Request("GET", "http://example.com")
+    prepped = catalog._session.prepare_request(req)
+    assert prepped.headers["Authorization"] == f"Bearer {TEST_TOKEN}"
     assert TEST_SCOPE in mock_request.last_request.text
 
 
@@ -335,14 +334,10 @@ def test_token_200_w_oauth2_server_uri(rest_mock: Mocker) -> None:
         status_code=200,
         request_headers=OAUTH_TEST_HEADERS,
     )
-    # pylint: disable=W0212
-    assert (
-        RestCatalog("rest", uri=TEST_URI, credential=TEST_CREDENTIALS, **{OAUTH2_SERVER_URI: OAUTH2_SERVER_URI})._session.headers[
-            "Authorization"
-        ]
-        == f"Bearer {TEST_TOKEN}"
-    )
-    # pylint: enable=W0212
+    catalog = RestCatalog("rest", uri=TEST_URI, credential=TEST_CREDENTIALS, **{OAUTH2_SERVER_URI: OAUTH2_SERVER_URI})
+    req = Request("GET", "http://example.com")
+    prepped = catalog._session.prepare_request(req)
+    assert prepped.headers["Authorization"] == f"Bearer {TEST_TOKEN}"
 
 
 @pytest.mark.filterwarnings(
@@ -3173,66 +3168,87 @@ def test_load_table_without_storage_credentials(
 # variables unless auth JSON strings are decoded.
 
 
-def test_rest_catalog_with_basic_auth_as_json_string(rest_mock: Mocker) -> None:
-    """When auth arrives as a JSON string (e.g. from an environment variable), it should be decoded correctly."""
-    import json
-
-    rest_mock.get(
-        f"{TEST_URI}v1/config",
-        json={"defaults": {}, "overrides": {}},
-        status_code=200,
-    )
-    auth_dict = {
-        "type": "basic",
-        "basic": {
-            "username": "one",
-            "password": "two",
-        },
-    }
-    catalog_properties = {
-        "uri": TEST_URI,
-        "auth": json.dumps(auth_dict),
-    }
-    catalog = RestCatalog("rest", **catalog_properties)
-    assert catalog.uri == TEST_URI
-
-    encoded_user_pass = base64.b64encode(b"one:two").decode()
-    expected_auth_header = f"Basic {encoded_user_pass}"
-    assert rest_mock.last_request.headers["Authorization"] == expected_auth_header
+import json
 
+import pytest
 
-def test_rest_catalog_with_oauth2_auth_as_json_string(requests_mock: Mocker) -> None:
-    """OAuth2 auth configured as a JSON string (e.g. from an environment variable) should work correctly."""
-    import json
 
-    requests_mock.post(
-        f"{TEST_URI}oauth2/token",
-        json={
-            "access_token": "MTQ0NjJkZmQ5OTM2NDE1ZTZjNGZmZjI3",
-            "token_type": "Bearer",
-            "expires_in": 3600,
-        },
-        status_code=200,
-    )
-    requests_mock.get(
-        f"{TEST_URI}v1/config",
-        json={"defaults": {}, "overrides": {}},
-        status_code=200,
-    )
-    auth_dict = {
-        "type": "oauth2",
-        "oauth2": {
-            "client_id": "some_client_id",
-            "client_secret": "some_client_secret",
-            "token_url": f"{TEST_URI}oauth2/token",
-        },
-    }
-    catalog_properties = {
-        "uri": TEST_URI,
-        "auth": json.dumps(auth_dict),
-    }
-    catalog = RestCatalog("rest", **catalog_properties)
-    assert catalog.uri == TEST_URI
+@pytest.mark.parametrize(
+    "auth_dict, expected_header, mocker_type",
+    [
+        # Basic auth
+        (
+            {"type": "basic", "basic": {"username": "one", "password": "two"}},
+            lambda: f"Basic {base64.b64encode(b'one:two').decode()}",
+            "rest_mock",
+        ),
+        # OAuth2 auth
+        (
+            {
+                "type": "oauth2",
+                "oauth2": {
+                    "client_id": "some_client_id",
+                    "client_secret": "some_client_secret",
+                    "token_url": f"{TEST_URI}oauth2/token",
+                },
+            },
+            None,  # OAuth2 does not set Authorization header immediately
+            "requests_mock",
+        ),
+        # OAuth2 with int fields
+        (
+            {
+                "type": "oauth2",
+                "oauth2": {
+                    "client_id": "id",
+                    "client_secret": "secret",
+                    "token_url": f"{TEST_URI}oauth2/token",
+                    "refresh_margin": 10,
+                    "expires_in": 3600,
+                },
+            },
+            None,
+            "requests_mock",
+        ),
+    ],
+)
+def test_rest_catalog_with_auth_json_string(requests_mock, rest_mock, auth_dict, expected_header, mocker_type):
+    """Test various auth configs as JSON string (from env var) are decoded and handled correctly."""
+    if mocker_type == "rest_mock":
+        rest_mock.get(
+            f"{TEST_URI}v1/config",
+            json={"defaults": {}, "overrides": {}},
+            status_code=200,
+        )
+        catalog_properties = {
+            "uri": TEST_URI,
+            "auth": json.dumps(auth_dict),
+        }
+        catalog = RestCatalog("rest", **catalog_properties)
+        assert catalog.uri == TEST_URI
+        if expected_header:
+            assert rest_mock.last_request.headers["Authorization"] == expected_header()
+    else:
+        requests_mock.post(
+            f"{TEST_URI}oauth2/token",
+            json={
+                "access_token": "MTQ0NjJkZmQ5OTM2NDE1ZTZjNGZmZjI3",
+                "token_type": "Bearer",
+                "expires_in": 3600,
+            },
+            status_code=200,
+        )
+        requests_mock.get(
+            f"{TEST_URI}v1/config",
+            json={"defaults": {}, "overrides": {}},
+            status_code=200,
+        )
+        catalog_properties = {
+            "uri": TEST_URI,
+            "auth": json.dumps(auth_dict),
+        }
+        catalog = RestCatalog("rest", **catalog_properties)
+        assert catalog.uri == TEST_URI
 
 
 def test_rest_catalog_with_invalid_json_auth_string() -> None: