port-ocean 0.12.2.dev10__py3-none-any.whl → 0.12.2.dev12__py3-none-any.whl

port_ocean/cli/commands/list_integrations.py

@@ -1,5 +1,7 @@
 # -*- coding: utf-8 -*-
-import httpx
+import asyncio
+
+import aiohttp
 
 from port_ocean.cli.commands.main import cli_start, console
 
@@ -9,16 +11,17 @@ def list_git_folders(owner: str, repo_name: str, path: str) -> list[str]:
     api_url = f"https://api.github.com/repos/{owner}/{repo_name}/contents/{path}"
 
     # Send a GET request to the API
-    response = httpx.get(api_url)
+    response = asyncio.run(aiohttp.ClientSession().get(api_url))
 
     # Check if the request was successful
-    if response.is_error:
+    if not response.ok:
+        text = asyncio.run(response.text())
         console.print(
-            f"[bold red]Failed to list folders.[/bold red] Status Code: {response.status_code}, Error: {response.text}"
+            f"[bold red]Failed to list folders.[/bold red] Status Code: {response.status}, Error: {text}"
         )
         exit(1)
 
-    contents = response.json()
+    contents = await response.json()
     folders = [item["name"] for item in contents if item["type"] == "dir"]
     return folders
 

port_ocean/cli/commands/pull.py

@@ -1,26 +1,26 @@
+import asyncio
 import os
-import shutil
-from io import BytesIO
 
+import aiohttp
 import click
-import httpx
 
 from port_ocean.cli.commands.main import cli_start, console
 
 
 def download_github_folder(
-    owner: str, repo_name: str, folder_path: str, destination_path: str
+        owner: str, repo_name: str, folder_path: str, destination_path: str
 ) -> None:
     # Construct the API URL to get the contents of the folder
     api_url = f"https://api.github.com/repos/{owner}/{repo_name}/contents/{folder_path}"
 
     # Send a GET request to the API
-    response = httpx.get(api_url)
+    response = asyncio.run(aiohttp.ClientSession().request("GET", api_url))
 
     # Check if the request was successful
     if response.is_error:
+        text = asyncio.run(response.text())
         console.print(
-            f"[bold red]Failed to download the folder `{folder_path}`.[/bold red] Status Code: {response.status_code}, Error: {response.text}"
+            f"[bold red]Failed to download the folder `{folder_path}`.[/bold red] Status Code: {response.status}, Error: {text}"
         )
         exit(1)
 
