Use Apache Kafka with Apache Airflow

Info

This page has not yet been updated for Airflow 3. The concepts shown are relevant, but some code may need to be updated. If you run any examples, take care to update import statements and watch for any other breaking changes.

Apache Kafka is an open source tool for handling event streaming. Combining Kafka and Airflow allows you to build powerful pipelines that integrate streaming data with batch processing. In this tutorial, you’ll learn how to install and use the Kafka Airflow provider to interact directly with Kafka topics.

Warning

While it is possible to manage a Kafka cluster with Airflow, be aware that Airflow itself should not be used for streaming or low-latency processes. See the Best practices section for more information.

Time to complete

This tutorial takes approximately 1 hour to complete.

Assumed knowledge

To get the most out of this tutorial, make sure you have an understanding of:

• The basics of Apache Kafka. See the official Introduction to Kafka.
• Airflow fundamentals, such as writing DAGs and defining tasks. See Get started with Apache Airflow.
• Airflow operators. See Operators 101.

Quickstart

If you have a GitHub account, you can use the quickstart repository for this tutorial, which automatically starts up Airflow, initiates a local Kafka cluster, and configures all necessary connections. Clone the quickstart repository and then skip to Step 6: Run the DAGs.

Prerequisites

• A Kafka cluster with a topic. This tutorial uses a cluster hosted by Confluent Cloud, which has a free trial option. See the Confluent documentation for how to create a Kafka cluster and topic in Confluent Cloud.
• The Astro CLI.

Info

To connect a local Kafka cluster to an Airflow instance running in Docker, set the following properties in your Kafka cluster’s server.properties file before starting your Kafka cluster:

listeners=PLAINTEXT://:9092,DOCKER_HACK://:19092
advertised.listeners=PLAINTEXT://localhost:9092,DOCKER_HACK://host.docker.internal:19092
listener.security.protocol.map=PLAINTEXT:PLAINTEXT,DOCKER_HACK:PLAINTEXT

You can learn more about connecting to local Kafka from within a Docker container in Confluent’s Documentation.

Step 1: Configure your Astro project

1. Create a new Astro project:

   1
   $ mkdir astro-kafka-tutorial && cd astro-kafka-tutorial
   2
   $ astro dev init

2. Add the following packages to your packages.txt file:

   build-essential
   librdkafka-dev

3. Add the following packages to your requirements.txt file:

   confluent-kafka==2.1.1
   apache-airflow-providers-apache-kafka==1.0.0

4. Run the following command to start your project in a local environment:

   1
   astro dev start

Step 2: Create two Kafka connections

The Kafka Airflow provider uses a Kafka connection assigned to the kafka_conn_id parameter of each operator to interact with a Kafka cluster. For this tutorial you define two Kafka connections, because two different consumers will be created.

1. In your web browser, go to localhost:8080 to access the Airflow UI.

2. Click Admin > Connections > + to create a new connection.

3. Name your connection kafka_default and select the Apache Kafka connection type. Provide the details for the connection to your Kafka cluster as JSON in the Extra field.

   If you connect to a local Kafka cluster created with the server.properties in the info box from the Prerequisites section, use the following configuration:

   1
   {
   2
       "bootstrap.servers": "kafka:19092",
   3
       "group.id": "group_1",
   4
       "security.protocol": "PLAINTEXT",
   5
       "auto.offset.reset": "beginning"
   6
   }

   The key-value pairs for your connection depend on what kind of Kafka cluster you are connecting to. Most operators in the Kafka Airflow provider mandate that you define the bootstrap.servers key. You can find a full list of optional connection parameters in the librdkafka documentation.

4. Click Save.

5. Create a second new connection.

6. Name your second connection kafka_listener and select the Apache Kafka connection type. Provide the same details as you did in Step 2, but set the group.id to group_2. You must have a second connection with a different group.id because the DAGs in this tutorial have two consuming tasks that consume messages from the same Kafka topic. Learn more in Kafka’s Consumer Configs documentation.

7. Click Save.

Step 3: Create a DAG with a producer and a consumer task

The Kafka Airflow provider package contains a ProduceToTopicOperator, which you can use to produce messages directly to a Kafka topic, and a ConsumeFromTopicOperator, which you can use to directly consume messages from a topic.

1. Create a new file in your dags folder called produce_consume_treats.py.

