connector-py 4.180.0__tar.gz → 4.181.0__tar.gz

{connector_py-4.180.0 → connector_py-4.181.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: connector-py
-Version: 4.180.0
+Version: 4.181.0
 Summary: An Abstract Tool to Perform Actions on Integrations.
 Author-email: teamlumos <security@lumos.com>
 License-Expression: Apache-2.0

connector_py-4.181.0/connector/__about__.py

@@ -0,0 +1 @@
+__version__ = "4.181.0"

{connector_py-4.180.0 → connector_py-4.181.0}/connector/client.py

@@ -1,5 +1,9 @@
 from .httpx_rewrite import AsyncClient, GqlHTTPXAsyncTransport, HTTPXAsyncTransport
-from .oai.base_clients import BaseGraphQLSession, BaseIntegrationClient
+from .oai.base_clients import (
+    BaseGraphQLSession,
+    BaseIntegrationClient,
+    RateLimitedHTTPXAsyncTransport,
+)
 from .oai.capability import (
     get_basic_auth,
     get_jwt_auth,
@@ -20,6 +24,7 @@ from .utils.sync_to_async import sync_to_async
 __all__ = [
     "GqlHTTPXAsyncTransport",
     "HTTPXAsyncTransport",
+    "RateLimitedHTTPXAsyncTransport",
     "BaseGraphQLSession",
     "BaseIntegrationClient",
     "get_basic_auth",

{connector_py-4.180.0 → connector_py-4.181.0}/connector/oai/base_clients.py

@@ -10,11 +10,11 @@ from connector_sdk_types.generated import ErrorCode
 from gql import Client
 from gql.client import AsyncClientSession
 from gql.dsl import DSLSchema
-from graphql import GraphQLSchema, build_client_schema, build_schema
+from graphql import DocumentNode, GraphQLSchema, build_client_schema, build_schema
 from httpx import Response
 from typing_extensions import Self
 
-from connector.httpx_rewrite import AsyncClient
+from connector.httpx_rewrite import AsyncClient, HTTPXAsyncTransport
 from connector.oai.capability import Request
 from connector.oai.errors import ConnectorError
 from connector.utils.rate_limiting import RateLimitConfig, RateLimiter
@@ -178,6 +178,75 @@ class RateLimitedClient(AsyncClient):
             await self.base_client.__aexit__(exc_type, exc_val, exc_tb)
 
 
+class RateLimitedHTTPXAsyncTransport(HTTPXAsyncTransport):
+    """A wrapper around HTTPXAsyncTransport that applies rate limiting to GraphQL requests."""
+
+    def __init__(self, base_transport: HTTPXAsyncTransport, rate_limit_config: RateLimitConfig):
+        # Copy all attributes from base transport, but exclude 'execute' to avoid shadowing our method
+        base_dict = {k: v for k, v in base_transport.__dict__.items() if k != "execute"}
+        self.__dict__.update(base_dict)
+        self.base_transport = base_transport
+        self.rate_limiter = RateLimiter[Callable[[], Any], Any](rate_limit_config)
+
+    async def connect(self):
+        """Connect the underlying transport and replace its client with a rate-limited one."""
+        await self.base_transport.connect()
+
+        # Replace the base transport's client with a rate-limited version
+        if hasattr(self.base_transport, "client") and self.base_transport.client:
+            # The transport's client should be our AsyncClient type, but we need to handle the type
+            if isinstance(self.base_transport.client, AsyncClient):
+                self.base_transport.client = RateLimitedClient(
+                    self.base_transport.client, self.rate_limiter.config
+                )
+
+        # Copy the client reference
+        self.client = self.base_transport.client
+
+    async def execute(
+        self,
+        document: DocumentNode,
+        variable_values: dict[str, Any] | None = None,
+        operation_name: str | None = None,
+        extra_args: dict[str, Any] | None = None,
+        upload_files: bool = False,
+    ):
+        """Execute a GraphQL request with rate limiting."""
+
+        async def request_func():
+            result = await self.base_transport.execute(
+                document=document,
+                variable_values=variable_values,
+                operation_name=operation_name,
+                extra_args=extra_args,
+                upload_files=upload_files,
+            )
+            return result
+
+        # Use the rate limiter to execute the request
+        responses = await self.rate_limiter.execute_requests([request_func], lambda x: x())
+        if responses:
+            return responses[0]
+
+        raise ConnectorError(
+            message="No response from GraphQL API",
+            error_code=ErrorCode.API_ERROR,
+        )
+
+    def get_state(self) -> tuple[RateLimitConfig, float]:
+        """Get the current rate limit state."""
+        return self.rate_limiter.config, self.rate_limiter.current_delay
+
+    async def close(self):
+        """Close the underlying transport."""
+        if hasattr(self.base_transport, "close"):
+            await self.base_transport.close()
+
+    def __getattr__(self, name):
+        """Delegate attribute access to the underlying base_transport."""
+        return getattr(self.base_transport, name)
+
+
 class BaseIntegrationClient:
     _http_client: AsyncClient | RateLimitedClient
     _rate_limit_config: RateLimitConfig | None = None
@@ -248,8 +317,11 @@ class BaseIntegrationClient:
 
 
 class BaseGraphQLSession(AsyncClientSession):
-    def __init__(self, args: Request):
-        super().__init__(client=self.build_client(args))
+    _rate_limit_config: RateLimitConfig | None = None
+
+    def __init__(self, args: Request, rate_limit_config: RateLimitConfig | None = None):
+        client = self.build_client(args, rate_limit_config)
+        super().__init__(client=client)
 
     async def __aenter__(self) -> Self:
         await self.client.__aenter__()
@@ -267,8 +339,31 @@ class BaseGraphQLSession(AsyncClientSession):
         pass
 
     @classmethod
-    def build_client(cls, args: Request) -> Client:
-        return Client(**cls.prepare_client_args(args))
+    def build_client(
+        cls, args: Request, rate_limit_config: RateLimitConfig | None = None
+    ) -> Client:
+        client_args = cls.prepare_client_args(args)
+
+        # Apply rate limiting if configured
+        rate_limiting = rate_limit_config or cls._rate_limit_config
+        if rate_limiting is not None and "transport" in client_args:
+            transport = client_args["transport"]
+            if isinstance(transport, HTTPXAsyncTransport):
+                client_args["transport"] = RateLimitedHTTPXAsyncTransport(transport, rate_limiting)
+
+        return Client(**client_args)
+
+    def get_current_rate_limits(self) -> tuple[RateLimitConfig | None, float]:
+        """
+        Get the current rate limit state.
+
+        Returns a tuple of the rate limit config and the current delay. (or None if the client is not rate limited)
+        """
+        if hasattr(self.client, "transport") and isinstance(
+            self.client.transport, RateLimitedHTTPXAsyncTransport
+        ):
+            return self.client.transport.get_state()
+        return None, 0
 
     @classmethod
     def load_schema(cls, schema_file_path: str | Path) -> GraphQLSchema:

{connector_py-4.180.0 → connector_py-4.181.0}/connector_py.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: connector-py
-Version: 4.180.0
+Version: 4.181.0
 Summary: An Abstract Tool to Perform Actions on Integrations.
 Author-email: teamlumos <security@lumos.com>
 License-Expression: Apache-2.0

{connector_py-4.180.0 → connector_py-4.181.0}/tests/oai/test_base_clients.py

@@ -11,15 +11,18 @@ from connector.generated import (
     TokenCredential,
 )
 from connector.oai.base_clients import (
+    BaseGraphQLSession,
     BaseIntegrationClient,
     BatchRequest,
     BatchRequests,
     RateLimitedClient,
+    RateLimitedHTTPXAsyncTransport,
 )
 from connector.oai.capability import Request, get_token_auth
 from connector.oai.errors import ConnectorError
 from connector.utils.httpx_auth import BearerAuth
-from connector.utils.rate_limiting import RateLimitConfig, RateLimitStrategy
+from connector.utils.rate_limiting import RateLimitConfig, RateLimiter, RateLimitStrategy
+from connector_sdk_types.generated import ErrorCode
 
 
 @pytest.fixture(autouse=True)
@@ -579,3 +582,365 @@ class TestIntegrationScenarios:
             assert len(responses) == 2
             assert responses[0].status_code == 200
             assert responses[1].status_code == 404
+
+
+class TestRateLimitedHTTPXAsyncTransport:
+    """Test cases for RateLimitedHTTPXAsyncTransport."""
+
+    @pytest.fixture
+    def mock_base_transport(self):
+        """Create a mock HTTPXAsyncTransport for testing."""
+        from types import SimpleNamespace
+
+        transport = SimpleNamespace()
+        transport.some_attr = "value"
+        return transport
+
+    async def test_rate_limited_graphql_execution(self, mock_base_transport, rate_limit_config):
+        """Test that GraphQL requests are executed through rate limiter."""
+        mock_result = {"data": {"test": "result"}}
+        mock_base_transport.execute = AsyncMock(return_value=mock_result)
+
+        transport = RateLimitedHTTPXAsyncTransport(mock_base_transport, rate_limit_config)
+
+        async def mock_execute_requests(self, requests, handler):
+            results = []
+            for req in requests:
+                result = await handler(req)
+                results.append(result)
+            return results
+
+        with patch.object(RateLimiter, "execute_requests", mock_execute_requests):
+            result = await transport.execute("query { test }", variable_values={})
+
+        assert result == mock_result
+        mock_base_transport.execute.assert_called_once()
+        call_args = mock_base_transport.execute.call_args
+        assert call_args.kwargs.get("document") == "query { test }"
+        assert call_args.kwargs.get("variable_values") == {}
+
+    async def test_connect_wraps_async_client(self, rate_limit_config):
+        """Test that connect wraps AsyncClient with RateLimitedClient."""
+        from connector.httpx_rewrite import AsyncClient
+
+        class MockTransport:
+            def __init__(self):
+                self.client = None
+                self.connect_called = False
+
+            async def connect(self):
+                self.connect_called = True
+
+        real_client = AsyncClient(base_url="https://example.com")
+        base_transport = MockTransport()
+        base_transport.client = real_client
+
+        transport = RateLimitedHTTPXAsyncTransport(base_transport, rate_limit_config)
+        await transport.connect()
+
+        assert base_transport.connect_called
+        assert isinstance(base_transport.client, RateLimitedClient)
+        assert base_transport.client.base_client is real_client
+        assert transport.client is base_transport.client
+
+    async def test_connect_with_non_async_client(self, mock_base_transport, rate_limit_config):
+        """Test that connect does not wrap non-AsyncClient instances."""
+        mock_client = MagicMock()
+        mock_base_transport.client = mock_client
+        mock_base_transport.connect = AsyncMock()
+
+        transport = RateLimitedHTTPXAsyncTransport(mock_base_transport, rate_limit_config)
+        await transport.connect()
+
+        mock_base_transport.connect.assert_called_once()
+        assert mock_base_transport.client is mock_client
+        assert not isinstance(mock_base_transport.client, RateLimitedClient)
+        assert transport.client is mock_client
+
+    async def test_connect_without_client(self, mock_base_transport, rate_limit_config):
+        """Test that connect handles case when base transport has no client."""
+        mock_base_transport.connect = AsyncMock()
+        mock_base_transport.client = None
+
+        transport = RateLimitedHTTPXAsyncTransport(mock_base_transport, rate_limit_config)
+        await transport.connect()
+
+        mock_base_transport.connect.assert_called_once()
+        assert transport.client is None
+
+    async def test_connect_without_client_attribute(self, rate_limit_config):
+        """Test that connect handles case when base transport doesn't have client attribute.
+
+        This tests the branch on line 195 where hasattr(self.base_transport, "client") returns False.
+        We use a mock that raises AttributeError when accessing the client attribute to test this branch.
+        """
+
+        class MockTransportWithoutClient:
+            def __init__(self):
+                self.connect_called = False
+
+            async def connect(self):
+                self.connect_called = True
+
+            def __getattribute__(self, name):
+                if name == "client":
+                    raise AttributeError(
+                        f"'{type(self).__name__}' object has no attribute 'client'"
+                    )
+                return super().__getattribute__(name)
+
+        mock_base_transport = MockTransportWithoutClient()
+
+        transport = RateLimitedHTTPXAsyncTransport(mock_base_transport, rate_limit_config)
+
+        assert not hasattr(transport.base_transport, "client")
+
+        with pytest.raises(AttributeError):
+            await transport.connect()
+
+    async def test_connect_isinstance_async_client_true(self, rate_limit_config):
+        """Test that connect wraps AsyncClient when isinstance check is True, covering line 197 positive branch."""
+        from connector.httpx_rewrite import AsyncClient
+
+        class MockTransport:
+            def __init__(self):
+                self.client = None
+                self.connect_called = False
+
+            async def connect(self):
+                self.connect_called = True
+
+        real_client = AsyncClient(base_url="https://example.com")
+        base_transport = MockTransport()
+        base_transport.client = real_client
+
+        transport = RateLimitedHTTPXAsyncTransport(base_transport, rate_limit_config)
+        await transport.connect()
+
+        assert base_transport.connect_called
+        assert isinstance(base_transport.client, RateLimitedClient)
+        assert base_transport.client.base_client is real_client
+        assert transport.client is base_transport.client
+
+    async def test_connect_isinstance_async_client_false(self, rate_limit_config):
+        """Test that connect does not wrap when isinstance check is False, covering line 197 negative branch."""
+        from connector.httpx_rewrite import AsyncClient
+
+        class MockTransport:
+            def __init__(self):
+                self.client = None
+                self.connect_called = False
+
+            async def connect(self):
+                self.connect_called = True
+
+        mock_client = MagicMock()
+        mock_client.some_attr = "value"
+        base_transport = MockTransport()
+        base_transport.client = mock_client
+
+        transport = RateLimitedHTTPXAsyncTransport(base_transport, rate_limit_config)
+        await transport.connect()
+
+        assert base_transport.connect_called
+        assert base_transport.client is mock_client
+        assert not isinstance(base_transport.client, RateLimitedClient)
+        assert not isinstance(base_transport.client, AsyncClient)
+        assert transport.client is mock_client
+
+    async def test_execute_with_all_parameters(self, mock_base_transport, rate_limit_config):
+        """Test execute method with all GraphQL parameters."""
+        mock_result = {"data": {"test": "result"}}
+        mock_base_transport.execute = AsyncMock(return_value=mock_result)
+
+        transport = RateLimitedHTTPXAsyncTransport(mock_base_transport, rate_limit_config)
+
+        async def mock_execute_requests(self, requests, handler):
+            results = []
+            for req in requests:
+                result = await handler(req)
+                results.append(result)
+            return results
+
+        with patch.object(RateLimiter, "execute_requests", mock_execute_requests):
+            result = await transport.execute(
+                "query { test }", variable_values={"var": "value"}, operation_name="TestQuery"
+            )
+
+        assert result == mock_result
+        mock_base_transport.execute.assert_called_once()
+        call_args = mock_base_transport.execute.call_args
+        assert call_args.kwargs.get("document") == "query { test }"
+        assert call_args.kwargs.get("variable_values") == {"var": "value"}
+        assert call_args.kwargs.get("operation_name") == "TestQuery"
+
+    async def test_execute_calls_request_func_through_handler(
+        self, mock_base_transport, rate_limit_config
+    ):
+        """Test that execute calls request_func() through the handler lambda (lines 209-214).
+
+        This verifies that:
+        1. A request_func closure is created (line 208-211) that calls base_transport.execute
+        2. The request_func is passed to rate_limiter.execute_requests (line 214)
+        3. The handler lambda (lambda x: x()) is used to call request_func (line 214)
+        4. request_func() executes and calls base_transport.execute
+        """
+        mock_result = {"data": {"test": "result"}}
+        mock_base_transport.execute = AsyncMock(return_value=mock_result)
+
+        transport = RateLimitedHTTPXAsyncTransport(mock_base_transport, rate_limit_config)
+
+        async def mock_execute_requests(self, requests, handler):
+            results = []
+            for req in requests:
+                result = await handler(req)
+                results.append(result)
+            return results
+
+        with patch.object(RateLimiter, "execute_requests", mock_execute_requests):
+            result = await transport.execute("query { test }", variable_values={"var": "value"})
+
+        assert result == mock_result
+        mock_base_transport.execute.assert_called_once()
+
+        call_args = mock_base_transport.execute.call_args
+        assert call_args.kwargs.get("document") == "query { test }"
+        assert call_args.kwargs.get("variable_values") == {"var": "value"}
+
+    async def test_execute_raises_on_no_response(self, mock_base_transport, rate_limit_config):
+        """Test that execute raises ConnectorError when rate limiter returns no response.
+
+        This tests the negative branch of line 215 where responses is falsy (empty list).
+        """
+        transport = RateLimitedHTTPXAsyncTransport(mock_base_transport, rate_limit_config)
+
+        async def empty_execute_requests(self, requests, handler):
+            return []
+
+        with patch.object(RateLimiter, "execute_requests", empty_execute_requests):
+            with pytest.raises(ConnectorError) as exc_info:
+                await transport.execute("query { test }")
+
+            assert exc_info.value.message == "No response from GraphQL API"
+            assert exc_info.value.error_code == ErrorCode.API_ERROR
+
+    async def test_execute_with_successful_response(self, mock_base_transport, rate_limit_config):
+        """Test that execute returns response when rate limiter returns a response.
+
+        This tests the positive branch of line 215 where responses is truthy (non-empty list).
+        We patch at the class level to ensure the mock is actually used.
+        """
+        mock_result = {"data": {"test": "result"}}
+        mock_base_transport.execute = AsyncMock(return_value=mock_result)
+
+        transport = RateLimitedHTTPXAsyncTransport(mock_base_transport, rate_limit_config)
+
+        async def mock_execute_requests(self, requests, handler):
+            results = []
+            for req in requests:
+                result = await handler(req)
+                results.append(result)
+            return results
+
+        with patch.object(RateLimiter, "execute_requests", mock_execute_requests):
+            result = await transport.execute("query { test }")
+
+        assert result == mock_result
+        mock_base_transport.execute.assert_called_once()
+
+    def test_delegation_and_state(self, mock_base_transport, rate_limit_config):
+        """Test attribute delegation and state retrieval."""
+        mock_base_transport.delegated_method = MagicMock(return_value="delegated_value")
+        transport = RateLimitedHTTPXAsyncTransport(mock_base_transport, rate_limit_config)
+
+        assert transport.some_attr == "value"
+        assert transport.delegated_method() == "delegated_value"
+
+        config, delay = transport.get_state()
+        assert config is rate_limit_config
+        assert isinstance(delay, float)
+
+
+class TestBaseGraphQLSession:
+    """Test cases for BaseGraphQLSession."""
+
+    class ConcreteTestGraphQLSession(BaseGraphQLSession):
+        """Concrete implementation for testing."""
+
+        @classmethod
+        def prepare_client_args(cls, args: Request) -> dict[str, Any]:
+            from connector.httpx_rewrite import HTTPXAsyncTransport
+
+            return {
+                "transport": HTTPXAsyncTransport(
+                    url="https://example.com/graphql",
+                    auth=BearerAuth(
+                        token=get_token_auth(args).token,
+                        token_prefix="",
+                        auth_header="Authorization",
+                    ),
+                ),
+            }
+
+    def test_build_client_applies_rate_limiting(self, sample_request, rate_limit_config):
+        """Test that build_client wraps transport with rate limiter when config provided."""
+        client = self.ConcreteTestGraphQLSession.build_client(sample_request, rate_limit_config)
+
+        assert client is not None
+        assert isinstance(client.transport, RateLimitedHTTPXAsyncTransport)
+        assert client.transport.rate_limiter.config is rate_limit_config
+
+        client_no_rate_limit = self.ConcreteTestGraphQLSession.build_client(sample_request)
+        assert client_no_rate_limit is not None
+        assert not isinstance(client_no_rate_limit.transport, RateLimitedHTTPXAsyncTransport)
+
+        class TestSessionWithClassConfig(self.ConcreteTestGraphQLSession):
+            _rate_limit_config = RateLimitConfig(
+                app_id="class-config",
+                requests_per_window=5,
+                window_seconds=30,
+            )
+
+        client_with_class_config = TestSessionWithClassConfig.build_client(sample_request)
+        assert isinstance(client_with_class_config.transport, RateLimitedHTTPXAsyncTransport)
+
+    def test_get_current_rate_limits(self, sample_request, rate_limit_config):
+        """Test rate limit state retrieval."""
+        client_with_rate_limit = self.ConcreteTestGraphQLSession.build_client(
+            sample_request, rate_limit_config
+        )
+        session_with_rate_limit = self.ConcreteTestGraphQLSession(sample_request, rate_limit_config)
+        session_with_rate_limit.client = client_with_rate_limit
+
+        config, delay = session_with_rate_limit.get_current_rate_limits()
+        assert config is rate_limit_config
+        assert isinstance(delay, float)
+
+        client_no_rate_limit = self.ConcreteTestGraphQLSession.build_client(sample_request)
+        session_no_rate_limit = self.ConcreteTestGraphQLSession(sample_request)
+        session_no_rate_limit.client = client_no_rate_limit
+
+        config, delay = session_no_rate_limit.get_current_rate_limits()
+        assert config is None
+        assert delay == 0
+
+    async def test_context_manager(self, sample_request):
+        """Test async context manager functionality."""
+        with patch("connector.oai.base_clients.Client") as mock_client_class:
+            mock_client = MagicMock()
+            mock_client.__aenter__ = AsyncMock(return_value=mock_client)
+            mock_client.__aexit__ = AsyncMock()
+            mock_client_class.return_value = mock_client
+
+            session = self.ConcreteTestGraphQLSession(sample_request)
+
+            async with session as ctx:
+                assert ctx is session
+
+            mock_client.__aenter__.assert_called_once()
+            mock_client.__aexit__.assert_called_once()
+
+            test_exception = ValueError("Test error")
+            with pytest.raises(ValueError, match="Test error"):
+                async with session:
+                    raise test_exception

connector_py-4.180.0/connector/__about__.py

@@ -1 +0,0 @@
-__version__ = "4.180.0"