@@ -28,23 +28,27 @@ def download_github_folder(
     if not os.path.exists(destination_path):
         os.makedirs(destination_path)
 
+    async def read_async():
+        async with aiohttp.ClientSession().request("GET", file_url) as file_response:
+            if file_response.status == 200:
+                with open(file_name, "wb") as file:
+                    file.write(await file_response.read())
+            else:
+                text = asyncio.run(file_response.text())
+                console.print(
+                    f"[bold red]Failed to download file `{content['name']}`.[/bold red] Status code: {file_response.status}, Error: {text}"
+                )
+                exit(1)
+
     # Iterate over the files and download them
-    repo_contents = response.json()
+    repo_contents = await response.json()
     for content in repo_contents:
         if content["type"] == "file":
             file_url = content["download_url"]
             file_name = os.path.join(destination_path, content["name"])
 
             # Download the file
-            with httpx.stream("GET", file_url) as file_response:
-                if file_response.status_code == 200:
-                    with open(file_name, "wb") as file:
-                        shutil.copyfileobj(BytesIO(file_response.content), file)
-                else:
-                    console.print(
-                        f"[bold red]Failed to download file `{content['name']}`.[/bold red] Status code: {file_response.status_code}, Error: {file_response.text}"
-                    )
-                    exit(1)
+            asyncio.run(read_async())
 
     console.print(f"Folder `{folder_path}` downloaded successfully!")
 

port_ocean/clients/port/authentication.py

@@ -1,6 +1,6 @@
 from typing import Any
 
-import httpx
+from aiohttp import ClientSession
 from loguru import logger
 from pydantic import BaseModel, Field, PrivateAttr
 
@@ -26,14 +26,14 @@ class TokenResponse(BaseModel):
 
 class PortAuthentication:
     def __init__(
-        self,
-        client: httpx.AsyncClient,
-        client_id: str,
-        client_secret: str,
-        api_url: str,
-        integration_identifier: str,
-        integration_type: str,
-        integration_version: str,
+            self,
+            client: ClientSession,
+            client_id: str,
+            client_secret: str,
+            api_url: str,
+            integration_identifier: str,
+            integration_type: str,
+            integration_version: str,
     ):
         self.client = client
         self.api_url = api_url
@@ -51,10 +51,9 @@ class PortAuthentication:
         response = await self.client.post(
             f"{self.api_url}/auth/access_token",
             json=credentials,
-            extensions={"retryable": True},
         )
-        handle_status_code(response)
-        return TokenResponse(**response.json())
+        await handle_status_code(response)
+        return TokenResponse(**(await response.json()))
 
     def user_agent(self, user_agent_type: UserAgentType | None = None) -> str:
         user_agent = f"port-ocean/{self.integration_type}/{self.integration_version}/{self.integration_identifier}"
@@ -64,7 +63,7 @@ class PortAuthentication:
         return user_agent
 
     async def headers(
-        self, user_agent_type: UserAgentType | None = None
+            self, user_agent_type: UserAgentType | None = None
     ) -> dict[Any, Any]:
         return {
             "Authorization": await self.token,

port_ocean/clients/port/client.py

@@ -1,3 +1,5 @@
+from typing import Any
+
 from loguru import logger
 
 from port_ocean.clients.port.authentication import PortAuthentication
@@ -13,7 +15,6 @@ from port_ocean.clients.port.utils import (
     get_internal_http_client,
 )
 from port_ocean.exceptions.clients import KafkaCredentialsNotFound
-from typing import Any
 
 
 class PortClient(
@@ -56,9 +57,9 @@ class PortClient(
         )
         if response.is_error:
             logger.error("Error getting kafka credentials")
-        handle_status_code(response)
+        await handle_status_code(response)
 
-        credentials = response.json().get("credentials")
+        credentials = (await response.json()).get("credentials")
 
         if credentials is None:
             raise KafkaCredentialsNotFound("No kafka credentials found")
@@ -72,10 +73,10 @@ class PortClient(
             f"{self.api_url}/organization", headers=await self.auth.headers()
         )
         if response.is_error:
-            logger.error(f"Error getting organization id, error: {response.text}")
-        handle_status_code(response)
+            logger.error(f"Error getting organization id, error: {await response.text()}")
+        await handle_status_code(response)
 
-        return response.json()["organization"]["id"]
+        return (await response.json())["organization"]["id"]
 
     async def update_integration_state(
         self, state: dict[str, Any], should_raise: bool = True, should_log: bool = True
@@ -87,8 +88,8 @@ class PortClient(
             headers=await self.auth.headers(),
             json=state,
         )
-        handle_status_code(response, should_raise, should_log)
+        await handle_status_code(response, should_raise, should_log)
         if response.is_success and should_log:
             logger.info("Integration resync state updated successfully")
 
-        return response.json().get("integration", {})
+        return (await response.json()).get("integration", {})

port_ocean/clients/port/mixins/blueprints.py

@@ -1,6 +1,6 @@
 from typing import Any
 
-import httpx
+import aiohttp
 from loguru import logger
 
 from port_ocean.clients.port.authentication import PortAuthentication
@@ -10,7 +10,7 @@ from port_ocean.core.models import Blueprint
 
 
 class BlueprintClientMixin:
-    def __init__(self, auth: PortAuthentication, client: httpx.AsyncClient):
+    def __init__(self, auth: PortAuthentication, client: aiohttp.ClientSession):
         self.auth = auth
         self.client = client
 
@@ -22,8 +22,8 @@ class BlueprintClientMixin:
             f"{self.auth.api_url}/blueprints/{identifier}",
             headers=await self.auth.headers(),
         )
-        handle_status_code(response, should_log=should_log)
-        return Blueprint.parse_obj(response.json()["blueprint"])
+        await handle_status_code(response, should_log=should_log)
+        return Blueprint.parse_obj((await response.json())["blueprint"])
 
     async def create_blueprint(
         self,
@@ -35,8 +35,8 @@ class BlueprintClientMixin:
         response = await self.client.post(
             f"{self.auth.api_url}/blueprints", headers=headers, json=raw_blueprint
         )
-        handle_status_code(response)
-        return response.json()["blueprint"]
+        await handle_status_code(response)
+        return (await response.json())["blueprint"]
 
     async def patch_blueprint(
         self,
@@ -51,7 +51,7 @@ class BlueprintClientMixin:
             headers=headers,
             json=raw_blueprint,
         )
-        handle_status_code(response)
+        await handle_status_code(response)
 
     async def delete_blueprint(
         self,
@@ -70,7 +70,7 @@ class BlueprintClientMixin:
                 f"{self.auth.api_url}/blueprints/{identifier}",
                 headers=headers,
             )
-            handle_status_code(response, should_raise)
+            await handle_status_code(response, should_raise)
             return None
         else:
             response = await self.client.delete(
@@ -78,8 +78,8 @@ class BlueprintClientMixin:
                 headers=await self.auth.headers(),
             )
 
-            handle_status_code(response, should_raise)
-            return response.json().get("migrationId", "")
+            await handle_status_code(response, should_raise)
+            return (await response.json()).get("migrationId", "")
 
     async def create_action(
         self, action: dict[str, Any], should_log: bool = True
@@ -91,7 +91,7 @@ class BlueprintClientMixin:
             headers=await self.auth.headers(),
         )
 
-        handle_status_code(response, should_log=should_log)
+        await handle_status_code(response, should_log=should_log)
 
     async def create_scorecard(
         self,
@@ -106,7 +106,7 @@ class BlueprintClientMixin:
             headers=await self.auth.headers(),
         )
 
-        handle_status_code(response, should_log=should_log)
+        await handle_status_code(response, should_log=should_log)
 
     async def create_page(
         self, page: dict[str, Any], should_log: bool = True
@@ -118,7 +118,7 @@ class BlueprintClientMixin:
             headers=await self.auth.headers(),
         )
 
-        handle_status_code(response, should_log=should_log)
+        await handle_status_code(response, should_log=should_log)
         return page
 
     async def delete_page(
@@ -132,4 +132,4 @@ class BlueprintClientMixin:
             headers=await self.auth.headers(),
         )
 
-        handle_status_code(response, should_raise)
+        await handle_status_code(response, should_raise)

port_ocean/clients/port/mixins/entities.py

@@ -2,7 +2,7 @@ import asyncio
 from typing import Any
 from urllib.parse import quote_plus
 
-import httpx
+import aiohttp
 from loguru import logger
 
 from port_ocean.clients.port.authentication import PortAuthentication
@@ -15,7 +15,7 @@ from port_ocean.core.models import Entity
 
 
 class EntityClientMixin:
-    def __init__(self, auth: PortAuthentication, client: httpx.AsyncClient):
+    def __init__(self, auth: PortAuthentication, client: aiohttp.ClientSession):
         self.auth = auth
         self.client = client
         # Semaphore is used to limit the number of concurrent requests to port, to avoid overloading it.
@@ -56,8 +56,8 @@ class EntityClientMixin:
                 f"entity: {entity.identifier} of "
                 f"blueprint: {entity.blueprint}"
             )
-        handle_status_code(response, should_raise)
-        result = response.json()
+        await handle_status_code(response, should_raise)
+        result = await response.json()
 
         result_entity = (
             Entity.parse_obj(result["entity"]) if result.get("entity") else entity
@@ -149,7 +149,7 @@ class EntityClientMixin:
                     f"blueprint: {entity.blueprint}"
                 )
 
-            handle_status_code(response, should_raise)
+            await handle_status_code(response, should_raise)
 
     async def batch_delete_entities(
         self,
@@ -204,11 +204,10 @@ class EntityClientMixin:
                 "exclude_calculated_properties": "true",
                 "include": ["blueprint", "identifier"],
             },
-            extensions={"retryable": True},
             timeout=30,
         )
-        handle_status_code(response)
-        return [Entity.parse_obj(result) for result in response.json()["entities"]]
+        await handle_status_code(response)
+        return [Entity.parse_obj(result) for result in (await response.json())["entities"]]
 
     async def search_batch_entities(
         self, user_agent_type: UserAgentType, entities_to_search: list[Entity]

port_ocean/clients/port/mixins/integrations.py

@@ -1,7 +1,7 @@
 from typing import Any, TYPE_CHECKING, Optional, TypedDict
 from urllib.parse import quote_plus
 
-import httpx
+import aiohttp
 from loguru import logger
 
 from port_ocean.clients.port.authentication import PortAuthentication
@@ -22,7 +22,7 @@ class IntegrationClientMixin:
         integration_identifier: str,
         integration_version: str,
         auth: PortAuthentication,
-        client: httpx.AsyncClient,
+        client: aiohttp.ClientSession,
     ):
         self.integration_identifier = integration_identifier
         self.integration_version = integration_version
@@ -30,7 +30,7 @@ class IntegrationClientMixin:
         self.client = client
         self._log_attributes: LogAttributes | None = None
 
-    async def _get_current_integration(self) -> httpx.Response:
+    async def _get_current_integration(self) -> aiohttp.ClientResponse:
         logger.info(f"Fetching integration with id: {self.integration_identifier}")
         response = await self.client.get(
             f"{self.auth.api_url}/integration/{self.integration_identifier}",
@@ -42,8 +42,8 @@ class IntegrationClientMixin:
         self, should_raise: bool = True, should_log: bool = True
     ) -> dict[str, Any]:
         response = await self._get_current_integration()
-        handle_status_code(response, should_raise, should_log)
-        return response.json().get("integration", {})
+        await handle_status_code(response, should_raise, should_log)
+        return (await response.json()).get("integration", {})
 
     async def get_log_attributes(self) -> LogAttributes:
         if self._log_attributes is None:
@@ -71,8 +71,8 @@ class IntegrationClientMixin:
         response = await self.client.post(
             f"{self.auth.api_url}/integration", headers=headers, json=json
         )
-        handle_status_code(response)
-        return response.json()["integration"]
+        await handle_status_code(response)
+        return (await response.json())["integration"]
 
     async def patch_integration(
         self,
@@ -96,8 +96,8 @@ class IntegrationClientMixin:
             headers=headers,
             json=json,
         )
-        handle_status_code(response)
-        return response.json()["integration"]
+        await handle_status_code(response)
+        return (await response.json())["integration"]
 
     async def ingest_integration_logs(self, logs: list[dict[str, Any]]) -> None:
         logger.debug("Ingesting logs")
@@ -110,7 +110,7 @@ class IntegrationClientMixin:
                 "logs": logs,
             },
         )