2. Copy and paste the following code into the produce_consume_treats.py file:

   1
   """
   2
   ### DAG which produces to and consumes from a Kafka cluster
   3
   4
   This DAG will produce messages consisting of several elements to a Kafka cluster and consume
   5
   them.
   6
   """
   7
   8
   from airflow.decorators import dag, task
   9
   from pendulum import datetime
   10
   from airflow.providers.apache.kafka.operators.produce import ProduceToTopicOperator
   11
   from airflow.providers.apache.kafka.operators.consume import ConsumeFromTopicOperator
   12
   import json
   13
   import random
   14
   15
   YOUR_NAME = "<your name>"
   16
   YOUR_PET_NAME = "<your (imaginary) pet name>"
   17
   NUMBER_OF_TREATS = 5
   18
   KAFKA_TOPIC = "my_topic"
   19
   20
   21
   def prod_function(num_treats, pet_name):
   22
       """Produces `num_treats` messages containing the pet's name, a randomly picked
   23
       pet mood post treat and whether or not it was the last treat in a series."""
   24
   25
       for i in range(num_treats):
   26
           final_treat = False
   27
           pet_mood_post_treat = random.choices(
   28
               ["content", "happy", "zoomy", "bouncy"], weights=[2, 2, 1, 1], k=1
   29
           )[0]
   30
           if i + 1 == num_treats:
   31
               final_treat = True
   32
           yield (
   33
               json.dumps(i),
   34
               json.dumps(
   35
                   {
   36
                       "pet_name": pet_name,
   37
                       "pet_mood_post_treat": pet_mood_post_treat,
   38
                       "final_treat": final_treat,
   39
                   }
   40
               ),
   41
           )
   42
   43
   44
   def consume_function(message, name):
   45
       "Takes in consumed messages and prints its contents to the logs."
   46
   47
       key = json.loads(message.key())
   48
       message_content = json.loads(message.value())
   49
       pet_name = message_content["pet_name"]
   50
       pet_mood_post_treat = message_content["pet_mood_post_treat"]
   51
       print(
   52
           f"Message #{key}: Hello {name}, your pet {pet_name} has consumed another treat and is now {pet_mood_post_treat}!"
   53
       )
   54
   55
   56
   @dag(
   57
       start_date=datetime(2023, 4, 1),
   58
       schedule=None,
   59
       catchup=False,
   60
       render_template_as_native_obj=True,
   61
   )
   62
   def produce_consume_treats():
   63
       @task
   64
       def get_your_pet_name(pet_name=None):
   65
           return pet_name
   66
   67
       @task
   68
       def get_number_of_treats(num_treats=None):
   69
           return num_treats
   70
   71
       @task
   72
       def get_pet_owner_name(your_name=None):
   73
           return your_name
   74
   75
       produce_treats = ProduceToTopicOperator(
   76
           task_id="produce_treats",
   77
           kafka_config_id="kafka_default",
   78
           topic=KAFKA_TOPIC,
   79
           producer_function=prod_function,
   80
           producer_function_args=["{{ ti.xcom_pull(task_ids='get_number_of_treats')}}"],
   81
           producer_function_kwargs={
   82
               "pet_name": "{{ ti.xcom_pull(task_ids='get_your_pet_name')}}"
   83
           },
   84
           poll_timeout=10,
   85
       )
   86
   87
       consume_treats = ConsumeFromTopicOperator(
   88
           task_id="consume_treats",
   89
           kafka_config_id="kafka_default",
   90
           topics=[KAFKA_TOPIC],
   91
           apply_function=consume_function,
   92
           apply_function_kwargs={
   93
               "name": "{{ ti.xcom_pull(task_ids='get_pet_owner_name')}}"
   94
           },
   95
           poll_timeout=20,
   96
           max_messages=20,
   97
           max_batch_size=20,
   98
       )
   99
   100
       [
   101
           get_your_pet_name(YOUR_PET_NAME),
   102
           get_number_of_treats(NUMBER_OF_TREATS),
   103
       ] >> produce_treats
   104
       get_pet_owner_name(YOUR_NAME) >> consume_treats
   105
   106
       produce_treats >> consume_treats
   107
   108
   109
   produce_consume_treats()

   This DAG produces messages to a Kafka topic (KAFKA_TOPIC) and consumes them.

   • The produce_treats task retrieves the number of treats (num_treats) to give to your pet from the upstream get_number_of_treats task. Then, the task supplies the number of treats to the producer_function as a positional argument with the producer_function_args parameter. In a similar process, the task also retrieves the name of your pet from the upstream get_your_pet_name task and provides it as a kwarg to producer_function_kwargs.
   • Next, the produce_treats task writes one message for every treat to a Kafka topic. Each message contains the pet name, a randomly picked pet mood after the treat has been given, and whether or not a treat was the last one in a series. The ProduceToTopicOperator accomplishes this by using a function passed to its producer_function parameter, which returns a generator containing key-value pairs.
   • The consume_treats task consumes messages from the same Kafka topic and modifies them to print a string to the logs using the callable provided to the apply_function parameter. This task also retrieves a value from an upstream task and supplies it as a kwarg to the apply_function with the apply_function_kwargs parameter.

3. Navigate to the Airflow UI (localhost:8080 if you are running Airflow locally) and manually run your DAG.

