Creating an Airflow Custom Hook for Reliable API Calls
Write your own Airflow custom Hook that makes API calls with retries, rate limits, and timeouts
Welcome to my 76th article on Medium. This is going to be a guide on crafting a custom Airflow hook to guarantee dependable API requests. In this blog, I will discuss how you can write your own custom Airflow Hook that is reliable and dynamic enough to handle failures🔥 and take care of rate limits ⚙️ and timeouts ⏲. Your data engineering workflows will gain resilience and effectiveness by incorporating these features.
Table of contents
· What do I mean when I say reliable? · The ReliableAPIHook, custom Hook for making API calls · The ReliableAPIHook in action, Dag
There are various ways to make API calls from your Airflow Dag. I have already discussed quite a few in my previous blogs.
This article is going to be my last one on this topic.
What do I mean when I say reliable?
- Retry on failure: When a request fails because of: 502 (Bad Gateway server error) 504 (Request Timed out) 429 (Rate limit exceed) 503 (service_unavailable), or 500 (internal_server_error), e.t.c
- Takes care of Rate limits
- Times out after response time crosses a threshold time
The ReliableHTTPBinAPIHook, custom Hook for making API calls
In this example, I have used the open-source REST API provided by www.httpbin.org. This piece of code is a hook; place it in the hooks package of your airflow repo in a file named http_bin.py
import logging
import time
from airflow.providers.http.hooks.http import HttpHook
from airflow.exceptions import AirflowException
import requests
from requests.exceptions import HTTPError
GET_ENDPOINT = "https://www.httpbin.org/get"
POST_ENDPOINT = "https://www.httpbin.org/post"
PUT_ENDPOINT = "https://www.httpbin.org/put"
PATCH_ENDPOINT = "https://www.httpbin.org/patch"
DELETE_ENDPOINT = "https://www.httpbin.org/delete"
def make_request(url, headers, data=None, method="GET", sleep_time=0.0, max_attempts=3, timeout=3600.00):
"""
Make requests to any API with an automatic retry and rate limit mechanism.
Raises:
HTTPError Exception: If response status code!= 200.
Returns:
Response: The response from the API call.
"""
# API is rate-limited, adjust the sleep time based on the rate limit of the API you are calling
time.sleep(sleep_time)
if method not in ["POST", "PATCH", "PUT", "GET", "DELETE"]:
raise ValueError("Method not supported. Available methods are: POST, PATCH, PUT, GET & DELETE")
kwargs = {"url": url, "headers": headers, "data": data, "timeout": timeout}
request_methods = {
'POST': requests.post,
'PATCH': requests.patch,
'PUT': requests.put,
'DELETE': requests.delete,
'GET': requests.get
}
for attempt in range(max_attempts):
try:
response = request_methods[method](**kwargs)
response.raise_for_status()
return response
except HTTPError as e:
if attempt < max_attempts - 1:
logging.warning(f"Retrying, Request failed, Error: {e}")
else:
raise e
class ReliableHTTPBinAPIHook:
"""
Make API Calls Reliable
:param conn_id: ID of your airflow connection, containing
Notion token in the 'password' field, and 'https' in the 'schema' field
:type conn_id: str
"""
def __init__(self, conn_id=None):
self.conn_id = conn_id
self.headers = {
# "Authorization": f"Bearer {self._get_token()}",
"Content-Type": "application/json",
}
# def _get_token(self):
# if self.conn_id:
# connection = HttpHook.get_connection(self.conn_id)
# password = connection.password
# return password
# else:
# raise AirflowException(
# "Cannot get token: No valid Notion token or conn_id supplied"
# )
def get_call(self):
return make_request(method="GET", url=GET_ENDPOINT, headers=self.headers, sleep_time=1, max_attempts=5, timeout=60)
def post_call(self):
return make_request(method="POST", url=POST_ENDPOINT, headers=self.headers, sleep_time=0.5)
def put_call(self):
return make_request(method="PUT", url=PUT_ENDPOINT, headers=self.headers, max_attempts=2)
def patch_call(self):
return make_request(method="PATCH", url=PATCH_ENDPOINT, headers=self.headers)
def delete_call(self):
return make_request(method="DELETE", url=DELETE_ENDPOINT, headers=self.headers, timeout=1)This code was good enough to suit my requirements; however, you can make minor modifications in the block of code below depending on what status code of the response you want to retry. For example,
for attempt in range(max_attempts):
try:
response = request_methods[method](**kwargs)
response.raise_for_status()
return response
except HTTPError as e:
if response.status_code == 502 and attempt < max_attempts - 1:
logging.warning(f"Retrying, Request failed, Error: {e}")
else:
raise eAuthorisation: In scenarios where your API requires authorization, uncomment the Authorization and _get_token part of the code. Make sure to set the connection in the Airflow Console properly, as discussed in my previous blogs attached above.
Rate Limit: Calculate the sleep time between each request based on the rate limit offered by your API.
Retry: Once your code is in production, you will encounter more 4XX errors than 5XX errors more often from the server for different reasons, as discussed above. It's a good practice to try 2-3 times before failing the job because these types of errors are typically fixable by simply retrying.
The ReliableHTTPBinAPIHook in action, Dag
import logging
from airflow import DAG
from airflow.operators.python import PythonOperator
from datetime import datetime
from hooks.http_bin import ReliableHTTPBinAPIHook
default_args = {
'owner': 'airflow',
'start_date': datetime(2023, 1, 1)
}
dag = DAG('Test.Reliable.API.Hook', default_args=default_args, schedule_interval=None, catchup=False)
def test_reliable_api_hook():
hook = ReliableHTTPBinAPIHook()
get_res = hook.get_call()
logging.info(f"GET: StatusCode={get_res.status_code}, Response={get_res.json()}")
post_res = hook.post_call()
logging.info(f"POST: StatusCode={post_res.status_code}, Response={post_res.json()}")
put_res = hook.put_call()
logging.info(f"PUT: StatusCode={put_res.status_code}, Response={put_res.json()}")
patch_res = hook.patch_call()
logging.info(f"PATCH: StatusCode={patch_res.status_code}, Response={patch_res.json()}")
del_res = hook.delete_call()
logging.info(f"DELETE: StatusCode={del_res.status_code}, Response={del_res.json()}")
with dag:
test_reliable_api_hook = PythonOperator(
task_id='test_reliable_api_hook',
python_callable=test_reliable_api_hook,
)
test_reliable_api_hook
Since these codes are in plain Python, don’t shy away from experimenting by modifying the code here and there.
📣 Checkout: Here is yet another article where I have discussed how you can write your own Airflow Custom Operator ⛑ and use this Hook. 🚀
Conclusion
You did a fantastic job implementing a custom Airflow hook for reliable API calls! By utilizing Airflow and incorporating features such as retries, rate limitations, and timeouts, you have been able to build robust data engineering operations. Please utilize the best practices, suggestions, and approaches provided in this tutorial to enhance data integration. With the knowledge you have acquired from my airflow series on making API calls, you can confidently tackle challenges and optimize Airflow!
Thanks for spending time on this article! Before you go:
- 👏 Applaud this article 🚀
- 📰 Subscribe to my upcoming blogs: Subscribe
- 🔔 Follow me on Medium, Linkedin, and Twitter