-        handle_status_code(response, should_log=False)
+        await handle_status_code(response, should_log=False)
         logger.debug("Logs successfully ingested")
 
     async def ingest_integration_kind_examples(
@@ -125,5 +125,5 @@ class IntegrationClientMixin:
                 "examples": sensitive_log_filter.mask_object(data, full_hide=True),
             },
         )
-        handle_status_code(response, should_log=should_log)
+        await handle_status_code(response, should_log=should_log)
         logger.debug(f"Examples for kind {kind} successfully ingested")

port_ocean/clients/port/mixins/migrations.py

@@ -1,6 +1,6 @@
 import asyncio
 
-import httpx
+import aiohttp
 from loguru import logger
 
 from port_ocean.clients.port.authentication import PortAuthentication
@@ -9,7 +9,7 @@ from port_ocean.core.models import Migration
 
 
 class MigrationClientMixin:
-    def __init__(self, auth: PortAuthentication, client: httpx.AsyncClient):
+    def __init__(self, auth: PortAuthentication, client: aiohttp.ClientSession):
         self.auth = auth
         self.client = client
 
@@ -28,9 +28,9 @@ class MigrationClientMixin:
             headers=headers,
         )
 
-        handle_status_code(response, should_raise=True)
+        await handle_status_code(response, should_raise=True)
 