4. View the produced events in your Kafka cluster. The following example screenshot shows four messages that have been produced to a topic called test_topic_1 in Confluent Cloud.

   

5. View the logs of your consume_treats task, which shows a list of the consumed events.

   

Info

If you defined a schema for your Kafka topic, the generator needs to return compatible objects. In this example, the generator produces a JSON value.

Tip

The ConsumeFromTopicOperator can replace classical sinks by containing the logic to write messages to a storage destination in its apply_function. This gives you the advantage of being able to use Airflow to schedule message consumption from a Kafka topic based on complex logic embedded in your wider data ecosystem. For example, you can write messages to S3 using the S3CreateObjectOperator, which depends on other upstream task having completed successfully, such as the creation of a specific S3 bucket.

Step 4: Create a listener DAG

Airflow can run a function when a specific message appears in your Kafka topic. The AwaitMessageTriggerFunctionSensor is a deferrable operator that listens to your Kafka topic for a message that fulfills specific criteria, which, when met, runs the callable provided to event_triggered_function. The TriggerDagRunOperator can be used within the event_triggered_function to initiate a run of a downstream DAG.

1. Create a new file in your dags folder called listen_to_the_stream.py.

2. Copy and paste the following code into the file:

   1
   """
   2
   ### DAG continuously listening to a Kafka topic for a specific message
   3
   4
   This DAG will always run and asynchronously monitor a Kafka topic for a message
   5
   which causes the funtion supplied to the `apply_function` parameter to return a value.
   6
   If a value is returned by the `apply_function`, the `event_triggered_function` is
   7
   executed. Afterwards the task will go into a deferred state again.
   8
   """
   9
   10
   from airflow.decorators import dag
   11
   from pendulum import datetime
   12
   from airflow.providers.apache.kafka.sensors.kafka import (
   13
       AwaitMessageTriggerFunctionSensor,
   14
   )
   15
   from airflow.operators.trigger_dagrun import TriggerDagRunOperator
   16
   import json
   17
   import uuid
   18
   19
   PET_MOODS_NEEDING_A_WALK = ["zoomy", "bouncy"]
   20
   KAFKA_TOPIC = "my_topic"
   21
   22
   23
   def listen_function(message, pet_moods_needing_a_walk=[]):
   24
       """Checks if the message received indicates a pet is in
   25
       a mood listed in `pet_moods_needing_a_walk` when they received the last
   26
       treat of a treat-series."""
   27
   28
       message_content = json.loads(message.value())
   29
       print(f"Full message: {message_content}")
   30
       pet_name = message_content["pet_name"]
   31
       pet_mood_post_treat = message_content["pet_mood_post_treat"]
   32
       final_treat = message_content["final_treat"]
   33
       if final_treat:
   34
           if pet_mood_post_treat in pet_moods_needing_a_walk:
   35
               return pet_name, pet_mood_post_treat
   36
   37
   38
   def event_triggered_function(event, **context):
   39
       "Kicks off a downstream DAG with conf and waits for its completion."
   40
   41
       pet_name = event[0]
   42
       pet_mood_post_treat = event[1]
   43
       print(
   44
           f"Due to {pet_name} being in a {pet_mood_post_treat} mood, a walk is being initiated..."
   45
       )
   46
       # use the TriggerDagRunOperator (TDRO) to kick off a downstream DAG
   47
       TriggerDagRunOperator(
   48
           trigger_dag_id="walking_my_pet",
   49
           task_id=f"triggered_downstream_dag_{uuid.uuid4()}",
   50
           wait_for_completion=True,  # wait for downstream DAG completion
   51
           conf={"pet_name": pet_name},
   52
           poke_interval=5,
   53
       ).execute(context)
   54
   55
       print(f"The walk has concluded and {pet_name} is now happily taking a nap!")
   56
   57
   58
   @dag(
   59
       start_date=datetime(2023, 4, 1),
   60
       schedule="@continuous",
   61
       max_active_runs=1,
   62
       catchup=False,
   63
       render_template_as_native_obj=True,
   64
   )
   65
   def listen_to_the_stream():
   66
       listen_for_mood = AwaitMessageTriggerFunctionSensor(
   67
           task_id="listen_for_mood",
   68
           kafka_config_id="kafka_listener",
   69
           topics=[KAFKA_TOPIC],
   70
           # the apply function will be used from within the triggerer, this is
   71
           # why it needs to be a dot notation string
   72
           apply_function="listen_to_the_stream.listen_function",
   73
           poll_interval=5,
   74
           poll_timeout=1,
   75
           apply_function_kwargs={"pet_moods_needing_a_walk": PET_MOODS_NEEDING_A_WALK},
   76
           event_triggered_function=event_triggered_function,
   77
       )
   78
   79
   80
   listen_to_the_stream()

   This DAG has one task called listen_for_mood which uses the AwaitMessageTriggerFunctionSensor to listen to messages in all topics supplied to its topics parameters. For each message that is consumed, the following actions are performed:

   • The listen_function supplied to the apply_function parameter of the AwaitMessageTriggerFunctionSensor consumes and processes the message. The listen_function is provided as a dot notation string, which is necessary because the Airflow triggerer component needs to access this function.
   • If the message consumed causes the listen_function to return a value, a TriggerEvent fires.
   • After a TriggerEvent fires, the AwaitMessageTriggerFunctionSensor executes the function provided to the event_triggered_function parameter. In this example, the event_triggered_function starts a downstream DAG using the .execute() method of the TriggerDagRunOperator.
   • After the event_triggered_function completes, the AwaitMessageTriggerFunctionSensor returns to a deferred state.

   The AwaitMessageTriggerFunctionSensor always runs and listens. If the task fails, like if a malformed message is consumed, the DAG completes as failed and automatically starts its next DAG run because of the @continuous schedule.

