Refactor APIRequest class to use Record type for headers and body

This commit is contained in:
Rodrigo Nader 2024-03-31 00:47:56 -03:00
commit c71c8f4fee

View file

@ -27,12 +27,12 @@ class APIRequest(CustomComponent):
"headers": { "headers": {
"display_name": "Headers", "display_name": "Headers",
"info": "The headers to send with the request.", "info": "The headers to send with the request.",
"input_types": ["dict"], "input_types": ["Record"],
}, },
"body": { "body": {
"display_name": "Body", "display_name": "Body",
"info": "The body to send with the request (for POST, PATCH, PUT).", "info": "The body to send with the request (for POST, PATCH, PUT).",
"input_types": ["dict"], "input_types": ["Record"],
}, },
"timeout": { "timeout": {
"display_name": "Timeout", "display_name": "Timeout",
@ -47,8 +47,8 @@ class APIRequest(CustomComponent):
client: httpx.AsyncClient, client: httpx.AsyncClient,
method: str, method: str,
url: str, url: str,
headers: Optional[dict] = None, headers: Optional[Record] = None,
body: Optional[dict] = None, body: Optional[Record] = None,
timeout: int = 5, timeout: int = 5,
) -> Record: ) -> Record:
method = method.upper() method = method.upper()
@ -58,7 +58,9 @@ class APIRequest(CustomComponent):
data = body if body else None data = body if body else None
payload = json.dumps(data) payload = json.dumps(data)
try: try:
response = await client.request(method, url, headers=headers, content=payload, timeout=timeout) response = await client.request(
method, url, headers=headers, content=payload, timeout=timeout
)
try: try:
result = response.json() result = response.json()
except Exception: except Exception:
@ -94,24 +96,31 @@ class APIRequest(CustomComponent):
self, self,
method: str, method: str,
urls: List[str], urls: List[str],
headers: Optional[dict] = None, headers: Optional[Record] = None,
body: Optional[List[Record]] = None, body: Optional[Record] = None,
timeout: int = 5, timeout: int = 5,
) -> List[Record]: ) -> List[Record]:
if headers is None: if headers is None:
headers = {} headers = {}
else:
headers = headers.data
bodies = [] bodies = []
if body: if body:
if isinstance(body, list): if isinstance(body, list):
bodies = [b.data for b in body] bodies = [b.data for b in body]
else: else:
bodies = [body.data] bodies = [body.data]
if len(urls) != len(bodies): if len(urls) != len(bodies):
# add bodies with None # add bodies with None
bodies += [None] * (len(urls) - len(bodies)) # type: ignore bodies += [None] * (len(urls) - len(bodies)) # type: ignore
async with httpx.AsyncClient() as client: async with httpx.AsyncClient() as client:
results = await asyncio.gather( results = await asyncio.gather(
*[self.make_request(client, method, u, headers, rec, timeout) for u, rec in zip(urls, bodies)] *[
self.make_request(client, method, u, headers, rec, timeout)
for u, rec in zip(urls, bodies)
]
) )
self.status = results self.status = results
return results return results