-        migration_status = response.json().get("migration", {}).get("status", None)
+        migration_status = (await response.json()).get("migration", {}).get("status", None)
         if (
             migration_status == "RUNNING"
             or migration_status == "INITIALIZING"
@@ -43,4 +43,4 @@ class MigrationClientMixin:
                 f"Migration with id: {migration_id} finished with status {migration_status}",
             )
 
-        return Migration.parse_obj(response.json()["migration"])
+        return Migration.parse_obj(await response.json()["migration"])

port_ocean/clients/port/retry_transport.py

@@ -1,51 +1,29 @@
-import asyncio
 from http import HTTPStatus
-from typing import TYPE_CHECKING, Any
 
-import httpx
+from aiohttp import ClientResponse
 
-from port_ocean.helpers.retry import RetryTransport
+from port_ocean.helpers.retry import RetryRequestClass
 
-if TYPE_CHECKING:
-    from port_ocean.clients.port.client import PortClient
 
+class TokenRetryRequestClass(RetryRequestClass):
+    async def _handle_unauthorized(self) -> None:
+        token = await self._session.port_client.auth.token
+        self.headers["Authorization"] = f"Bearer {token}"
 
-class TokenRetryTransport(RetryTransport):
-    def __init__(self, port_client: "PortClient", **kwargs: Any) -> None:
-        super().__init__(**kwargs)
-        self.port_client = port_client
-
-    async def _handle_unauthorized(self, response: httpx.Response) -> None:
-        token = await self.port_client.auth.token
-        response.headers["Authorization"] = f"Bearer {token}"
-
-    def is_token_error(self, response: httpx.Response) -> bool:
+    def is_token_error(self, response: ClientResponse) -> bool:
         return (
-            response.status_code == HTTPStatus.UNAUTHORIZED
-            and "/auth/access_token" not in str(response.request.url)
-            and self.port_client.auth.last_token_object is not None
-            and self.port_client.auth.last_token_object.expired
+                response.status == HTTPStatus.UNAUTHORIZED
+                and "/auth/access_token" not in str(self.url)
+                and self._session.port_client.auth.last_token_object is not None
+                and self._session.port_client.auth.last_token_object.expired
         )
 