Info

When working locally, you need to restart your Airflow instance to apply changes to the apply_function of the AwaitMessageTriggerFunctionSensor because the function is imported into the Triggerer, which does not periodically restart. To restart Airflow, run astro dev restart in your terminal. Changes to the event_triggered_function of the AwaitMessageTriggerFunctionSensor do not require a restart of your Airflow instance.

On Astro, the Triggerer is restarted automatically when a new image is deployed, but not on dag-only deploys, see Deploy DAGs to Astro.

Step 5: Create a downstream DAG

The event_triggered_function of the AwaitMessageTriggerFunctionSensor operator starts a downstream DAG. This example shows how to implement a dependency based on messages that appear in your Kafka topic.

1. Create a new file in your dags folder called walking_my_pet.py.

2. Copy and paste the following code into the file:

   1
   """
   2
   ### Simple DAG that runs one task with params
   3
   4
   This DAG uses one string type param and uses it in a python decorated task.
   5
   """
   6
   7
   from airflow.decorators import dag, task
   8
   from pendulum import datetime
   9
   from airflow.models.param import Param
   10
   import random
   11
   12
   13
   @dag(
   14
       start_date=datetime(2023, 4, 1),
   15
       schedule=None,
   16
       catchup=False,
   17
       render_template_as_native_obj=True,
   18
       params={"pet_name": Param("Undefined!", type="string")},
   19
   )
   20
   def walking_my_pet():
   21
       @task
   22
       def walking_your_pet(**context):
   23
           pet_name = context["params"]["pet_name"]
   24
           minutes = random.randint(2, 10)
   25
           print(f"{pet_name} has been on a {minutes} minute walk!")
   26
   27
       walking_your_pet()
   28
   29
   30
   walking_my_pet()

   This DAG acts as a downstream dependency to the listen_to_the_stream DAG. You can add any tasks to this DAG.

Step 6: Run the DAGs

Now that all three DAGs are ready, run them to see how they work together.

1. Make sure you unpause all DAGs in the Airflow UI and that your Kafka cluster is running.

2. The listen_to_the_stream DAG immediately starts running after it unpauses and the listen_for_mood task goes into a Deferred state, which is indicated with a purple square in the Airflow UI.

   

3. Manually run the produce_consume_treats DAG to give your pet some treats and produce a few messages to the Kafka cluster.

4. Check the logs of the listen_for_mood task in the listen_to_the_stream DAG to see if a message fitting the criteria defined by the listen_function has been detected. You might need to run the produce_consume_treats DAG a couple of times for a message to appear.

   If the TriggerEvent of the listen_for_mood task fires, the listen_for_mood task logs show the walking_my_pet DAG initiating.

   

5. Finally, check the logs of the walking_my_pet task to see how long your pet enjoyed their walk!

Best practices

Apache Kafka is a tool optimized for streaming messages at high frequencies, for example in an IoT application. Airflow is designed to handle orchestration of data pipelines in batches.

Astronomer recommends to combine these two open source tools by handling low-latency processes with Kafka and data orchestration with Airflow.

Common patterns include:

• Configuring a Kafka cluster with a blob storage like S3 as a sink. Batch process data from S3 at regular intervals.
• Using the ProduceToTopicOperator in Airflow to produce messages to a Kafka cluster as one of several producers.
• Consuming data from a Kafka cluster via the ConsumeFromTopicOperator in batches using the apply function to extract and load information to a blob storage or data warehouse.
• Listening for specific messages in a data stream running through a Kafka cluster using the AwaitMessageTriggerFunctionSensor to trigger downstream tasks after the message appears.

Conclusion

Congratulations! You used the Kafka Airflow provider to directly interact with a Kafka topic from within Apache Airflow.