-    async def _should_retry_async(self, response: httpx.Response) -> bool:
+    async def _should_retry_async(self, response: ClientResponse) -> bool:
         if self.is_token_error(response):
             if self._logger:
                 self._logger.info(
                     "Got unauthorized response, trying to refresh token before retrying"
                 )
-            await self._handle_unauthorized(response)
+            await self._handle_unauthorized()
             return True
         return await super()._should_retry_async(response)
-
-    def _should_retry(self, response: httpx.Response) -> bool:
-        if self.is_token_error(response):
-            if self._logger:
-                self._logger.info(
-                    "Got unauthorized response, trying to refresh token before retrying"
-                )
-            asyncio.get_running_loop().run_until_complete(
-                self._handle_unauthorized(response)
-            )
-
-            return True
-        return super()._should_retry(response)

port_ocean/clients/port/utils.py

@@ -1,11 +1,12 @@
+from asyncio import ensure_future
 from typing import TYPE_CHECKING
 
-import httpx
+import aiohttp
 from loguru import logger
 from werkzeug.local import LocalStack, LocalProxy
 
-from port_ocean.clients.port.retry_transport import TokenRetryTransport
-from port_ocean.helpers.async_client import OceanAsyncClient
+from port_ocean.helpers.retry import RetryRequestClass
+from port_ocean.utils.signal import signal_handler
 
 if TYPE_CHECKING:
     from port_ocean.clients.port.client import PortClient
@@ -20,33 +21,41 @@ PORT_HTTP_MAX_CONNECTIONS_LIMIT = 200
 PORT_HTTP_MAX_KEEP_ALIVE_CONNECTIONS = 50
 PORT_HTTP_TIMEOUT = 60.0
 
-PORT_HTTPX_TIMEOUT = httpx.Timeout(PORT_HTTP_TIMEOUT)
-PORT_HTTPX_LIMITS = httpx.Limits(
-    max_connections=PORT_HTTP_MAX_CONNECTIONS_LIMIT,
-    max_keepalive_connections=PORT_HTTP_MAX_KEEP_ALIVE_CONNECTIONS,
-)
+TIMEOUT = aiohttp.ClientTimeout(total=PORT_HTTP_TIMEOUT)
 
-_http_client: LocalStack[httpx.AsyncClient] = LocalStack()
+_http_client: LocalStack[aiohttp.ClientSession] = LocalStack()
 
 
-def _get_http_client_context(port_client: "PortClient") -> httpx.AsyncClient:
+class PortRetryRequestClass(RetryRequestClass):
+    RETRYABLE_ROUTES = frozenset(["/auth/access_token", "/entities/search"])
+
+    def _is_retryable(self) -> bool:
+        return super()._is_retryable() or any([route in self.url.path for route in self.RETRYABLE_ROUTES])
+
+
+class OceanPortAsyncClient(aiohttp.ClientSession):
+    def __init__(self, port_client: "PortClient", *args,
+                 **kwargs):
+        self._port_client = port_client
+        super().__init__(*args, **kwargs)
+
+
+def _get_http_client_context(port_client: "PortClient") -> aiohttp.ClientSession:
     client = _http_client.top
     if client is None:
-        client = OceanAsyncClient(
-            TokenRetryTransport,
-            transport_kwargs={"port_client": port_client},
-            timeout=PORT_HTTPX_TIMEOUT,
-            limits=PORT_HTTPX_LIMITS,
-        )
+        AIOHTTP_CONNECTOR = aiohttp.TCPConnector(limit=PORT_HTTP_MAX_CONNECTIONS_LIMIT, force_close=True)
+        client = OceanPortAsyncClient(port_client, timeout=TIMEOUT, request_class=PortRetryRequestClass,
+                                      connector=AIOHTTP_CONNECTOR,
+                                      trust_env=True)
         _http_client.push(client)
-
+        signal_handler.register(lambda: ensure_future(client.close()))
     return client
 
 
-_port_internal_async_client: httpx.AsyncClient = None  # type: ignore
+_port_internal_async_client: aiohttp.ClientSession = None  # type: ignore
 
 
-def get_internal_http_client(port_client: "PortClient") -> httpx.AsyncClient:
+def get_internal_http_client(port_client: "PortClient") -> aiohttp.ClientSession:
     global _port_internal_async_client
     if _port_internal_async_client is None:
         _port_internal_async_client = LocalProxy(
@@ -56,12 +65,12 @@ def get_internal_http_client(port_client: "PortClient") -> httpx.AsyncClient:
     return _port_internal_async_client
 
 
-def handle_status_code(
-    response: httpx.Response, should_raise: bool = True, should_log: bool = True
+async def handle_status_code(
+        response: aiohttp.ClientResponse, should_raise: bool = True, should_log: bool = True
 ) -> None:
-    if should_log and response.is_error:
+    if should_log and not response.ok:
         logger.error(
-            f"Request failed with status code: {response.status_code}, Error: {response.text}"
+            f"Request failed with status code: {response.status}, Error: {await response.text()}"
         )
     if should_raise:
         response.raise_for_status()

port_ocean/core/defaults/clean.py

@@ -1,7 +1,7 @@
 import asyncio
 from typing import Type
 
-import httpx
+import aiohttp
 from loguru import logger
 
 from port_ocean.context.ocean import ocean
@@ -66,6 +66,6 @@ async def _clean_defaults(
                     for migration_id in migration_ids
                 )
             )
-    except httpx.HTTPStatusError as e:
-        logger.error(f"Failed to delete blueprints: {e.response.text}.")
+    except aiohttp.ClientResponseError as e:
+        logger.error(f"Failed to delete blueprints: {e.message}.")
         raise e

port_ocean/core/defaults/common.py

@@ -2,7 +2,7 @@ import json
 from pathlib import Path
 from typing import Type, Any, TypedDict, Optional
 
-import httpx
+import aiohttp
 import yaml
 from pydantic import BaseModel, Field
 from starlette import status
@@ -39,8 +39,8 @@ async def is_integration_exists(port_client: PortClient) -> bool:
     try:
         await port_client.get_current_integration(should_log=False)
         return True
-    except httpx.HTTPStatusError as e:
-        if e.response.status_code != status.HTTP_404_NOT_FOUND:
+    except aiohttp.ClientResponseError as e:
+        if e.status != status.HTTP_404_NOT_FOUND:
             raise e
 
     return False