You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@airflow.apache.org by "dylanbstorey (via GitHub)" <gi...@apache.org> on 2023/04/07 06:58:28 UTC
[GitHub] [airflow] dylanbstorey opened a new pull request, #30175: Add provider for Apache Kafka
dylanbstorey opened a new pull request, #30175:
URL: https://github.com/apache/airflow/pull/30175
<!--
Thank you for contributing! Please make sure that your code changes
are covered with tests. And in case of new features or big changes
remember to adjust the documentation.
Feel free to ping committers for the review!
In case of an existing issue, reference it using one of the following:
related: #21767, #12388, #10660, #7407, #1415
How to write a good git commit message:
http://chris.beams.io/posts/git-commit/
-->
---
Pulls in a series of integrations to Kafka from `airflow-provider-kafka` (https://pypi.org/project/airflow-provider-kafka/) to core airflow.
I am the maintainer of that package and have permission to submit to OSS Airflow.
**^ Add meaningful description above**
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [airflow] dylanbstorey commented on a diff in pull request #30175: Add provider for Apache Kafka
Posted by "dylanbstorey (via GitHub)" <gi...@apache.org>.
dylanbstorey commented on code in PR #30175:
URL: https://github.com/apache/airflow/pull/30175#discussion_r1167626661
##########
airflow/providers/apache/kafka/operators/consume.py:
##########
@@ -0,0 +1,184 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+from functools import partial
+from typing import Any, Callable, Sequence
+
+from airflow.exceptions import AirflowException
+from airflow.models import BaseOperator
+from airflow.providers.apache.kafka.hooks.consume import KafkaConsumerHook
+from airflow.utils.module_loading import import_string
+
+VALID_COMMIT_CADENCE = {"never", "end_of_batch", "end_of_operator"}
+
+
+class ConsumeFromTopicOperator(BaseOperator):
+ """An operator that consumes from Kafka a topic(s) and processing the messages.
+
+ The operator creates a Kafka consumer that reads a batch of messages from the cluster and processes them
+ using the user supplied callable function. The consumer will continue to read in batches until it reaches
+ the end of the log or reads a maximum number of messages is reached.
+
+ :param kafka_config_id: The connection object to use, defaults to "kafka_default"
+ :param topics: A list of topics or regex patterns the consumer should subscribe to.
+ :param apply_function: The function that should be applied to fetched one at a time.
+ name of dag file executing the function and the function name delimited by a `.`
+ :param apply_function_batch: The function that should be applied to a batch of messages fetched. Can not
+ be used with `apply_function`. Intended for transactional workloads where an expensive task might
+ be called before or after operations on the messages are taken.
+ :param apply_function_args: Additional arguments that should be applied to the callable, defaults to None
+ :param apply_function_kwargs: Additional key word arguments that should be applied to the callable
+ defaults to None
+ :param commit_cadence: When consumers should commit offsets ("never", "end_of_batch","end_of_operator"),
+ defaults to "end_of_operator";
+ if end_of_operator, the commit() is called based on the max_messages arg. Commits are made after the
+ operator has processed the apply_function method for the maximum messages in the operator.
+ if end_of_batch, the commit() is called based on the max_batch_size arg. Commits are made after each
+ batch has processed by the apply_function method for all messages in the batch.
+ if never, close() is called without calling the commit() method.
+ :param max_messages: The maximum total number of messages an operator should read from Kafka,
+ defaults to None
+ :param max_batch_size: The maximum number of messages a consumer should read when polling,
+ defaults to 1000
+ :param poll_timeout: How long the Kafka consumer should wait before determining no more messages are
+ available, defaults to 60
+
+ .. seealso::
+ For more information on how to use this operator, take a look at the guide:
+ :ref:`howto/operator:ConsumeFromTopicOperator`
+ """
+
+ BLUE = "#ffefeb"
+ ui_color = BLUE
+ template_fields = (
+ "topics",
+ "apply_function",
+ "apply_function_args",
+ "apply_function_kwargs",
+ "kafka_config_id",
+ )
+
+ def __init__(
+ self,
+ topics: str | Sequence[str],
+ kafka_config_id: str = "kafka_default",
+ apply_function: Callable[..., Any] | str | None = None,
+ apply_function_batch: Callable[..., Any] | str | None = None,
+ apply_function_args: Sequence[Any] | None = None,
+ apply_function_kwargs: dict[Any, Any] | None = None,
+ commit_cadence: str | None = "end_of_operator",
+ max_messages: int | None = None,
+ max_batch_size: int = 1000,
+ poll_timeout: float = 60,
+ **kwargs: Any,
+ ) -> None:
+
+ super().__init__(**kwargs)
+
+ self.topics = topics
+ self.apply_function = apply_function
+ self.apply_function_batch = apply_function_batch
+ self.apply_function_args = apply_function_args or ()
+ self.apply_function_kwargs = apply_function_kwargs or {}
+ self.kafka_config_id = kafka_config_id
+ self.commit_cadence = commit_cadence
+ self.max_messages = max_messages or True
+ self.max_batch_size = max_batch_size
+ self.poll_timeout = poll_timeout
+
+ if self.commit_cadence not in VALID_COMMIT_CADENCE:
+ raise AirflowException(
+ f"commit_cadence must be one of {VALID_COMMIT_CADENCE}. Got {self.commit_cadence}"
+ )
+
+ if self.max_messages and self.max_batch_size > self.max_messages:
+ self.log.warning(
+ "max_batch_size (%s) > max_messages (%s). Setting max_messages to %s ",
+ self.max_batch_size,
+ self.max_messages,
+ self.max_batch_size,
+ )
+
+ if self.commit_cadence == "never":
+ self.commit_cadence = None
+
+ if apply_function and apply_function_batch:
+ raise AirflowException(
+ "One of apply_function or apply_function_batch must be supplied, not both."
+ )
+
+ def execute(self, context) -> Any:
+
+ consumer = KafkaConsumerHook(topics=self.topics, kafka_config_id=self.kafka_config_id).get_consumer()
+
+ if isinstance(self.apply_function, str):
+ self.apply_function = import_string(self.apply_function)
+
+ if isinstance(self.apply_function_batch, str):
+ self.apply_function_batch = import_string(self.apply_function_batch)
+
+ if self.apply_function:
+ apply_callable = partial(
+ self.apply_function, *self.apply_function_args, **self.apply_function_kwargs # type: ignore
+ )
+
+ if self.apply_function_batch:
+ apply_callable = partial(
+ self.apply_function_batch, # type: ignore
+ *self.apply_function_args,
+ **self.apply_function_kwargs,
+ )
+
+ messages_left = self.max_messages
+ messages_processed = 0
+
+ while (
+ messages_left > 0
+ ): # bool(True > 0) == True in the case where self.max_messages isn't set by the user
+
+ if not isinstance(messages_left, bool):
+ batch_size = self.max_batch_size if messages_left > self.max_batch_size else messages_left
+ else:
+ batch_size = self.max_batch_size
+
+ msgs = consumer.consume(num_messages=batch_size, timeout=self.poll_timeout)
+ messages_left -= len(msgs)
Review Comment:
https://github.com/dylanbstorey/airflow/blob/92d93e9be7371efe4203005cd30981de1a2243d8/airflow/providers/apache/kafka/operators/consume.py#L152
Adding some flavor to the doc string.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [airflow] dylanbstorey commented on a diff in pull request #30175: Add provider for Apache Kafka
Posted by "dylanbstorey (via GitHub)" <gi...@apache.org>.
dylanbstorey commented on code in PR #30175:
URL: https://github.com/apache/airflow/pull/30175#discussion_r1167626974
##########
airflow/providers/apache/kafka/sensors/kafka.py:
##########
Review Comment:
I've had these patterns around since ~1.10.2 and chose not to use them even in limited settings until deferrable existed because they were catastrophically bad. Kafka and Airflow already have plenty of footguns available - I'd prefer not to hand out more.
We can revisit this conversation as KIP-500 reaches more widespread adoption and I get some time to test some how Airflow 2.0 does with this pattern though.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [airflow] dylanbstorey commented on a diff in pull request #30175: Add provider for Apache Kafka
Posted by "dylanbstorey (via GitHub)" <gi...@apache.org>.
dylanbstorey commented on code in PR #30175:
URL: https://github.com/apache/airflow/pull/30175#discussion_r1167628969
##########
docs/apache-airflow-providers-apache-kafka/connections/kafka.rst:
##########
@@ -0,0 +1,38 @@
+ .. Licensed to the Apache Software Foundation (ASF) under one
+ or more contributor license agreements. See the NOTICE file
+ distributed with this work for additional information
+ regarding copyright ownership. The ASF licenses this file
+ to you under the Apache License, Version 2.0 (the
+ "License"); you may not use this file except in compliance
+ with the License. You may obtain a copy of the License at
+
+ .. http://www.apache.org/licenses/LICENSE-2.0
+
+ .. Unless required by applicable law or agreed to in writing,
+ software distributed under the License is distributed on an
+ "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ KIND, either express or implied. See the License for the
+ specific language governing permissions and limitations
+ under the License.
+
+.. _howto/connection: kafka
+
+Apache Kafka Connection
+========================
+
+The Apache Kafka connection type configures a connection to Apache Kafka via the ``confluent-kafka`` Python package.
+
+Default Connection IDs
+----------------------
+
+Kafka hooks and operators use ``kafka_default`` by default.
Review Comment:
Just an oversite - updated.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [airflow] dylanbstorey commented on a diff in pull request #30175: Add provider for Apache Kafka
Posted by "dylanbstorey (via GitHub)" <gi...@apache.org>.
dylanbstorey commented on code in PR #30175:
URL: https://github.com/apache/airflow/pull/30175#discussion_r1167627281
##########
airflow/providers/apache/kafka/hooks/base.py:
##########
@@ -0,0 +1,71 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+from typing import Any
+
+from confluent_kafka.admin import AdminClient
+
+from airflow.compat.functools import cached_property
+from airflow.hooks.base import BaseHook
+
+
+class KafkaHook(BaseHook):
Review Comment:
🤷 - no strong feelings.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [airflow] hussein-awala commented on a diff in pull request #30175: Add provider for Apache Kafka
Posted by "hussein-awala (via GitHub)" <gi...@apache.org>.
hussein-awala commented on code in PR #30175:
URL: https://github.com/apache/airflow/pull/30175#discussion_r1167634245
##########
airflow/providers/apache/kafka/operators/consume.py:
##########
@@ -0,0 +1,184 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+from functools import partial
+from typing import Any, Callable, Sequence
+
+from airflow.exceptions import AirflowException
+from airflow.models import BaseOperator
+from airflow.providers.apache.kafka.hooks.consume import KafkaConsumerHook
+from airflow.utils.module_loading import import_string
+
+VALID_COMMIT_CADENCE = {"never", "end_of_batch", "end_of_operator"}
+
+
+class ConsumeFromTopicOperator(BaseOperator):
+ """An operator that consumes from Kafka a topic(s) and processing the messages.
+
+ The operator creates a Kafka consumer that reads a batch of messages from the cluster and processes them
+ using the user supplied callable function. The consumer will continue to read in batches until it reaches
+ the end of the log or reads a maximum number of messages is reached.
+
+ :param kafka_config_id: The connection object to use, defaults to "kafka_default"
+ :param topics: A list of topics or regex patterns the consumer should subscribe to.
+ :param apply_function: The function that should be applied to fetched one at a time.
+ name of dag file executing the function and the function name delimited by a `.`
+ :param apply_function_batch: The function that should be applied to a batch of messages fetched. Can not
+ be used with `apply_function`. Intended for transactional workloads where an expensive task might
+ be called before or after operations on the messages are taken.
+ :param apply_function_args: Additional arguments that should be applied to the callable, defaults to None
+ :param apply_function_kwargs: Additional key word arguments that should be applied to the callable
+ defaults to None
+ :param commit_cadence: When consumers should commit offsets ("never", "end_of_batch","end_of_operator"),
+ defaults to "end_of_operator";
+ if end_of_operator, the commit() is called based on the max_messages arg. Commits are made after the
+ operator has processed the apply_function method for the maximum messages in the operator.
+ if end_of_batch, the commit() is called based on the max_batch_size arg. Commits are made after each
+ batch has processed by the apply_function method for all messages in the batch.
+ if never, close() is called without calling the commit() method.
+ :param max_messages: The maximum total number of messages an operator should read from Kafka,
+ defaults to None
+ :param max_batch_size: The maximum number of messages a consumer should read when polling,
+ defaults to 1000
+ :param poll_timeout: How long the Kafka consumer should wait before determining no more messages are
+ available, defaults to 60
+
+ .. seealso::
+ For more information on how to use this operator, take a look at the guide:
+ :ref:`howto/operator:ConsumeFromTopicOperator`
+ """
+
+ BLUE = "#ffefeb"
+ ui_color = BLUE
+ template_fields = (
+ "topics",
+ "apply_function",
+ "apply_function_args",
+ "apply_function_kwargs",
+ "kafka_config_id",
+ )
+
+ def __init__(
+ self,
+ topics: str | Sequence[str],
+ kafka_config_id: str = "kafka_default",
+ apply_function: Callable[..., Any] | str | None = None,
+ apply_function_batch: Callable[..., Any] | str | None = None,
+ apply_function_args: Sequence[Any] | None = None,
+ apply_function_kwargs: dict[Any, Any] | None = None,
+ commit_cadence: str | None = "end_of_operator",
+ max_messages: int | None = None,
+ max_batch_size: int = 1000,
+ poll_timeout: float = 60,
+ **kwargs: Any,
+ ) -> None:
+
+ super().__init__(**kwargs)
+
+ self.topics = topics
+ self.apply_function = apply_function
+ self.apply_function_batch = apply_function_batch
+ self.apply_function_args = apply_function_args or ()
+ self.apply_function_kwargs = apply_function_kwargs or {}
+ self.kafka_config_id = kafka_config_id
+ self.commit_cadence = commit_cadence
+ self.max_messages = max_messages or True
Review Comment:
I understand what happen with the current code, my question was more about what is supposed to happen.
> Functionally this means "read to the end of the log and exit when kafka has no more messages".
Can you please add it to the doc to avoid the confusion?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [airflow] dylanbstorey commented on a diff in pull request #30175: Add provider for Apache Kafka
Posted by "dylanbstorey (via GitHub)" <gi...@apache.org>.
dylanbstorey commented on code in PR #30175:
URL: https://github.com/apache/airflow/pull/30175#discussion_r1163359179
##########
airflow/providers/apache/kafka/hooks/consume.py:
##########
@@ -0,0 +1,50 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+from typing import Sequence
+
+from confluent_kafka import Consumer
+
+from airflow.providers.apache.kafka.hooks.base import KafkaHook
+
+
+class KafkaConsumerHook(KafkaHook):
+ """
+ A hook for creating a Kafka Consumer
+
+ :param kafka_config_id: The connection object to use, defaults to "kafka_config_default"
+ :param topics: A list of topics to subscribe to.
+ """
+
+ def __init__(self, topics: Sequence[str], kafka_config_id=KafkaHook.default_conn_name) -> None:
+
+ super().__init__(kafka_config_id=kafka_config_id)
+ self.topics = topics
+ if not self.get_conn().get("group.id", None):
+ raise ValueError("The 'group.id' parameter must be set in the config dictionary'. Got <None>")
+
+ if not self.get_conn().get("bootstrap.servers", None):
Review Comment:
Good call outs.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [airflow] dylanbstorey commented on a diff in pull request #30175: Add provider for Apache Kafka
Posted by "dylanbstorey (via GitHub)" <gi...@apache.org>.
dylanbstorey commented on code in PR #30175:
URL: https://github.com/apache/airflow/pull/30175#discussion_r1167630051
##########
airflow/providers/apache/kafka/hooks/client.py:
##########
@@ -0,0 +1,65 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+from typing import Any, Sequence
+
+from confluent_kafka.admin import AdminClient, NewTopic
+
+from airflow.providers.apache.kafka.hooks.base import KafkaHook
+
+
+class KafkaAdminClientHook(KafkaHook):
+ """
+ A hook for interacting with the Kafka Cluster
+
+ :param kafka_config_id: The connection object to use, defaults to "kafka_default"
+ """
+
+ def __init__(self, kafka_config_id=KafkaHook.default_conn_name) -> None:
+ super().__init__(kafka_config_id=kafka_config_id)
+
+ def get_admin_client(self) -> AdminClient:
+ """returns an AdminClient for communicating with the cluster
+
+ :return: an interactive admin client for the Kafka cluster
+ :rtype: AdminClient
+ """
+ client = AdminClient(self.get_conn)
+ return client
+
+ def create_topic(
+ self,
+ topics: Sequence[Sequence[Any]],
Review Comment:
>the topics list format and arguments order is not clear in the docstring.
Added flavor.
>IMHO we can create a dataclass in the provider with the 3 used fields, and the user can import it to provide the list, also we can support a list of NewTopic as we do in the KPO, in this case the can also use the kafka NewTopic.
Valid strategies all. I don't really have strong feelings about it beyond we can delay the decision to discuss more.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [airflow] potiuk commented on pull request #30175: Add provider for Apache Kafka
Posted by "potiuk (via GitHub)" <gi...@apache.org>.
potiuk commented on PR #30175:
URL: https://github.com/apache/airflow/pull/30175#issuecomment-1508915771
You will need to resolve conflicts (again I am afraid)
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [airflow] dylanbstorey commented on a diff in pull request #30175: Add provider for Apache Kafka
Posted by "dylanbstorey (via GitHub)" <gi...@apache.org>.
dylanbstorey commented on code in PR #30175:
URL: https://github.com/apache/airflow/pull/30175#discussion_r1167626112
##########
airflow/providers/apache/kafka/hooks/consume.py:
##########
Review Comment:
The consumer is cleaned up on exit automatically.
##########
airflow/providers/apache/kafka/hooks/produce.py:
##########
Review Comment:
The consumer is cleaned up on exit automatically.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [airflow] dylanbstorey commented on pull request #30175: Add provider for Apache Kafka
Posted by "dylanbstorey (via GitHub)" <gi...@apache.org>.
dylanbstorey commented on PR #30175:
URL: https://github.com/apache/airflow/pull/30175#issuecomment-1474601653
> How does a user define an Airflow Connection to use these hooks?
The config dict is the primary mechanism for describing the connection.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [airflow] dylanbstorey commented on a diff in pull request #30175: Add provider for Apache Kafka
Posted by "dylanbstorey (via GitHub)" <gi...@apache.org>.
dylanbstorey commented on code in PR #30175:
URL: https://github.com/apache/airflow/pull/30175#discussion_r1167627397
##########
airflow/providers/apache/kafka/operators/consume.py:
##########
@@ -0,0 +1,184 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+from functools import partial
+from typing import Any, Callable, Sequence
+
+from airflow.exceptions import AirflowException
+from airflow.models import BaseOperator
+from airflow.providers.apache.kafka.hooks.consume import KafkaConsumerHook
+from airflow.utils.module_loading import import_string
+
+VALID_COMMIT_CADENCE = {"never", "end_of_batch", "end_of_operator"}
+
+
+class ConsumeFromTopicOperator(BaseOperator):
+ """An operator that consumes from Kafka a topic(s) and processing the messages.
+
+ The operator creates a Kafka consumer that reads a batch of messages from the cluster and processes them
+ using the user supplied callable function. The consumer will continue to read in batches until it reaches
+ the end of the log or reads a maximum number of messages is reached.
+
+ :param kafka_config_id: The connection object to use, defaults to "kafka_default"
+ :param topics: A list of topics or regex patterns the consumer should subscribe to.
+ :param apply_function: The function that should be applied to fetched one at a time.
+ name of dag file executing the function and the function name delimited by a `.`
+ :param apply_function_batch: The function that should be applied to a batch of messages fetched. Can not
+ be used with `apply_function`. Intended for transactional workloads where an expensive task might
+ be called before or after operations on the messages are taken.
+ :param apply_function_args: Additional arguments that should be applied to the callable, defaults to None
+ :param apply_function_kwargs: Additional key word arguments that should be applied to the callable
+ defaults to None
+ :param commit_cadence: When consumers should commit offsets ("never", "end_of_batch","end_of_operator"),
+ defaults to "end_of_operator";
+ if end_of_operator, the commit() is called based on the max_messages arg. Commits are made after the
+ operator has processed the apply_function method for the maximum messages in the operator.
+ if end_of_batch, the commit() is called based on the max_batch_size arg. Commits are made after each
+ batch has processed by the apply_function method for all messages in the batch.
+ if never, close() is called without calling the commit() method.
+ :param max_messages: The maximum total number of messages an operator should read from Kafka,
+ defaults to None
+ :param max_batch_size: The maximum number of messages a consumer should read when polling,
+ defaults to 1000
+ :param poll_timeout: How long the Kafka consumer should wait before determining no more messages are
+ available, defaults to 60
+
+ .. seealso::
+ For more information on how to use this operator, take a look at the guide:
+ :ref:`howto/operator:ConsumeFromTopicOperator`
+ """
+
+ BLUE = "#ffefeb"
+ ui_color = BLUE
+ template_fields = (
+ "topics",
+ "apply_function",
+ "apply_function_args",
+ "apply_function_kwargs",
+ "kafka_config_id",
+ )
+
+ def __init__(
+ self,
+ topics: str | Sequence[str],
+ kafka_config_id: str = "kafka_default",
+ apply_function: Callable[..., Any] | str | None = None,
+ apply_function_batch: Callable[..., Any] | str | None = None,
+ apply_function_args: Sequence[Any] | None = None,
+ apply_function_kwargs: dict[Any, Any] | None = None,
+ commit_cadence: str | None = "end_of_operator",
+ max_messages: int | None = None,
+ max_batch_size: int = 1000,
+ poll_timeout: float = 60,
+ **kwargs: Any,
+ ) -> None:
+
+ super().__init__(**kwargs)
+
+ self.topics = topics
+ self.apply_function = apply_function
+ self.apply_function_batch = apply_function_batch
+ self.apply_function_args = apply_function_args or ()
+ self.apply_function_kwargs = apply_function_kwargs or {}
+ self.kafka_config_id = kafka_config_id
+ self.commit_cadence = commit_cadence
+ self.max_messages = max_messages or True
Review Comment:
```python
>>> y = 0 or True
>>> y
True
```
^that.
Functionally this means "read to the end of the log and exit when kafka has no more messages".
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [airflow] dylanbstorey commented on a diff in pull request #30175: Add provider for Apache Kafka
Posted by "dylanbstorey (via GitHub)" <gi...@apache.org>.
dylanbstorey commented on code in PR #30175:
URL: https://github.com/apache/airflow/pull/30175#discussion_r1167625595
##########
airflow/providers/apache/kafka/hooks/base.py:
##########
@@ -0,0 +1,71 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+from typing import Any
+
+from confluent_kafka.admin import AdminClient
+
+from airflow.compat.functools import cached_property
+from airflow.hooks.base import BaseHook
+
+
+class KafkaHook(BaseHook):
+ """A base hook for interacting with Apache Kafka"""
+
+ conn_name_attr = "kafka_config_id"
+ default_conn_name = "kafka_default"
+ conn_type = "kafka"
+ hook_name = "Apache Kafka"
+
+ def __init__(self, kafka_config_id=default_conn_name, *args, **kwargs):
+ """Initialize our Base"""
+ super().__init__()
+ self.kafka_config_id = kafka_config_id
+ self.get_conn
+
+ @staticmethod
+ def get_ui_field_behaviour() -> dict[str, Any]:
+ """Returns custom field behaviour"""
+ return {
+ "hidden_fields": ["schema", "login", "password", "port", "host"],
+ "relabeling": {"extra": "Config Dict"},
+ "placeholders": {
+ "extra": '{"bootstrap.servers": "localhost:9092"}',
+ },
+ }
+
+ @cached_property
+ def get_conn(self) -> Any:
+ """get the configuration object"""
+ config = self.get_connection(self.kafka_config_id).extra_dejson
+
+ if not (config.get("bootstrap.servers", None)):
+ raise ValueError("config['bootstrap.servers'] must be provided.")
+
+ return config
Review Comment:
I don't think adding an additional layer of misdirection here makes much sense at this time.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [airflow] eladkal closed pull request #30175: Add provider for Apache Kafka
Posted by "eladkal (via GitHub)" <gi...@apache.org>.
eladkal closed pull request #30175: Add provider for Apache Kafka
URL: https://github.com/apache/airflow/pull/30175
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [airflow] dylanbstorey commented on a diff in pull request #30175: Add provider for Apache Kafka
Posted by "dylanbstorey (via GitHub)" <gi...@apache.org>.
dylanbstorey commented on code in PR #30175:
URL: https://github.com/apache/airflow/pull/30175#discussion_r1140902267
##########
airflow/providers/apache/kafka/shared_utils.py:
##########
@@ -0,0 +1,53 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+import importlib
+from typing import Any, Callable
+
+
+def _get_callable(function_string: str) -> Callable:
Review Comment:
there is.
The doc string indicates "attr/class", while true it is pretty much the same code - i'm not sure if the intent will stay consistent in the future.
So - if you feel strongly about it / don't think it will a problem in the future I'm happy to change it
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [airflow] dylanbstorey commented on a diff in pull request #30175: Add provider for Apache Kafka
Posted by "dylanbstorey (via GitHub)" <gi...@apache.org>.
dylanbstorey commented on code in PR #30175:
URL: https://github.com/apache/airflow/pull/30175#discussion_r1140899328
##########
tests/providers/apache/kafka/hooks/test_admin_client.py:
##########
@@ -0,0 +1,69 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+import logging
+import unittest
+
+import pytest
+
+from airflow import AirflowException
+
+# Import Hook
+from airflow.providers.apache.kafka.hooks.admin_client import KafkaAdminClientHook
+
+log = logging.getLogger(__name__)
+
+
+class TestSampleHook(unittest.TestCase):
Review Comment:
i don't know what you mean. All of these tests are pytest compatible.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [airflow] dylanbstorey commented on a diff in pull request #30175: Add provider for Apache Kafka
Posted by "dylanbstorey (via GitHub)" <gi...@apache.org>.
dylanbstorey commented on code in PR #30175:
URL: https://github.com/apache/airflow/pull/30175#discussion_r1141075350
##########
airflow/providers/apache/kafka/CHANGELOG.md:
##########
Review Comment:
> connections.rst / operators.rst
Not sure if you can provide something a bit more concrete here on what you'd like to see differently.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [airflow] hussein-awala commented on a diff in pull request #30175: Add provider for Apache Kafka
Posted by "hussein-awala (via GitHub)" <gi...@apache.org>.
hussein-awala commented on code in PR #30175:
URL: https://github.com/apache/airflow/pull/30175#discussion_r1167632186
##########
airflow/providers/apache/kafka/hooks/base.py:
##########
@@ -0,0 +1,71 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+from typing import Any
+
+from confluent_kafka.admin import AdminClient
+
+from airflow.compat.functools import cached_property
+from airflow.hooks.base import BaseHook
+
+
+class KafkaHook(BaseHook):
+ """A base hook for interacting with Apache Kafka"""
+
+ conn_name_attr = "kafka_config_id"
+ default_conn_name = "kafka_default"
+ conn_type = "kafka"
+ hook_name = "Apache Kafka"
+
+ def __init__(self, kafka_config_id=default_conn_name, *args, **kwargs):
+ """Initialize our Base"""
+ super().__init__()
+ self.kafka_config_id = kafka_config_id
+ self.get_conn
+
+ @staticmethod
+ def get_ui_field_behaviour() -> dict[str, Any]:
+ """Returns custom field behaviour"""
+ return {
+ "hidden_fields": ["schema", "login", "password", "port", "host"],
+ "relabeling": {"extra": "Config Dict"},
+ "placeholders": {
+ "extra": '{"bootstrap.servers": "localhost:9092"}',
+ },
+ }
+
+ @cached_property
+ def get_conn(self) -> Any:
+ """get the configuration object"""
+ config = self.get_connection(self.kafka_config_id).extra_dejson
+
+ if not (config.get("bootstrap.servers", None)):
+ raise ValueError("config['bootstrap.servers'] must be provided.")
+
+ return config
Review Comment:
I don't know if it's documented somewhere or not, but `get_conn()` is the method that actually establishes the connection to the external system. This method should take care of authenticating with the system, and returning a connection object that can be used to execute commands against that system.
I prefer to check with @potiuk and @ashb before marked it as resolved
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [airflow] hussein-awala commented on a diff in pull request #30175: Add provider for Apache Kafka
Posted by "hussein-awala (via GitHub)" <gi...@apache.org>.
hussein-awala commented on code in PR #30175:
URL: https://github.com/apache/airflow/pull/30175#discussion_r1167482795
##########
airflow/providers/apache/kafka/hooks/client.py:
##########
Review Comment:
can you implement the method `__exit__` to auto-close the client?
##########
airflow/providers/apache/kafka/hooks/base.py:
##########
@@ -0,0 +1,71 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+from typing import Any
+
+from confluent_kafka.admin import AdminClient
+
+from airflow.compat.functools import cached_property
+from airflow.hooks.base import BaseHook
+
+
+class KafkaHook(BaseHook):
+ """A base hook for interacting with Apache Kafka"""
+
+ conn_name_attr = "kafka_config_id"
+ default_conn_name = "kafka_default"
+ conn_type = "kafka"
+ hook_name = "Apache Kafka"
+
+ def __init__(self, kafka_config_id=default_conn_name, *args, **kwargs):
+ """Initialize our Base"""
+ super().__init__()
+ self.kafka_config_id = kafka_config_id
+ self.get_conn
+
+ @staticmethod
+ def get_ui_field_behaviour() -> dict[str, Any]:
+ """Returns custom field behaviour"""
+ return {
+ "hidden_fields": ["schema", "login", "password", "port", "host"],
+ "relabeling": {"extra": "Config Dict"},
+ "placeholders": {
+ "extra": '{"bootstrap.servers": "localhost:9092"}',
+ },
+ }
+
+ @cached_property
+ def get_conn(self) -> Any:
+ """get the configuration object"""
+ config = self.get_connection(self.kafka_config_id).extra_dejson
+
+ if not (config.get("bootstrap.servers", None)):
+ raise ValueError("config['bootstrap.servers'] must be provided.")
+
+ return config
Review Comment:
In the `get_conn` method, I think we should return an authenticated client instead of just the configuration object.
If there is a common part in the `get_conn` method between the hooks and each hook requires a different client, you can create a method `_get_client` that creates a kafka client from the config dict, and override it each subclass hook to return the appropriate client.
Here's an example:
```python
def _get_client(config):
...
@cached_property
def get_conn(self) -> Any:
"""get the configuration object"""
config = self.get_connection(self.kafka_config_id).extra_dejson
if not (config.get("bootstrap.servers", None)):
raise ValueError("config['bootstrap.servers'] must be provided.")
return self._get_client(config)
```
in each subclass hook you just need to override the method `_get_client`, and use `get_conn` which returns the appropriate client.
##########
airflow/providers/apache/kafka/hooks/client.py:
##########
@@ -0,0 +1,65 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+from typing import Any, Sequence
+
+from confluent_kafka.admin import AdminClient, NewTopic
+
+from airflow.providers.apache.kafka.hooks.base import KafkaHook
+
+
+class KafkaAdminClientHook(KafkaHook):
+ """
+ A hook for interacting with the Kafka Cluster
+
+ :param kafka_config_id: The connection object to use, defaults to "kafka_default"
+ """
+
+ def __init__(self, kafka_config_id=KafkaHook.default_conn_name) -> None:
+ super().__init__(kafka_config_id=kafka_config_id)
+
+ def get_admin_client(self) -> AdminClient:
+ """returns an AdminClient for communicating with the cluster
+
+ :return: an interactive admin client for the Kafka cluster
+ :rtype: AdminClient
+ """
+ client = AdminClient(self.get_conn)
+ return client
+
+ def create_topic(
+ self,
+ topics: Sequence[Sequence[Any]],
Review Comment:
the `topics` list format and arguments order is not clear in the docstring.
IMHO we can create a dataclass in the provider with the 3 used fields, and the user can import it to provide the list, also we can support a list of `NewTopic` as we do in the KPO, in this case the can also use the kafka `NewTopic`.
##########
airflow/providers/apache/kafka/hooks/consume.py:
##########
Review Comment:
same comment for `__exit__`
##########
airflow/providers/apache/kafka/sensors/kafka.py:
##########
@@ -0,0 +1,219 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+from typing import Any, Callable, Sequence
+
+from airflow.models import BaseOperator
+from airflow.providers.apache.kafka.triggers.await_message import AwaitMessageTrigger
+
+VALID_COMMIT_CADENCE = {"never", "end_of_batch", "end_of_operator"}
+
+
+class AwaitMessageSensor(BaseOperator):
+ """An Airflow operator that defers until a specific message is published to Kafka.
+
+ The operator creates a consumer that reads the Kafka log until it encounters a positive event.
+
+ The behavior of the consumer for this trigger is as follows:
+ - poll the Kafka topics for a message
+ - if no message returned, sleep
+ - process the message with provided callable and commit the message offset
+ - if callable returns any data, raise a TriggerEvent with the return data
+ - else continue to next message
+ - return event (as default xcom or specific xcom key)
+
+
+ :param kafka_config_id: The connection object to use, defaults to "kafka_default"
+ :param topics: Topics (or topic regex) to use for reading from
+ :param apply_function: The function to apply to messages to determine if an event occurred. As a dot
+ notation string.
+ :param apply_function_args: Arguments to be applied to the processing function,
+ defaults to None
+ :param apply_function_kwargs: Key word arguments to be applied to the processing function,
+ defaults to None
+ :param poll_timeout: How long the kafka consumer should wait for a message to arrive from the kafka
+ cluster,defaults to 1
+ :param poll_interval: How long the kafka consumer should sleep after reaching the end of the Kafka log,
+ defaults to 5
+ :param xcom_push_key: the name of a key to push the returned message to, defaults to None
+
+
+ """
+
+ BLUE = "#ffefeb"
+ ui_color = BLUE
+
+ template_fields = (
+ "topics",
+ "apply_function",
+ "apply_function_args",
+ "apply_function_kwargs",
+ "kafka_config_id",
+ )
+
+ def __init__(
+ self,
+ topics: Sequence[str],
+ apply_function: str,
+ kafka_config_id: str = "kafka_default",
+ apply_function_args: Sequence[Any] | None = None,
+ apply_function_kwargs: dict[Any, Any] | None = None,
+ poll_timeout: float = 1,
+ poll_interval: float = 5,
+ xcom_push_key=None,
+ **kwargs: Any,
+ ) -> None:
+
+ super().__init__(**kwargs)
+
+ self.topics = topics
+ self.apply_function = apply_function
+ self.apply_function_args = apply_function_args
+ self.apply_function_kwargs = apply_function_kwargs
+ self.kafka_config_id = kafka_config_id
+ self.poll_timeout = poll_timeout
+ self.poll_interval = poll_interval
+ self.xcom_push_key = xcom_push_key
+
+ def execute(self, context) -> Any:
+
+ self.defer(
+ trigger=AwaitMessageTrigger(
+ topics=self.topics,
+ apply_function=self.apply_function,
+ apply_function_args=self.apply_function_args,
+ apply_function_kwargs=self.apply_function_kwargs,
+ kafka_config_id=self.kafka_config_id,
+ poll_timeout=self.poll_timeout,
+ poll_interval=self.poll_interval,
+ ),
+ method_name="execute_complete",
+ )
+
+ def execute_complete(self, context, event=None):
+ if self.xcom_push_key:
+ self.xcom_push(context, key=self.xcom_push_key, value=event)
+ return event
+
+
+class AwaitMessageTriggerFunctionSensor(BaseOperator):
+ """An Airflow operator that defers until a specific message is published to
+ Kafka, then triggers a registered function
+
+
+ The behavior of the consumer for this trigger is as follows:
+ - poll the Kafka topics for a message
+ - if no message returned, sleep
+ - process the message with provided callable and commit the message offset
+ - if callable returns any data, raise a TriggerEvent with the return data
+ - else continue to next message
+ - return event (as default xcom or specific xcom key)
+
+
+ :param kafka_config_id: The connection object to use, defaults to "kafka_default"
+ :param topics: Topics (or topic regex) to use for reading from
+ :param apply_function: The function to apply to messages to determine if an event occurred. As a dot
+ notation string.
+ :param event_triggered_function: The callable to trigger once the apply_function encounters a
+ positive event.
+ :param apply_function_args: Arguments to be applied to the processing function, defaults to None
+ :param apply_function_kwargs: Key word arguments to be applied to the processing function,
+ defaults to None
+ :param poll_timeout: How long the kafka consumer should wait for a message to arrive from the kafka
+ cluster, defaults to 1
+ :param poll_interval: How long the kafka consumer should sleep after reaching the end of the Kafka log,
+ defaults to 5
+ :param xcom_push_key: the name of a key to push the returned message to, defaults to None
Review Comment:
It is not used in this sensor, missing implementation or extra doc?
##########
airflow/providers/apache/kafka/hooks/client.py:
##########
@@ -0,0 +1,65 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+from typing import Any, Sequence
+
+from confluent_kafka.admin import AdminClient, NewTopic
+
+from airflow.providers.apache.kafka.hooks.base import KafkaHook
+
+
+class KafkaAdminClientHook(KafkaHook):
+ """
+ A hook for interacting with the Kafka Cluster
+
+ :param kafka_config_id: The connection object to use, defaults to "kafka_default"
+ """
+
+ def __init__(self, kafka_config_id=KafkaHook.default_conn_name) -> None:
+ super().__init__(kafka_config_id=kafka_config_id)
+
+ def get_admin_client(self) -> AdminClient:
+ """returns an AdminClient for communicating with the cluster
+
+ :return: an interactive admin client for the Kafka cluster
+ :rtype: AdminClient
+ """
+ client = AdminClient(self.get_conn)
+ return client
+
+ def create_topic(
+ self,
+ topics: Sequence[Sequence[Any]],
+ ) -> None:
+ """creates a topic
+
+ :param topics: a list of topics to create
+ """
+ admin_client = self.get_admin_client()
+
+ new_topics = [NewTopic(t[0], num_partitions=t[1], replication_factor=t[2]) for t in topics]
+
+ futures = admin_client.create_topics(new_topics)
+
+ for t, f in futures.items():
+ try:
+ f.result()
+ self.log.info("The topic %s has been created.", t)
+ except Exception as e:
+ if e.args[0].name() == "TOPIC_ALREADY_EXISTS":
+ self.log.warning("The topic %s already exists.", t)
Review Comment:
```suggestion
except KafkaException as e:
if e.args[0].name == "TOPIC_ALREADY_EXISTS":
self.log.warning("The topic %s already exists.", t)
else:
raise
```
##########
airflow/providers/apache/kafka/operators/consume.py:
##########
@@ -0,0 +1,184 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+from functools import partial
+from typing import Any, Callable, Sequence
+
+from airflow.exceptions import AirflowException
+from airflow.models import BaseOperator
+from airflow.providers.apache.kafka.hooks.consume import KafkaConsumerHook
+from airflow.utils.module_loading import import_string
+
+VALID_COMMIT_CADENCE = {"never", "end_of_batch", "end_of_operator"}
+
+
+class ConsumeFromTopicOperator(BaseOperator):
+ """An operator that consumes from Kafka a topic(s) and processing the messages.
+
+ The operator creates a Kafka consumer that reads a batch of messages from the cluster and processes them
+ using the user supplied callable function. The consumer will continue to read in batches until it reaches
+ the end of the log or reads a maximum number of messages is reached.
+
+ :param kafka_config_id: The connection object to use, defaults to "kafka_default"
+ :param topics: A list of topics or regex patterns the consumer should subscribe to.
+ :param apply_function: The function that should be applied to fetched one at a time.
+ name of dag file executing the function and the function name delimited by a `.`
+ :param apply_function_batch: The function that should be applied to a batch of messages fetched. Can not
+ be used with `apply_function`. Intended for transactional workloads where an expensive task might
+ be called before or after operations on the messages are taken.
+ :param apply_function_args: Additional arguments that should be applied to the callable, defaults to None
+ :param apply_function_kwargs: Additional key word arguments that should be applied to the callable
+ defaults to None
+ :param commit_cadence: When consumers should commit offsets ("never", "end_of_batch","end_of_operator"),
+ defaults to "end_of_operator";
+ if end_of_operator, the commit() is called based on the max_messages arg. Commits are made after the
+ operator has processed the apply_function method for the maximum messages in the operator.
+ if end_of_batch, the commit() is called based on the max_batch_size arg. Commits are made after each
+ batch has processed by the apply_function method for all messages in the batch.
+ if never, close() is called without calling the commit() method.
+ :param max_messages: The maximum total number of messages an operator should read from Kafka,
+ defaults to None
+ :param max_batch_size: The maximum number of messages a consumer should read when polling,
+ defaults to 1000
+ :param poll_timeout: How long the Kafka consumer should wait before determining no more messages are
+ available, defaults to 60
+
+ .. seealso::
+ For more information on how to use this operator, take a look at the guide:
+ :ref:`howto/operator:ConsumeFromTopicOperator`
+ """
+
+ BLUE = "#ffefeb"
+ ui_color = BLUE
+ template_fields = (
+ "topics",
+ "apply_function",
+ "apply_function_args",
+ "apply_function_kwargs",
+ "kafka_config_id",
+ )
+
+ def __init__(
+ self,
+ topics: str | Sequence[str],
+ kafka_config_id: str = "kafka_default",
+ apply_function: Callable[..., Any] | str | None = None,
+ apply_function_batch: Callable[..., Any] | str | None = None,
+ apply_function_args: Sequence[Any] | None = None,
+ apply_function_kwargs: dict[Any, Any] | None = None,
+ commit_cadence: str | None = "end_of_operator",
+ max_messages: int | None = None,
+ max_batch_size: int = 1000,
+ poll_timeout: float = 60,
+ **kwargs: Any,
+ ) -> None:
+
+ super().__init__(**kwargs)
+
+ self.topics = topics
+ self.apply_function = apply_function
+ self.apply_function_batch = apply_function_batch
+ self.apply_function_args = apply_function_args or ()
+ self.apply_function_kwargs = apply_function_kwargs or {}
+ self.kafka_config_id = kafka_config_id
+ self.commit_cadence = commit_cadence
+ self.max_messages = max_messages or True
+ self.max_batch_size = max_batch_size
+ self.poll_timeout = poll_timeout
+
+ if self.commit_cadence not in VALID_COMMIT_CADENCE:
+ raise AirflowException(
+ f"commit_cadence must be one of {VALID_COMMIT_CADENCE}. Got {self.commit_cadence}"
+ )
+
+ if self.max_messages and self.max_batch_size > self.max_messages:
+ self.log.warning(
+ "max_batch_size (%s) > max_messages (%s). Setting max_messages to %s ",
+ self.max_batch_size,
+ self.max_messages,
+ self.max_batch_size,
+ )
+
+ if self.commit_cadence == "never":
+ self.commit_cadence = None
+
+ if apply_function and apply_function_batch:
+ raise AirflowException(
+ "One of apply_function or apply_function_batch must be supplied, not both."
+ )
+
+ def execute(self, context) -> Any:
+
+ consumer = KafkaConsumerHook(topics=self.topics, kafka_config_id=self.kafka_config_id).get_consumer()
+
+ if isinstance(self.apply_function, str):
+ self.apply_function = import_string(self.apply_function)
+
+ if isinstance(self.apply_function_batch, str):
+ self.apply_function_batch = import_string(self.apply_function_batch)
+
+ if self.apply_function:
+ apply_callable = partial(
+ self.apply_function, *self.apply_function_args, **self.apply_function_kwargs # type: ignore
+ )
+
+ if self.apply_function_batch:
+ apply_callable = partial(
+ self.apply_function_batch, # type: ignore
+ *self.apply_function_args,
+ **self.apply_function_kwargs,
+ )
+
+ messages_left = self.max_messages
+ messages_processed = 0
+
+ while (
+ messages_left > 0
+ ): # bool(True > 0) == True in the case where self.max_messages isn't set by the user
+
+ if not isinstance(messages_left, bool):
+ batch_size = self.max_batch_size if messages_left > self.max_batch_size else messages_left
+ else:
+ batch_size = self.max_batch_size
+
+ msgs = consumer.consume(num_messages=batch_size, timeout=self.poll_timeout)
+ messages_left -= len(msgs)
Review Comment:
int(True) is 1
```python
>>> messages_left = True
>>> messages_left -= 1
>>> messages_left
0
```
what should happen when `max_messages` is None? should we process a single batch? if it is the case, can you update the docstring? if not (and I think so), this should be fixed
##########
airflow/providers/apache/kafka/sensors/kafka.py:
##########
Review Comment:
Should we implement the no deferrable sensors and support defferable mode as an option as we do in the other sensors?
##########
airflow/providers/apache/kafka/hooks/base.py:
##########
@@ -0,0 +1,71 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+from typing import Any
+
+from confluent_kafka.admin import AdminClient
+
+from airflow.compat.functools import cached_property
+from airflow.hooks.base import BaseHook
+
+
+class KafkaHook(BaseHook):
Review Comment:
I was wondering if the name `KafkaBaseHook` would be more suitable for this class.
##########
airflow/providers/apache/kafka/hooks/produce.py:
##########
Review Comment:
same comment for `__exit__`
##########
airflow/providers/apache/kafka/operators/consume.py:
##########
@@ -0,0 +1,184 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+from functools import partial
+from typing import Any, Callable, Sequence
+
+from airflow.exceptions import AirflowException
+from airflow.models import BaseOperator
+from airflow.providers.apache.kafka.hooks.consume import KafkaConsumerHook
+from airflow.utils.module_loading import import_string
+
+VALID_COMMIT_CADENCE = {"never", "end_of_batch", "end_of_operator"}
+
+
+class ConsumeFromTopicOperator(BaseOperator):
+ """An operator that consumes from Kafka a topic(s) and processing the messages.
+
+ The operator creates a Kafka consumer that reads a batch of messages from the cluster and processes them
+ using the user supplied callable function. The consumer will continue to read in batches until it reaches
+ the end of the log or reads a maximum number of messages is reached.
+
+ :param kafka_config_id: The connection object to use, defaults to "kafka_default"
+ :param topics: A list of topics or regex patterns the consumer should subscribe to.
+ :param apply_function: The function that should be applied to fetched one at a time.
+ name of dag file executing the function and the function name delimited by a `.`
+ :param apply_function_batch: The function that should be applied to a batch of messages fetched. Can not
+ be used with `apply_function`. Intended for transactional workloads where an expensive task might
+ be called before or after operations on the messages are taken.
+ :param apply_function_args: Additional arguments that should be applied to the callable, defaults to None
+ :param apply_function_kwargs: Additional key word arguments that should be applied to the callable
+ defaults to None
+ :param commit_cadence: When consumers should commit offsets ("never", "end_of_batch","end_of_operator"),
+ defaults to "end_of_operator";
+ if end_of_operator, the commit() is called based on the max_messages arg. Commits are made after the
+ operator has processed the apply_function method for the maximum messages in the operator.
+ if end_of_batch, the commit() is called based on the max_batch_size arg. Commits are made after each
+ batch has processed by the apply_function method for all messages in the batch.
+ if never, close() is called without calling the commit() method.
+ :param max_messages: The maximum total number of messages an operator should read from Kafka,
+ defaults to None
+ :param max_batch_size: The maximum number of messages a consumer should read when polling,
+ defaults to 1000
+ :param poll_timeout: How long the Kafka consumer should wait before determining no more messages are
+ available, defaults to 60
+
+ .. seealso::
+ For more information on how to use this operator, take a look at the guide:
+ :ref:`howto/operator:ConsumeFromTopicOperator`
+ """
+
+ BLUE = "#ffefeb"
+ ui_color = BLUE
+ template_fields = (
+ "topics",
+ "apply_function",
+ "apply_function_args",
+ "apply_function_kwargs",
+ "kafka_config_id",
+ )
+
+ def __init__(
+ self,
+ topics: str | Sequence[str],
+ kafka_config_id: str = "kafka_default",
+ apply_function: Callable[..., Any] | str | None = None,
+ apply_function_batch: Callable[..., Any] | str | None = None,
+ apply_function_args: Sequence[Any] | None = None,
+ apply_function_kwargs: dict[Any, Any] | None = None,
+ commit_cadence: str | None = "end_of_operator",
+ max_messages: int | None = None,
+ max_batch_size: int = 1000,
+ poll_timeout: float = 60,
+ **kwargs: Any,
+ ) -> None:
+
+ super().__init__(**kwargs)
+
+ self.topics = topics
+ self.apply_function = apply_function
+ self.apply_function_batch = apply_function_batch
+ self.apply_function_args = apply_function_args or ()
+ self.apply_function_kwargs = apply_function_kwargs or {}
+ self.kafka_config_id = kafka_config_id
+ self.commit_cadence = commit_cadence
+ self.max_messages = max_messages or True
Review Comment:
what should happen when the user provide 0?
##########
airflow/providers/apache/kafka/operators/consume.py:
##########
@@ -0,0 +1,184 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+from functools import partial
+from typing import Any, Callable, Sequence
+
+from airflow.exceptions import AirflowException
+from airflow.models import BaseOperator
+from airflow.providers.apache.kafka.hooks.consume import KafkaConsumerHook
+from airflow.utils.module_loading import import_string
+
+VALID_COMMIT_CADENCE = {"never", "end_of_batch", "end_of_operator"}
Review Comment:
is it private or public set? and I see it defined in the other operators/sensors, we may define it in a utils module. Also it looks more like an Enum.
WDYT?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [airflow] dylanbstorey commented on a diff in pull request #30175: Add provider for Apache Kafka
Posted by "dylanbstorey (via GitHub)" <gi...@apache.org>.
dylanbstorey commented on code in PR #30175:
URL: https://github.com/apache/airflow/pull/30175#discussion_r1167626974
##########
airflow/providers/apache/kafka/sensors/kafka.py:
##########
Review Comment:
I've had these patterns around since ~1.10.2 and chose not to use them even in limited settings until deferrable existed because they were catastrophically bad. Kafka and Airflow already have plenty of footguns available - I'd prefer not to hand out more.
We can revisit this conversation as KIP-500 reaches more widespread adoption and I get some time to test some how Airflow 2.0 does with this pattern though.
Also in over a year in operation not a single user has asked for this feature so I'm happy to wait until we see a demand.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [airflow] dylanbstorey commented on a diff in pull request #30175: Add provider for Apache Kafka
Posted by "dylanbstorey (via GitHub)" <gi...@apache.org>.
dylanbstorey commented on code in PR #30175:
URL: https://github.com/apache/airflow/pull/30175#discussion_r1167626075
##########
airflow/providers/apache/kafka/hooks/client.py:
##########
Review Comment:
This client doesn't have a "close" method associated with it.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [airflow] kaxil commented on a diff in pull request #30175: Add provider for Apache Kafka
Posted by "kaxil (via GitHub)" <gi...@apache.org>.
kaxil commented on code in PR #30175:
URL: https://github.com/apache/airflow/pull/30175#discussion_r1153149735
##########
airflow/providers/apache/kafka/hooks/base.py:
##########
@@ -0,0 +1,71 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+from typing import Any
+
+from confluent_kafka.admin import AdminClient
+
+from airflow.hooks.base import BaseHook
+
+
+class KafkaHook(BaseHook):
+ """A base hook for interacting with Apache Kafka"""
+
+ conn_name_attr = "kafka_config_id"
+ default_conn_name = "kafka_config_default"
+ conn_type = "kafka"
+ hook_name = "Apache Kafka"
+
+ def __init__(self, kafka_config_id=default_conn_name, *args, **kwargs):
+ """Initialize our Base"""
+ super().__init__()
+ self.kafka_config_id = kafka_config_id
+ self.get_conn()
Review Comment:
Do we need to call this here, is it here for validation?
##########
airflow/providers/apache/kafka/hooks/base.py:
##########
@@ -0,0 +1,71 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+from typing import Any
+
+from confluent_kafka.admin import AdminClient
+
+from airflow.hooks.base import BaseHook
+
+
+class KafkaHook(BaseHook):
+ """A base hook for interacting with Apache Kafka"""
+
+ conn_name_attr = "kafka_config_id"
+ default_conn_name = "kafka_config_default"
+ conn_type = "kafka"
+ hook_name = "Apache Kafka"
+
+ def __init__(self, kafka_config_id=default_conn_name, *args, **kwargs):
+ """Initialize our Base"""
+ super().__init__()
+ self.kafka_config_id = kafka_config_id
+ self.get_conn()
+
+ @staticmethod
+ def get_ui_field_behaviour() -> dict[str, Any]:
+ """Returns custom field behaviour"""
+ return {
+ "hidden_fields": ["schema", "login", "password", "port", "host"],
+ "relabeling": {"extra": "Config Dict"},
+ "placeholders": {
+ "extra": '{"bootstrap.servers": "localhost:9092"}',
+ },
+ }
+
+ def get_conn(self) -> dict[Any, Any]:
+ """get the configuration object"""
+ config = self.get_connection(self.kafka_config_id).extra_dejson
+
+ if not (config.get("bootstrap.servers", None)):
+ raise ValueError("config['bootsrap.servers'] must be provided.")
+
+ return config
+
+ def test_connection(self) -> tuple[bool, str]:
+ """Test Connectivity from the UI"""
+ try:
+ t = AdminClient(self.get_conn(), timeout=10).list_topics()
+ if t:
+ return True, "Connection successful."
+ except Exception as e:
+ False, str(e)
+
+ return False, "Failed to establish connection."
+
+ pass
Review Comment:
```suggestion
```
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [airflow] dylanbstorey commented on a diff in pull request #30175: Add provider for Apache Kafka
Posted by "dylanbstorey (via GitHub)" <gi...@apache.org>.
dylanbstorey commented on code in PR #30175:
URL: https://github.com/apache/airflow/pull/30175#discussion_r1163374803
##########
docs/apache-airflow-providers-apache-kafka/connections/kafka.rst:
##########
@@ -0,0 +1,38 @@
+ .. Licensed to the Apache Software Foundation (ASF) under one
+ or more contributor license agreements. See the NOTICE file
+ distributed with this work for additional information
+ regarding copyright ownership. The ASF licenses this file
+ to you under the Apache License, Version 2.0 (the
+ "License"); you may not use this file except in compliance
+ with the License. You may obtain a copy of the License at
+
+ .. http://www.apache.org/licenses/LICENSE-2.0
+
+ .. Unless required by applicable law or agreed to in writing,
+ software distributed under the License is distributed on an
+ "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ KIND, either express or implied. See the License for the
+ specific language governing permissions and limitations
+ under the License.
+
+.. _howto/connection: kafka
+
+Apache Kafka Connection
+========================
+
+The Apache Kafka connection type configures a connection to Apache Kafka via the ``confluent-kafka`` Python package.
+
+Default Connection IDs
+----------------------
+
+Kafka hooks and operators use ``kafka_default`` by default.
+
+Configuring the Connection
+--------------------------
+
+Connections are configured as a json serializable string within provided to the ``extra`` field. A full list of parameters
+are described in the `Confluent Kafka python library <https://github.com/confluentinc/librdkafka/blob/master/CONFIGURATION.md>`_.
+
+If you are defining the Airflow connection from the Airflow UI, the ``extra`` field will be renamed to ``Config Dict``.
+
+Most operators and hooks will check that at the minimum ``bootstrap.servers`` key exists and has a value set to be valid.
Review Comment:
I'd really prefer to keep it "all in the config dict" so it matches confluent's defined interface directly.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [airflow] dylanbstorey commented on a diff in pull request #30175: Add provider for Apache Kafka
Posted by "dylanbstorey (via GitHub)" <gi...@apache.org>.
dylanbstorey commented on code in PR #30175:
URL: https://github.com/apache/airflow/pull/30175#discussion_r1141074907
##########
airflow/providers/apache/kafka/CHANGELOG.md:
##########
Review Comment:
> operators/_.py file name preferably be on service level and not action level. The action is the operator itself. That way we can have several operators organized in a single file.
I have a vague feeling about what you might be alluding to here and yeah. I'll re namespace everything into just `operators`, `hooks`, `triggers` to cut down on the repetitive nature of the intermediate module.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [airflow] josh-fell commented on a diff in pull request #30175: Add provider for Apache Kafka
Posted by "josh-fell (via GitHub)" <gi...@apache.org>.
josh-fell commented on code in PR #30175:
URL: https://github.com/apache/airflow/pull/30175#discussion_r1157200175
##########
airflow/providers/apache/kafka/example_dags/__init__.py:
##########
Review Comment:
The concept of Example DAGs has transitioned to System Tests via [AIP-47](https://cwiki.apache.org/confluence/display/AIRFLOW/AIP-47+New+design+of+Airflow+System+Tests). Since these DAGs exist in `system/providers/apache/kafka/` already, the `example_dags` directory can be removed.
##########
airflow/providers/apache/kafka/hooks/produce.py:
##########
@@ -0,0 +1,41 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+from confluent_kafka import Producer
+
+from airflow.providers.apache.kafka.hooks.base import KafkaHook
+
+
+class KafkaProducerHook(KafkaHook):
+ """
+ A hook for creating a Kafka Producer
+ :param kafka_config_id: The connection object to use, defaults to "kafka_default"
+ """
+
+ def __init__(self, kafka_config_id=KafkaHook.default_conn_name) -> None:
+ super().__init__(kafka_config_id=kafka_config_id)
+
+ if not self.get_conn().get("bootstrap.servers", None):
Review Comment:
Similarly, this validation is done within `get_conn()` itself. This doesn't look to be needed.
##########
airflow/providers/apache/kafka/operators/consume.py:
##########
@@ -0,0 +1,184 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+from functools import partial
+from typing import Any, Callable, Sequence
+
+from airflow.exceptions import AirflowException
+from airflow.models import BaseOperator
+from airflow.providers.apache.kafka.hooks.consume import KafkaConsumerHook
+from airflow.utils.module_loading import import_string
+
+VALID_COMMIT_CADENCE = {"never", "end_of_batch", "end_of_operator"}
+
+
+class ConsumeFromTopicOperator(BaseOperator):
+ """An operator that consumes from Kafka a topic(s) and processing the messages.
+
+ The operator creates a Kafka consumer that reads a batch of messages from the cluster and processes them
+ using the user supplied callable function. The consumer will continue to read in batches until it reaches
+ the end of the log or reads a maximum number of messages is reached.
+
+ :param kafka_config_id: The connection object to use, defaults to "kafka_default"
+ :param topics: A list of topics or regex patterns the consumer should subscribe to.
+ :param apply_function: The function that should be applied to fetched one at a time.
+ name of dag file executing the function and the function name delimited by a `.`
+ :param apply_function_batch: The function that should be applied to a batch of messages fetched. Can not
+ be used with `apply_function`. Intended for transactional workloads where an expensive task might
+ be called before or after operations on the messages are taken.
+ :param apply_function_args: Additional arguments that should be applied to the callable, defaults to None
+ :param apply_function_kwargs: Additional key word arguments that should be applied to the callable
+ defaults to None
+ :param commit_cadence: When consumers should commit offsets ("never", "end_of_batch","end_of_operator"),
+ defaults to "end_of_operator";
+ if end_of_operator, the commit() is called based on the max_messages arg. Commits are made after the
+ operator has processed the apply_function method for the maximum messages in the operator.
+ if end_of_batch, the commit() is called based on the max_batch_size arg. Commits are made after each
+ batch has processed by the apply_function method for all messages in the batch.
+ if never, close() is called without calling the commit() method.
+ :param max_messages: The maximum total number of messages an operator should read from Kafka,
+ defaults to None
+ :param max_batch_size: The maximum number of messages a consumer should read when polling,
+ defaults to 1000
+ :param poll_timeout: How long the Kafka consumer should wait before determining no more messages are
+ available, defaults to 60
+
+ .. seealso::
+ For more information on how to use this operator, take a look at the guide:
+ :ref:`howto/operator:ConsumeFromTopicOperator`
+ """
+
+ BLUE = "#ffefeb"
+ ui_color = BLUE
+ template_fields = (
+ "topics",
+ "apply_function",
+ "apply_function_args",
+ "apply_function_kwargs",
+ "kafka_config_id",
+ )
+
+ def __init__(
+ self,
+ topics: str | Sequence[str],
+ kafka_config_id: str = "kafka_default",
+ apply_function: Callable[..., Any] | str | None = None,
+ apply_function_batch: Callable[..., Any] | str | None = None,
+ apply_function_args: Sequence[Any] | None = None,
+ apply_function_kwargs: dict[Any, Any] | None = None,
+ commit_cadence: str | None = "end_of_operator",
+ max_messages: int | None = None,
+ max_batch_size: int = 1000,
+ poll_timeout: float = 60,
+ **kwargs: Any,
+ ) -> None:
+
+ super().__init__(**kwargs)
+
+ self.topics = topics
+ self.apply_function = apply_function
+ self.apply_function_batch = apply_function_batch
+ self.apply_function_args = apply_function_args or ()
+ self.apply_function_kwargs = apply_function_kwargs or {}
+ self.kafka_config_id = kafka_config_id
+ self.commit_cadence = commit_cadence
+ self.max_messages = max_messages or True
+ self.max_batch_size = max_batch_size
+ self.poll_timeout = poll_timeout
+
+ if self.commit_cadence not in VALID_COMMIT_CADENCE:
+ raise AirflowException(
+ f"commit_cadence must be one of {VALID_COMMIT_CADENCE}. Got {self.commit_cadence}"
+ )
+
+ if self.max_messages and self.max_batch_size > self.max_messages:
+ self.log.warning(
+ "max_batch_size (%s) > max_messages" " (%s). Setting max_messages to" " %s ",
Review Comment:
```suggestion
"max_batch_size (%s) > max_messages (%s). Setting max_messages to %s ",
```
Probably some funny business from `black` re-formatting.
##########
airflow/providers/apache/kafka/hooks/produce.py:
##########
@@ -0,0 +1,41 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+from confluent_kafka import Producer
+
+from airflow.providers.apache.kafka.hooks.base import KafkaHook
+
+
+class KafkaProducerHook(KafkaHook):
+ """
+ A hook for creating a Kafka Producer
+ :param kafka_config_id: The connection object to use, defaults to "kafka_default"
Review Comment:
```suggestion
A hook for creating a Kafka Producer
:param kafka_config_id: The connection object to use, defaults to "kafka_default"
```
Otherwise the parameters won't render correctly in the Python API docs.
##########
airflow/providers/apache/kafka/operators/consume.py:
##########
@@ -0,0 +1,184 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+from functools import partial
+from typing import Any, Callable, Sequence
+
+from airflow.exceptions import AirflowException
+from airflow.models import BaseOperator
+from airflow.providers.apache.kafka.hooks.consume import KafkaConsumerHook
+from airflow.utils.module_loading import import_string
+
+VALID_COMMIT_CADENCE = {"never", "end_of_batch", "end_of_operator"}
+
+
+class ConsumeFromTopicOperator(BaseOperator):
+ """An operator that consumes from Kafka a topic(s) and processing the messages.
+
+ The operator creates a Kafka consumer that reads a batch of messages from the cluster and processes them
+ using the user supplied callable function. The consumer will continue to read in batches until it reaches
+ the end of the log or reads a maximum number of messages is reached.
+
+ :param kafka_config_id: The connection object to use, defaults to "kafka_default"
+ :param topics: A list of topics or regex patterns the consumer should subscribe to.
+ :param apply_function: The function that should be applied to fetched one at a time.
+ name of dag file executing the function and the function name delimited by a `.`
+ :param apply_function_batch: The function that should be applied to a batch of messages fetched. Can not
+ be used with `apply_function`. Intended for transactional workloads where an expensive task might
+ be called before or after operations on the messages are taken.
+ :param apply_function_args: Additional arguments that should be applied to the callable, defaults to None
+ :param apply_function_kwargs: Additional key word arguments that should be applied to the callable
+ defaults to None
+ :param commit_cadence: When consumers should commit offsets ("never", "end_of_batch","end_of_operator"),
+ defaults to "end_of_operator";
+ if end_of_operator, the commit() is called based on the max_messages arg. Commits are made after the
+ operator has processed the apply_function method for the maximum messages in the operator.
+ if end_of_batch, the commit() is called based on the max_batch_size arg. Commits are made after each
+ batch has processed by the apply_function method for all messages in the batch.
+ if never, close() is called without calling the commit() method.
+ :param max_messages: The maximum total number of messages an operator should read from Kafka,
+ defaults to None
+ :param max_batch_size: The maximum number of messages a consumer should read when polling,
+ defaults to 1000
+ :param poll_timeout: How long the Kafka consumer should wait before determining no more messages are
+ available, defaults to 60
+
+ .. seealso::
+ For more information on how to use this operator, take a look at the guide:
+ :ref:`howto/operator:ConsumeFromTopicOperator`
+ """
+
+ BLUE = "#ffefeb"
+ ui_color = BLUE
+ template_fields = (
+ "topics",
+ "apply_function",
+ "apply_function_args",
+ "apply_function_kwargs",
+ "kafka_config_id",
+ )
+
+ def __init__(
+ self,
+ topics: str | Sequence[str],
+ kafka_config_id: str = "kafka_default",
+ apply_function: Callable[..., Any] | str | None = None,
+ apply_function_batch: Callable[..., Any] | str | None = None,
+ apply_function_args: Sequence[Any] | None = None,
+ apply_function_kwargs: dict[Any, Any] | None = None,
+ commit_cadence: str | None = "end_of_operator",
+ max_messages: int | None = None,
+ max_batch_size: int = 1000,
+ poll_timeout: float = 60,
+ **kwargs: Any,
+ ) -> None:
+
+ super().__init__(**kwargs)
+
+ self.topics = topics
+ self.apply_function = apply_function
+ self.apply_function_batch = apply_function_batch
+ self.apply_function_args = apply_function_args or ()
+ self.apply_function_kwargs = apply_function_kwargs or {}
+ self.kafka_config_id = kafka_config_id
+ self.commit_cadence = commit_cadence
+ self.max_messages = max_messages or True
+ self.max_batch_size = max_batch_size
+ self.poll_timeout = poll_timeout
+
+ if self.commit_cadence not in VALID_COMMIT_CADENCE:
+ raise AirflowException(
+ f"commit_cadence must be one of {VALID_COMMIT_CADENCE}. Got {self.commit_cadence}"
+ )
+
+ if self.max_messages and self.max_batch_size > self.max_messages:
+ self.log.warning(
+ "max_batch_size (%s) > max_messages" " (%s). Setting max_messages to" " %s ",
+ self.max_batch_size,
+ self.max_messages,
+ self.max_batch_size,
+ )
+
+ if self.commit_cadence == "never":
+ self.commit_cadence = None
+
+ if apply_function and apply_function_batch:
+ raise AirflowException(
+ "One of apply_function or apply_function_batch must be supplied, not both."
+ )
+
+ def execute(self, context) -> Any:
+
+ consumer = KafkaConsumerHook(topics=self.topics, kafka_config_id=self.kafka_config_id).get_consumer()
+
+ if isinstance(self.apply_function, str):
+ self.apply_function = import_string(self.apply_function)
+
+ if isinstance(self.apply_function_batch, str):
+ self.apply_function_batch = import_string(self.apply_function_batch)
+
+ if self.apply_function:
+ apply_callable = partial(
+ self.apply_function, *self.apply_function_args, **self.apply_function_kwargs # type: ignore
+ )
+
+ if self.apply_function_batch:
+ apply_callable = partial(
+ self.apply_function_batch, # type: ignore
+ *self.apply_function_args,
+ **self.apply_function_kwargs,
+ )
+
+ messages_left = self.max_messages
+ messages_processed = 0
+
+ while (
+ messages_left > 0
+ ): # bool(True > 0) == True in the case where self.max_messages isn't set by the user
+
+ if not isinstance(messages_left, bool):
+ batch_size = self.max_batch_size if messages_left > self.max_batch_size else messages_left
+ else:
+ batch_size = self.max_batch_size
+
+ msgs = consumer.consume(num_messages=batch_size, timeout=self.poll_timeout)
+ messages_left -= len(msgs)
+ messages_processed += len(msgs)
+
+ if not msgs: # No messages + messages_left is being used.
+ self.log.info("Reached end of log. Exiting.")
+ break
+
+ if self.apply_function:
+ for m in msgs:
+ apply_callable(m)
+
+ if self.apply_function_batch:
+ apply_callable(msgs)
+
+ if self.commit_cadence == "end_of_batch":
+ self.log.info("committing offset at %s", self.commit_cadence)
+ consumer.commit()
+
+ if self.commit_cadence:
+ self.log.info("committing offset at %s", self.commit_cadence)
+ consumer.commit()
+
+ consumer.close()
+
+ return
Review Comment:
```suggestion
```
Not necessary and the return type annotation for `execute()` can be updated as well.
##########
airflow/providers/apache/kafka/operators/produce.py:
##########
@@ -0,0 +1,130 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+import logging
+from functools import partial
+from typing import Any, Callable, Sequence
+
+from airflow.exceptions import AirflowException
+from airflow.models import BaseOperator
+from airflow.providers.apache.kafka.hooks.produce import KafkaProducerHook
+from airflow.utils.module_loading import import_string
+
+local_logger = logging.getLogger("airflow")
+
+
+def acked(err, msg):
+ if err is not None:
+ local_logger.error(f"Failed to deliver message: {err}")
+ else:
+ local_logger.info(
+ f"Produced record to topic {msg.topic()} partition [{msg.partition()}] @ offset {msg.offset()}"
+ )
+
+
+class ProduceToTopicOperator(BaseOperator):
+ """An operator that produces messages to a Kafka topic
+
+ Registers a producer to a kafka topic and publishes messages to the log.
+
+ :param kafka_config_id: The connection object to use, defaults to "kafka_default"
+ :param topic: The topic the producer should produce to, defaults to None
+ :param producer_function: The function that generates key/value pairs as messages for production,
+ defaults to None
+ :param producer_function_args: Additional arguments to be applied to the producer callable,
+ defaults to None
+ :param producer_function_kwargs: Additional keyword arguments to be applied to the producer callable,
+ defaults to None
+ :param delivery_callback: The callback to apply after delivery(or failure) of a message, defaults to None
+ :param synchronous: If writes to kafka should be fully synchronous, defaults to True
+ :param poll_timeout: How long of a delay should be applied when calling poll after production to kafka,
+ defaults to 0
+ :raises AirflowException: _description_
+
+ .. seealso::
+ For more information on how to use this operator, take a look at the guide:
+ :ref:`howto/operator:ProduceToTopicOperator`
+ """
+
+ template_fields = (
+ "topic",
+ "producer_function",
+ "producer_function_args",
+ "producer_function_kwargs",
+ "kafka_config_id",
+ )
+
+ def __init__(
+ self,
+ topic: str,
+ producer_function: str | Callable[..., Any],
+ kafka_config_id: str = "kafka_default",
+ producer_function_args: Sequence[Any] | None = None,
+ producer_function_kwargs: dict[Any, Any] | None = None,
+ delivery_callback: str | None = None,
+ synchronous: bool = True,
+ poll_timeout: float = 0,
+ **kwargs: Any,
+ ) -> None:
+
+ super().__init__(**kwargs)
+
+ if delivery_callback:
+ dc = import_string(delivery_callback)
+ else:
+ dc = acked
+
+ self.kafka_config_id = kafka_config_id
+ self.topic = topic
+ self.producer_function = producer_function
+ self.producer_function_args = producer_function_args or ()
+ self.producer_function_kwargs = producer_function_kwargs or {}
+ self.delivery_callback = dc
+ self.synchronous = synchronous
+ self.poll_timeout = poll_timeout
+
+ if not (self.topic and self.producer_function):
+ raise AirflowException(
+ "topic and producer_function must be provided. Got topic="
+ + f"{self.topic} and producer_function={self.producer_function}"
Review Comment:
```suggestion
f"{self.topic} and producer_function={self.producer_function}"
```
The `+` operator isn't needed. The string will still concatenate as expected because of [string literal concatentation](https://docs.python.org/2/reference/lexical_analysis.html#string-literal-concatenation).
##########
tests/providers/apache/kafka/hooks/test_admin_client.py:
##########
@@ -0,0 +1,71 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+import json
+import logging
+
+import pytest
+
+from airflow.models import Connection
+from airflow.providers.apache.kafka.hooks.base import KafkaHook
+from airflow.providers.apache.kafka.hooks.client import KafkaAdminClientHook
+from airflow.utils import db
+
+log = logging.getLogger(__name__)
+
+
+class TestSampleHook:
+ """
+ Test Admin Client Hook.
+ """
+
+ def setup_method(self):
+ db.merge_conn(
+ Connection(
+ conn_id="kafka_default",
+ conn_type="kafka",
+ extra=json.dumps(
+ {"socket.timeout.ms": 10, "bootstrap.servers": "localhost:9092", "group.id": "test_group"}
+ ),
+ )
+ )
+
+ db.merge_conn(
+ Connection(
+ conn_id="kafka_bad",
+ conn_type="kafka",
+ extra=json.dumps({"socket.timeout.ms": 10}),
+ )
+ )
+
+ def test_init(self):
+ """test initialization of AdminClientHook"""
+
+ # Standard Init
+ KafkaAdminClientHook(kafka_config_id="kafka_default")
+
+ # # Not Enough Args
+ with pytest.raises(ValueError):
+ KafkaAdminClientHook(kafka_config_id="kafka_bad")
+
+ def test_create_topic(self, mocker):
Review Comment:
What is the assertion being tested here?
##########
tests/providers/apache/kafka/hooks/test_consumer.py:
##########
@@ -0,0 +1,62 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+import json
+
+import pytest
+
+from airflow.models import Connection
+
+# Import Hook
+from airflow.providers.apache.kafka.hooks.consume import KafkaConsumerHook
+from airflow.utils import db
+
+
+class TestConsumerHook:
+ """
+ Test consumer hook.
+ """
+
+ def setup_method(self):
+ db.merge_conn(
+ Connection(
+ conn_id="kafka_default",
+ conn_type="kafka",
+ extra=json.dumps(
+ {"socket.timeout.ms": 10, "bootstrap.servers": "localhost:9092", "group.id": "test_group"}
+ ),
+ )
+ )
+
+ db.merge_conn(
+ Connection(
+ conn_id="kafka_bad",
+ conn_type="kafka",
+ extra=json.dumps({}),
+ )
+ )
+
+ def test_init(self):
+ """test initialization of AdminClientHook"""
+
+ # Standard Init
+ KafkaConsumerHook(["test_1"], kafka_config_id="kafka_default")
Review Comment:
Should probably at least assert instance attrs of the hook are set as expected.
##########
airflow/providers/apache/kafka/operators/produce.py:
##########
@@ -0,0 +1,130 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+import logging
+from functools import partial
+from typing import Any, Callable, Sequence
+
+from airflow.exceptions import AirflowException
+from airflow.models import BaseOperator
+from airflow.providers.apache.kafka.hooks.produce import KafkaProducerHook
+from airflow.utils.module_loading import import_string
+
+local_logger = logging.getLogger("airflow")
+
+
+def acked(err, msg):
+ if err is not None:
+ local_logger.error(f"Failed to deliver message: {err}")
+ else:
+ local_logger.info(
+ f"Produced record to topic {msg.topic()} partition [{msg.partition()}] @ offset {msg.offset()}"
+ )
+
+
+class ProduceToTopicOperator(BaseOperator):
+ """An operator that produces messages to a Kafka topic
+
+ Registers a producer to a kafka topic and publishes messages to the log.
+
+ :param kafka_config_id: The connection object to use, defaults to "kafka_default"
+ :param topic: The topic the producer should produce to, defaults to None
+ :param producer_function: The function that generates key/value pairs as messages for production,
+ defaults to None
+ :param producer_function_args: Additional arguments to be applied to the producer callable,
+ defaults to None
+ :param producer_function_kwargs: Additional keyword arguments to be applied to the producer callable,
+ defaults to None
+ :param delivery_callback: The callback to apply after delivery(or failure) of a message, defaults to None
+ :param synchronous: If writes to kafka should be fully synchronous, defaults to True
+ :param poll_timeout: How long of a delay should be applied when calling poll after production to kafka,
+ defaults to 0
+ :raises AirflowException: _description_
+
+ .. seealso::
+ For more information on how to use this operator, take a look at the guide:
+ :ref:`howto/operator:ProduceToTopicOperator`
+ """
+
+ template_fields = (
+ "topic",
+ "producer_function",
+ "producer_function_args",
+ "producer_function_kwargs",
+ "kafka_config_id",
+ )
+
+ def __init__(
+ self,
+ topic: str,
+ producer_function: str | Callable[..., Any],
+ kafka_config_id: str = "kafka_default",
+ producer_function_args: Sequence[Any] | None = None,
+ producer_function_kwargs: dict[Any, Any] | None = None,
+ delivery_callback: str | None = None,
+ synchronous: bool = True,
+ poll_timeout: float = 0,
+ **kwargs: Any,
+ ) -> None:
+
+ super().__init__(**kwargs)
+
+ if delivery_callback:
+ dc = import_string(delivery_callback)
+ else:
+ dc = acked
+
+ self.kafka_config_id = kafka_config_id
+ self.topic = topic
+ self.producer_function = producer_function
+ self.producer_function_args = producer_function_args or ()
+ self.producer_function_kwargs = producer_function_kwargs or {}
+ self.delivery_callback = dc
+ self.synchronous = synchronous
+ self.poll_timeout = poll_timeout
+
+ if not (self.topic and self.producer_function):
+ raise AirflowException(
+ "topic and producer_function must be provided. Got topic="
+ + f"{self.topic} and producer_function={self.producer_function}"
+ )
+
+ return
+
+ def execute(self, context) -> Any:
Review Comment:
```suggestion
def execute(self, context) -> None:
```
##########
tests/providers/apache/kafka/operators/test_consume.py:
##########
Review Comment:
What are the assertions in these tests either positive or negative? Same comment in test_produce.py as well.
##########
airflow/providers/apache/kafka/hooks/consume.py:
##########
@@ -0,0 +1,50 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+from typing import Sequence
+
+from confluent_kafka import Consumer
+
+from airflow.providers.apache.kafka.hooks.base import KafkaHook
+
+
+class KafkaConsumerHook(KafkaHook):
+ """
+ A hook for creating a Kafka Consumer
+
+ :param kafka_config_id: The connection object to use, defaults to "kafka_config_default"
+ :param topics: A list of topics to subscribe to.
+ """
+
+ def __init__(self, topics: Sequence[str], kafka_config_id=KafkaHook.default_conn_name) -> None:
+
+ super().__init__(kafka_config_id=kafka_config_id)
+ self.topics = topics
+ if not self.get_conn().get("group.id", None):
+ raise ValueError("The 'group.id' parameter must be set in the config dictionary'. Got <None>")
+
+ if not self.get_conn().get("bootstrap.servers", None):
Review Comment:
+1 for this suggestion in the base hook.
Also, `get_conn()` already raises a `ValueError` if "bootstrap.servers" doesn't exist in the connection extras.
##########
tests/providers/apache/kafka/hooks/test_admin_client.py:
##########
Review Comment:
Some coverage on `get_conn()` (or if it turns into a cached_property) in all of the hook tests generally (even in test_consumer.py and test_producer.py) would be great.
##########
tests/providers/apache/kafka/hooks/test_admin_client.py:
##########
@@ -0,0 +1,71 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+import json
+import logging
+
+import pytest
+
+from airflow.models import Connection
+from airflow.providers.apache.kafka.hooks.base import KafkaHook
+from airflow.providers.apache.kafka.hooks.client import KafkaAdminClientHook
+from airflow.utils import db
+
+log = logging.getLogger(__name__)
+
+
+class TestSampleHook:
+ """
+ Test Admin Client Hook.
+ """
+
+ def setup_method(self):
+ db.merge_conn(
+ Connection(
+ conn_id="kafka_default",
+ conn_type="kafka",
+ extra=json.dumps(
+ {"socket.timeout.ms": 10, "bootstrap.servers": "localhost:9092", "group.id": "test_group"}
+ ),
+ )
+ )
+
+ db.merge_conn(
+ Connection(
+ conn_id="kafka_bad",
+ conn_type="kafka",
+ extra=json.dumps({"socket.timeout.ms": 10}),
+ )
+ )
+
+ def test_init(self):
+ """test initialization of AdminClientHook"""
+
+ # Standard Init
+ KafkaAdminClientHook(kafka_config_id="kafka_default")
+
+ # # Not Enough Args
+ with pytest.raises(ValueError):
+ KafkaAdminClientHook(kafka_config_id="kafka_bad")
+
+ def test_create_topic(self, mocker):
Review Comment:
WDYT about adding a test re: topics that already exist?
##########
docs/apache-airflow-providers-apache-kafka/connections/kafka.rst:
##########
@@ -0,0 +1,38 @@
+ .. Licensed to the Apache Software Foundation (ASF) under one
+ or more contributor license agreements. See the NOTICE file
+ distributed with this work for additional information
+ regarding copyright ownership. The ASF licenses this file
+ to you under the Apache License, Version 2.0 (the
+ "License"); you may not use this file except in compliance
+ with the License. You may obtain a copy of the License at
+
+ .. http://www.apache.org/licenses/LICENSE-2.0
+
+ .. Unless required by applicable law or agreed to in writing,
+ software distributed under the License is distributed on an
+ "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ KIND, either express or implied. See the License for the
+ specific language governing permissions and limitations
+ under the License.
+
+.. _howto/connection: kafka
+
+Apache Kafka Connection
+========================
+
+The Apache Kafka connection type configures a connection to Apache Kafka via the ``confluent-kafka`` Python package.
+
+Default Connection IDs
+----------------------
+
+Kafka hooks and operators use ``kafka_default`` by default.
+
+Configuring the Connection
+--------------------------
+
+Connections are configured as a json serializable string within provided to the ``extra`` field. A full list of parameters
+are described in the `Confluent Kafka python library <https://github.com/confluentinc/librdkafka/blob/master/CONFIGURATION.md>`_.
+
+If you are defining the Airflow connection from the Airflow UI, the ``extra`` field will be renamed to ``Config Dict``.
+
+Most operators and hooks will check that at the minimum ``bootstrap.servers`` key exists and has a value set to be valid.
Review Comment:
Since checking that "bootstrap.servers" exists in done within `get_conn()`, this config would be required for all hook, operator, and sensor use right?
If so, I think it would be wise to make Bootstrap Servers an explicit field in the Airflow UI too.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [airflow] dylanbstorey commented on a diff in pull request #30175: Add provider for Apache Kafka
Posted by "dylanbstorey (via GitHub)" <gi...@apache.org>.
dylanbstorey commented on code in PR #30175:
URL: https://github.com/apache/airflow/pull/30175#discussion_r1167627022
##########
airflow/providers/apache/kafka/sensors/kafka.py:
##########
@@ -0,0 +1,219 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+from typing import Any, Callable, Sequence
+
+from airflow.models import BaseOperator
+from airflow.providers.apache.kafka.triggers.await_message import AwaitMessageTrigger
+
+VALID_COMMIT_CADENCE = {"never", "end_of_batch", "end_of_operator"}
+
+
+class AwaitMessageSensor(BaseOperator):
+ """An Airflow operator that defers until a specific message is published to Kafka.
+
+ The operator creates a consumer that reads the Kafka log until it encounters a positive event.
+
+ The behavior of the consumer for this trigger is as follows:
+ - poll the Kafka topics for a message
+ - if no message returned, sleep
+ - process the message with provided callable and commit the message offset
+ - if callable returns any data, raise a TriggerEvent with the return data
+ - else continue to next message
+ - return event (as default xcom or specific xcom key)
+
+
+ :param kafka_config_id: The connection object to use, defaults to "kafka_default"
+ :param topics: Topics (or topic regex) to use for reading from
+ :param apply_function: The function to apply to messages to determine if an event occurred. As a dot
+ notation string.
+ :param apply_function_args: Arguments to be applied to the processing function,
+ defaults to None
+ :param apply_function_kwargs: Key word arguments to be applied to the processing function,
+ defaults to None
+ :param poll_timeout: How long the kafka consumer should wait for a message to arrive from the kafka
+ cluster,defaults to 1
+ :param poll_interval: How long the kafka consumer should sleep after reaching the end of the Kafka log,
+ defaults to 5
+ :param xcom_push_key: the name of a key to push the returned message to, defaults to None
+
+
+ """
+
+ BLUE = "#ffefeb"
+ ui_color = BLUE
+
+ template_fields = (
+ "topics",
+ "apply_function",
+ "apply_function_args",
+ "apply_function_kwargs",
+ "kafka_config_id",
+ )
+
+ def __init__(
+ self,
+ topics: Sequence[str],
+ apply_function: str,
+ kafka_config_id: str = "kafka_default",
+ apply_function_args: Sequence[Any] | None = None,
+ apply_function_kwargs: dict[Any, Any] | None = None,
+ poll_timeout: float = 1,
+ poll_interval: float = 5,
+ xcom_push_key=None,
+ **kwargs: Any,
+ ) -> None:
+
+ super().__init__(**kwargs)
+
+ self.topics = topics
+ self.apply_function = apply_function
+ self.apply_function_args = apply_function_args
+ self.apply_function_kwargs = apply_function_kwargs
+ self.kafka_config_id = kafka_config_id
+ self.poll_timeout = poll_timeout
+ self.poll_interval = poll_interval
+ self.xcom_push_key = xcom_push_key
+
+ def execute(self, context) -> Any:
+
+ self.defer(
+ trigger=AwaitMessageTrigger(
+ topics=self.topics,
+ apply_function=self.apply_function,
+ apply_function_args=self.apply_function_args,
+ apply_function_kwargs=self.apply_function_kwargs,
+ kafka_config_id=self.kafka_config_id,
+ poll_timeout=self.poll_timeout,
+ poll_interval=self.poll_interval,
+ ),
+ method_name="execute_complete",
+ )
+
+ def execute_complete(self, context, event=None):
+ if self.xcom_push_key:
+ self.xcom_push(context, key=self.xcom_push_key, value=event)
+ return event
+
+
+class AwaitMessageTriggerFunctionSensor(BaseOperator):
+ """An Airflow operator that defers until a specific message is published to
+ Kafka, then triggers a registered function
+
+
+ The behavior of the consumer for this trigger is as follows:
+ - poll the Kafka topics for a message
+ - if no message returned, sleep
+ - process the message with provided callable and commit the message offset
+ - if callable returns any data, raise a TriggerEvent with the return data
+ - else continue to next message
+ - return event (as default xcom or specific xcom key)
+
+
+ :param kafka_config_id: The connection object to use, defaults to "kafka_default"
+ :param topics: Topics (or topic regex) to use for reading from
+ :param apply_function: The function to apply to messages to determine if an event occurred. As a dot
+ notation string.
+ :param event_triggered_function: The callable to trigger once the apply_function encounters a
+ positive event.
+ :param apply_function_args: Arguments to be applied to the processing function, defaults to None
+ :param apply_function_kwargs: Key word arguments to be applied to the processing function,
+ defaults to None
+ :param poll_timeout: How long the kafka consumer should wait for a message to arrive from the kafka
+ cluster, defaults to 1
+ :param poll_interval: How long the kafka consumer should sleep after reaching the end of the Kafka log,
+ defaults to 5
+ :param xcom_push_key: the name of a key to push the returned message to, defaults to None
Review Comment:
Good eye.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [airflow] dylanbstorey commented on a diff in pull request #30175: Add provider for Apache Kafka
Posted by "dylanbstorey (via GitHub)" <gi...@apache.org>.
dylanbstorey commented on code in PR #30175:
URL: https://github.com/apache/airflow/pull/30175#discussion_r1167627397
##########
airflow/providers/apache/kafka/operators/consume.py:
##########
@@ -0,0 +1,184 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+from functools import partial
+from typing import Any, Callable, Sequence
+
+from airflow.exceptions import AirflowException
+from airflow.models import BaseOperator
+from airflow.providers.apache.kafka.hooks.consume import KafkaConsumerHook
+from airflow.utils.module_loading import import_string
+
+VALID_COMMIT_CADENCE = {"never", "end_of_batch", "end_of_operator"}
+
+
+class ConsumeFromTopicOperator(BaseOperator):
+ """An operator that consumes from Kafka a topic(s) and processing the messages.
+
+ The operator creates a Kafka consumer that reads a batch of messages from the cluster and processes them
+ using the user supplied callable function. The consumer will continue to read in batches until it reaches
+ the end of the log or reads a maximum number of messages is reached.
+
+ :param kafka_config_id: The connection object to use, defaults to "kafka_default"
+ :param topics: A list of topics or regex patterns the consumer should subscribe to.
+ :param apply_function: The function that should be applied to fetched one at a time.
+ name of dag file executing the function and the function name delimited by a `.`
+ :param apply_function_batch: The function that should be applied to a batch of messages fetched. Can not
+ be used with `apply_function`. Intended for transactional workloads where an expensive task might
+ be called before or after operations on the messages are taken.
+ :param apply_function_args: Additional arguments that should be applied to the callable, defaults to None
+ :param apply_function_kwargs: Additional key word arguments that should be applied to the callable
+ defaults to None
+ :param commit_cadence: When consumers should commit offsets ("never", "end_of_batch","end_of_operator"),
+ defaults to "end_of_operator";
+ if end_of_operator, the commit() is called based on the max_messages arg. Commits are made after the
+ operator has processed the apply_function method for the maximum messages in the operator.
+ if end_of_batch, the commit() is called based on the max_batch_size arg. Commits are made after each
+ batch has processed by the apply_function method for all messages in the batch.
+ if never, close() is called without calling the commit() method.
+ :param max_messages: The maximum total number of messages an operator should read from Kafka,
+ defaults to None
+ :param max_batch_size: The maximum number of messages a consumer should read when polling,
+ defaults to 1000
+ :param poll_timeout: How long the Kafka consumer should wait before determining no more messages are
+ available, defaults to 60
+
+ .. seealso::
+ For more information on how to use this operator, take a look at the guide:
+ :ref:`howto/operator:ConsumeFromTopicOperator`
+ """
+
+ BLUE = "#ffefeb"
+ ui_color = BLUE
+ template_fields = (
+ "topics",
+ "apply_function",
+ "apply_function_args",
+ "apply_function_kwargs",
+ "kafka_config_id",
+ )
+
+ def __init__(
+ self,
+ topics: str | Sequence[str],
+ kafka_config_id: str = "kafka_default",
+ apply_function: Callable[..., Any] | str | None = None,
+ apply_function_batch: Callable[..., Any] | str | None = None,
+ apply_function_args: Sequence[Any] | None = None,
+ apply_function_kwargs: dict[Any, Any] | None = None,
+ commit_cadence: str | None = "end_of_operator",
+ max_messages: int | None = None,
+ max_batch_size: int = 1000,
+ poll_timeout: float = 60,
+ **kwargs: Any,
+ ) -> None:
+
+ super().__init__(**kwargs)
+
+ self.topics = topics
+ self.apply_function = apply_function
+ self.apply_function_batch = apply_function_batch
+ self.apply_function_args = apply_function_args or ()
+ self.apply_function_kwargs = apply_function_kwargs or {}
+ self.kafka_config_id = kafka_config_id
+ self.commit_cadence = commit_cadence
+ self.max_messages = max_messages or True
Review Comment:
```python
>>> y = 0 or True
>>> y
True
```
^that.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [airflow] dylanbstorey commented on pull request #30175: Add provider for Apache Kafka
Posted by "dylanbstorey (via GitHub)" <gi...@apache.org>.
dylanbstorey commented on PR #30175:
URL: https://github.com/apache/airflow/pull/30175#issuecomment-1509982464
>And for the responsibility of the method get_conn, I prefer to have a small discussion about it.
By all means but your inital comment boiled down to "i prefer this" and as you point out this is to your knowledge at least an un documented convention. It also doesn't impact the external API or require a breaking change so I think "we can revisit this in the future" is an appropriate response.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [airflow] hussein-awala commented on pull request #30175: Add provider for Apache Kafka
Posted by "hussein-awala (via GitHub)" <gi...@apache.org>.
hussein-awala commented on PR #30175:
URL: https://github.com/apache/airflow/pull/30175#issuecomment-1514031935
Need to merge main and resolve the conflicts
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [airflow] dylanbstorey commented on pull request #30175: Add provider for Apache Kafka
Posted by "dylanbstorey (via GitHub)" <gi...@apache.org>.
dylanbstorey commented on PR #30175:
URL: https://github.com/apache/airflow/pull/30175#issuecomment-1510546320
@raphaelauv
I'm going to need more to go on.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [airflow] potiuk commented on pull request #30175: Add provider for Apache Kafka
Posted by "potiuk (via GitHub)" <gi...@apache.org>.
potiuk commented on PR #30175:
URL: https://github.com/apache/airflow/pull/30175#issuecomment-1513749534
FYI. The Kafka Provider has been accepted by Lazy Consensus now, and I see we still have some iterations to complete, but this can now be done with the realm of being "blessed" by the community :) https://lists.apache.org/thread/y85s9ql68xyszt1ph0ljdz312nr2pvtw
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [airflow] josh-fell commented on a diff in pull request #30175: Add provider for Apache Kafka
Posted by "josh-fell (via GitHub)" <gi...@apache.org>.
josh-fell commented on code in PR #30175:
URL: https://github.com/apache/airflow/pull/30175#discussion_r1140900552
##########
tests/providers/apache/kafka/hooks/test_admin_client.py:
##########
@@ -0,0 +1,69 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+import logging
+import unittest
+
+import pytest
+
+from airflow import AirflowException
+
+# Import Hook
+from airflow.providers.apache.kafka.hooks.admin_client import KafkaAdminClientHook
+
+log = logging.getLogger(__name__)
+
+
+class TestSampleHook(unittest.TestCase):
Review Comment:
Relatedly: https://github.com/apache/airflow/issues/29305. Tests shouldn't be unittest cases.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [airflow] josh-fell commented on a diff in pull request #30175: Add provider for Apache Kafka
Posted by "josh-fell (via GitHub)" <gi...@apache.org>.
josh-fell commented on code in PR #30175:
URL: https://github.com/apache/airflow/pull/30175#discussion_r1140876165
##########
airflow/providers/apache/kafka/hooks/consumer.py:
##########
@@ -0,0 +1,60 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+from typing import Any, Sequence
+
+from confluent_kafka import Consumer
+
+from airflow.exceptions import AirflowException
+from airflow.hooks.base import BaseHook
+
+
+class KafkaConsumerHook(BaseHook):
+ """KafkaConsumerHook
+
+ A hook for creating a Kafka Consumer
+ :param config: A config dictionary to use with confluent_kafka library, defaults to None
+ """
+
+
+ def __init__(
+ self,
+ topics: Sequence[str],
+ config: dict[Any, Any] | None = None,
+ ) -> None:
+ super().__init__()
+
+ self.config: dict[Any, Any] = config or {}
+ self.topics = topics
+
+ if not self.config.get("group.id", None):
+ raise AirflowException(
+ "The 'group.id' parameter must be set in the config dictionary'. Got <None>"
+ )
+
+ if not self.config.get("bootstrap.servers", None):
+ raise AirflowException("config['bootsrap.servers'] must be provided.")
Review Comment:
Should `group.id` and `bootstrap.servers` be their own explicit params for this hook?
##########
airflow/providers/apache/kafka/operators/consume_from_topic.py:
##########
@@ -0,0 +1,170 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+from functools import partial
+from typing import Any, Callable, Sequence
+
+from airflow.exceptions import AirflowException
+from airflow.models import BaseOperator
+from airflow.providers.apache.kafka.hooks.consumer import KafkaConsumerHook
+from airflow.providers.apache.kafka.shared_utils import _get_callable
+
+VALID_COMMIT_CADENCE = {"never", "end_of_batch", "end_of_operator"}
+
+
+class ConsumeFromTopicOperator(BaseOperator):
+ """ConsumeFromTopicOperator An operator that consumes from Kafka a topic(s) and processing the messages.
+
+ The operator creates a Kafka consumer that reads a batch of messages from the cluster and processes them
+ using the user supplied callable function. The consumer will continue to read in batches until it reaches
+ the end of the log or reads a maximum number of messages is reached.
+
+ :param topics: A list of topics or regex patterns the consumer should subsrcribe to.
+ :param apply_function: The function that should be applied to fetched one at a time.
+ name of dag file executing the function and the function name delimited by a `.`
+ :param apply_function_batch: The function that should be applied to a batch of messages fetched. Can not
+ be used with `apply_function`. Intended for transactional workloads where an expensive task might
+ be called before or after operations on the messages are taken.
+ :param apply_function_args: Additional arguments that should be applied to the callable, defaults to None
+ :param apply_function_kwargs: Additional key word arguments that should be applied to the callable
+ , defaults to None
+ :param consumer_config: the config dictionary for the kafka client (additional information available on the
+ confluent-python-kafka documentation), defaults to None, defaults to None
+ :param commit_cadence: When consumers should commit offsets ("never", "end_of_batch","end_of_operator"), defaults to "end_of_operator";
+ if end_of_operator, the commit() is called based on the max_messsages arg. Commits are made after the operator has processed the apply_function method for the maximum messages in the operator.
+ if end_of_batch, the commit() is called based on the max_batch_size arg. Commits are made after each batch has processed by the apply_function method for all messages in the batch.
+ if never, close() is called without calling the commit() method.
+ :param max_messages: The maximum total number of messages an operator should read from Kafka,
+ defaults to None
+ :param max_batch_size: The maximum number of messages a consumer should read when polling, defaults to 1000
+ :param poll_timeout: How long the Kafka consumer should wait before determining no more messages are available,
+ defaults to 60
+ """
+
+ BLUE = "#ffefeb"
+ ui_color = BLUE
+ template_fields = (
+ "topics",
+ "apply_function",
+ "apply_function_args",
+ "apply_function_kwargs",
+ "config",
+ )
+
+ def __init__(
+ self,
+ topics: str | Sequence[str],
+ apply_function: Callable[..., Any] | str | None = None,
+ apply_function_batch: Callable[..., Any] | str | None = None,
+ apply_function_args: Sequence[Any] | None = None,
+ apply_function_kwargs: dict[Any, Any] | None = None,
+ consumer_config: dict[Any, Any] | None = None,
+ commit_cadence: str | None = "end_of_operator",
Review Comment:
```suggestion
commit_cadence: str = "end_of_operator",
```
##########
airflow/providers/apache/kafka/hooks/consumer.py:
##########
@@ -0,0 +1,60 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+from typing import Any, Sequence
+
+from confluent_kafka import Consumer
+
+from airflow.exceptions import AirflowException
+from airflow.hooks.base import BaseHook
+
+
+class KafkaConsumerHook(BaseHook):
+ """KafkaConsumerHook
+
+ A hook for creating a Kafka Consumer
+ :param config: A config dictionary to use with confluent_kafka library, defaults to None
Review Comment:
```suggestion
A hook for creating a Kafka Consumer
:param config: A config dictionary to use with confluent_kafka library, defaults to None
```
The blank line is needed between the description and parameter directives so the parameters are properly rendered in the Python API docs. This comment applies to other instances in the provider as well.
##########
airflow/providers/apache/kafka/hooks/admin_client.py:
##########
@@ -0,0 +1,74 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+from typing import Any, Sequence
+
+from confluent_kafka.admin import AdminClient, NewTopic
+
+from airflow.exceptions import AirflowException
+from airflow.hooks.base import BaseHook
+
+
+class KafkaAdminClientHook(BaseHook):
+ """KafkaAdminClientHook
+
+ A hook for interacting with the Kafka Cluster
Review Comment:
```suggestion
"""
A hook for interacting with the Kafka Cluster
```
Having the hook name at the beginning of the docstring seems redundant if reading in the Python API docs. Same comment applies to other operators/hooks/triggers/methods/functions in this provider too.
##########
airflow/providers/apache/kafka/hooks/admin_client.py:
##########
@@ -0,0 +1,74 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+from typing import Any, Sequence
+
+from confluent_kafka.admin import AdminClient, NewTopic
+
+from airflow.exceptions import AirflowException
+from airflow.hooks.base import BaseHook
+
+
+class KafkaAdminClientHook(BaseHook):
+ """KafkaAdminClientHook
+
+ A hook for interacting with the Kafka Cluster
+
+ :param config: A config dictionary to use with confluent_kafka library, defaults to None
+ """
+
+ def __init__(
+ self,
+ config: dict[Any, Any] | None = None,
+ ) -> None:
+ super().__init__()
+
+ self.config: dict[Any, Any] = config or {}
+
+ if not (self.config.get("bootstrap.servers", None)):
+ raise AirflowException("config['bootsrap.servers'] must be provided.")
Review Comment:
WDYT about breaking out `bootstrap_servers` into its own parameter since it's required configuration? Then the `config` param can include _optional_ configuration for the underlying lib.
##########
airflow/providers/apache/kafka/operators/await_message.py:
##########
@@ -0,0 +1,107 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+from typing import Any, Sequence
+
+from airflow.models import BaseOperator
+from airflow.providers.apache.kafka.triggers.await_message import AwaitMessageTrigger
+
+VALID_COMMIT_CADENCE = {"never", "end_of_batch", "end_of_operator"}
+
+
+class AwaitKafkaMessageOperator(BaseOperator):
+ """AwaitKafkaMessageOperator An Airflow operator that defers until a specific message is published to Kafka.
+
+ The operator creates a consumer that reads the Kafka log until it encounters a positive event.
+
+ The behavior of the consumer for this trigger is as follows:
+ - poll the Kafka topics for a message
+ - if no message returned, sleep
+ - process the message with provided callable and commit the message offset
+ - if callable returns any data, raise a TriggerEvent with the return data
+ - else continue to next message
+ - return event (as default xcom or specific xcom key)
+
+ :param topics: Topics (or topic regex) to use for reading from
+ :param apply_function: The functoin to apply to messages to determine if an event occurred. As a dot
+ notation string.
Review Comment:
```suggestion
:param apply_function: The functoin to apply to messages to determine if an event occurred. As a dot
notation string.
```
I forget how picky Sphinx is, but IIRC if the param description goes beyond the first line, each subsequent line needs to be indented. You can [build the docs locally with Breeze](https://github.com/apache/airflow/blob/main/BREEZE.rst#building-the-documentation) and confirm.
##########
airflow/providers/apache/kafka/hooks/producer.py:
##########
@@ -0,0 +1,55 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+from typing import Any
+
+from confluent_kafka import Producer
+
+from airflow.exceptions import AirflowException
+from airflow.hooks.base import BaseHook
+
+
+class KafkaProducerHook(BaseHook):
+ """KafkaProducerHook
+
+ A hook for creating a Kafka Producer
+ :param config: A config dictionary to use with confluent_kafka library, defaults to None
+ """
+
+ def __init__(
+ self,
+ config: dict[Any, Any] | None = None,
+ no_broker: bool = False,
+ ) -> None:
+ super().__init__()
+
+ self.config: dict[Any, Any] = config or {}
+
+ if not self.config.get("bootstrap.servers", None):
Review Comment:
Same comment here re: explicit parameter for a required configuration.
##########
airflow/providers/apache/kafka/hooks/consumer.py:
##########
@@ -0,0 +1,60 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+from typing import Any, Sequence
+
+from confluent_kafka import Consumer
+
+from airflow.exceptions import AirflowException
+from airflow.hooks.base import BaseHook
+
+
+class KafkaConsumerHook(BaseHook):
+ """KafkaConsumerHook
+
+ A hook for creating a Kafka Consumer
+ :param config: A config dictionary to use with confluent_kafka library, defaults to None
+ """
+
+
+ def __init__(
+ self,
+ topics: Sequence[str],
+ config: dict[Any, Any] | None = None,
+ ) -> None:
+ super().__init__()
+
+ self.config: dict[Any, Any] = config or {}
+ self.topics = topics
+
+ if not self.config.get("group.id", None):
+ raise AirflowException(
+ "The 'group.id' parameter must be set in the config dictionary'. Got <None>"
+ )
+
+ if not self.config.get("bootstrap.servers", None):
+ raise AirflowException("config['bootsrap.servers'] must be provided.")
+
+ def get_consumer(self) -> Consumer:
Review Comment:
WDYT about splitting this method into two:
1. `get_consumer()` which retrieves the `Consumer` object
2. `subscribe_to_topics()` which uses the `Consumer` object and subscribes to an input set of topics?
The subscribe functionality could be used atomically if folks wanted to use it and also use the `Consumer` object via `get_consumer()` for other actions which might not have an explicit hook method in this provider.
##########
airflow/providers/apache/kafka/operators/consume_from_topic.py:
##########
@@ -0,0 +1,170 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+from functools import partial
+from typing import Any, Callable, Sequence
+
+from airflow.exceptions import AirflowException
+from airflow.models import BaseOperator
+from airflow.providers.apache.kafka.hooks.consumer import KafkaConsumerHook
+from airflow.providers.apache.kafka.shared_utils import _get_callable
+
+VALID_COMMIT_CADENCE = {"never", "end_of_batch", "end_of_operator"}
+
+
+class ConsumeFromTopicOperator(BaseOperator):
+ """ConsumeFromTopicOperator An operator that consumes from Kafka a topic(s) and processing the messages.
+
+ The operator creates a Kafka consumer that reads a batch of messages from the cluster and processes them
+ using the user supplied callable function. The consumer will continue to read in batches until it reaches
+ the end of the log or reads a maximum number of messages is reached.
+
+ :param topics: A list of topics or regex patterns the consumer should subsrcribe to.
+ :param apply_function: The function that should be applied to fetched one at a time.
+ name of dag file executing the function and the function name delimited by a `.`
+ :param apply_function_batch: The function that should be applied to a batch of messages fetched. Can not
+ be used with `apply_function`. Intended for transactional workloads where an expensive task might
+ be called before or after operations on the messages are taken.
+ :param apply_function_args: Additional arguments that should be applied to the callable, defaults to None
+ :param apply_function_kwargs: Additional key word arguments that should be applied to the callable
+ , defaults to None
+ :param consumer_config: the config dictionary for the kafka client (additional information available on the
+ confluent-python-kafka documentation), defaults to None, defaults to None
+ :param commit_cadence: When consumers should commit offsets ("never", "end_of_batch","end_of_operator"), defaults to "end_of_operator";
+ if end_of_operator, the commit() is called based on the max_messsages arg. Commits are made after the operator has processed the apply_function method for the maximum messages in the operator.
+ if end_of_batch, the commit() is called based on the max_batch_size arg. Commits are made after each batch has processed by the apply_function method for all messages in the batch.
+ if never, close() is called without calling the commit() method.
+ :param max_messages: The maximum total number of messages an operator should read from Kafka,
+ defaults to None
+ :param max_batch_size: The maximum number of messages a consumer should read when polling, defaults to 1000
+ :param poll_timeout: How long the Kafka consumer should wait before determining no more messages are available,
+ defaults to 60
+ """
+
+ BLUE = "#ffefeb"
+ ui_color = BLUE
+ template_fields = (
+ "topics",
+ "apply_function",
+ "apply_function_args",
+ "apply_function_kwargs",
+ "config",
+ )
+
+ def __init__(
+ self,
+ topics: str | Sequence[str],
+ apply_function: Callable[..., Any] | str | None = None,
+ apply_function_batch: Callable[..., Any] | str | None = None,
+ apply_function_args: Sequence[Any] | None = None,
+ apply_function_kwargs: dict[Any, Any] | None = None,
+ consumer_config: dict[Any, Any] | None = None,
+ commit_cadence: str | None = "end_of_operator",
+ max_messages: int | None = None,
+ max_batch_size: int = 1000,
+ poll_timeout: float | None = 60,
+ **kwargs: Any,
+ ) -> None:
+
+ super().__init__(**kwargs)
+
+ self.topics = topics
+ self.apply_function = apply_function or None
+ self.apply_function_batch = apply_function_batch or None
+ self.apply_function_args = apply_function_args or ()
+ self.apply_function_kwargs = apply_function_kwargs or {}
+ self.config = consumer_config or {}
+ self.commit_cadence = commit_cadence
+ self.max_messages = max_messages or True
+ self.max_batch_size = max_batch_size
+ self.poll_timeout = poll_timeout
+
+ if self.commit_cadence not in VALID_COMMIT_CADENCE:
+ raise AirflowException(
Review Comment:
WDYT about raising a `ValueError` instead?
##########
airflow/providers/apache/kafka/operators/consume_from_topic.py:
##########
@@ -0,0 +1,170 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+from functools import partial
+from typing import Any, Callable, Sequence
+
+from airflow.exceptions import AirflowException
+from airflow.models import BaseOperator
+from airflow.providers.apache.kafka.hooks.consumer import KafkaConsumerHook
+from airflow.providers.apache.kafka.shared_utils import _get_callable
+
+VALID_COMMIT_CADENCE = {"never", "end_of_batch", "end_of_operator"}
+
+
+class ConsumeFromTopicOperator(BaseOperator):
+ """ConsumeFromTopicOperator An operator that consumes from Kafka a topic(s) and processing the messages.
+
+ The operator creates a Kafka consumer that reads a batch of messages from the cluster and processes them
+ using the user supplied callable function. The consumer will continue to read in batches until it reaches
+ the end of the log or reads a maximum number of messages is reached.
+
+ :param topics: A list of topics or regex patterns the consumer should subsrcribe to.
+ :param apply_function: The function that should be applied to fetched one at a time.
+ name of dag file executing the function and the function name delimited by a `.`
+ :param apply_function_batch: The function that should be applied to a batch of messages fetched. Can not
+ be used with `apply_function`. Intended for transactional workloads where an expensive task might
+ be called before or after operations on the messages are taken.
+ :param apply_function_args: Additional arguments that should be applied to the callable, defaults to None
+ :param apply_function_kwargs: Additional key word arguments that should be applied to the callable
+ , defaults to None
+ :param consumer_config: the config dictionary for the kafka client (additional information available on the
+ confluent-python-kafka documentation), defaults to None, defaults to None
+ :param commit_cadence: When consumers should commit offsets ("never", "end_of_batch","end_of_operator"), defaults to "end_of_operator";
+ if end_of_operator, the commit() is called based on the max_messsages arg. Commits are made after the operator has processed the apply_function method for the maximum messages in the operator.
+ if end_of_batch, the commit() is called based on the max_batch_size arg. Commits are made after each batch has processed by the apply_function method for all messages in the batch.
+ if never, close() is called without calling the commit() method.
+ :param max_messages: The maximum total number of messages an operator should read from Kafka,
+ defaults to None
+ :param max_batch_size: The maximum number of messages a consumer should read when polling, defaults to 1000
+ :param poll_timeout: How long the Kafka consumer should wait before determining no more messages are available,
+ defaults to 60
+ """
+
+ BLUE = "#ffefeb"
+ ui_color = BLUE
+ template_fields = (
+ "topics",
+ "apply_function",
+ "apply_function_args",
+ "apply_function_kwargs",
+ "config",
+ )
+
+ def __init__(
+ self,
+ topics: str | Sequence[str],
+ apply_function: Callable[..., Any] | str | None = None,
+ apply_function_batch: Callable[..., Any] | str | None = None,
+ apply_function_args: Sequence[Any] | None = None,
+ apply_function_kwargs: dict[Any, Any] | None = None,
+ consumer_config: dict[Any, Any] | None = None,
+ commit_cadence: str | None = "end_of_operator",
+ max_messages: int | None = None,
+ max_batch_size: int = 1000,
+ poll_timeout: float | None = 60,
+ **kwargs: Any,
+ ) -> None:
+
+ super().__init__(**kwargs)
+
+ self.topics = topics
+ self.apply_function = apply_function or None
+ self.apply_function_batch = apply_function_batch or None
+ self.apply_function_args = apply_function_args or ()
+ self.apply_function_kwargs = apply_function_kwargs or {}
+ self.config = consumer_config or {}
+ self.commit_cadence = commit_cadence
+ self.max_messages = max_messages or True
Review Comment:
Why does `max_messages` default to `True` if int types are the expected arguments?
Maybe there should be an explicit parameter which is the equivalent of `max_messages==True` like `fetch_all_messages`? Gives users the option to either specify a maximum number of messages to consume or consume all possible messages from the topic.
##########
airflow/providers/apache/kafka/operators/consume_from_topic.py:
##########
@@ -0,0 +1,170 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+from functools import partial
+from typing import Any, Callable, Sequence
+
+from airflow.exceptions import AirflowException
+from airflow.models import BaseOperator
+from airflow.providers.apache.kafka.hooks.consumer import KafkaConsumerHook
+from airflow.providers.apache.kafka.shared_utils import _get_callable
+
+VALID_COMMIT_CADENCE = {"never", "end_of_batch", "end_of_operator"}
+
+
+class ConsumeFromTopicOperator(BaseOperator):
+ """ConsumeFromTopicOperator An operator that consumes from Kafka a topic(s) and processing the messages.
+
+ The operator creates a Kafka consumer that reads a batch of messages from the cluster and processes them
+ using the user supplied callable function. The consumer will continue to read in batches until it reaches
+ the end of the log or reads a maximum number of messages is reached.
+
+ :param topics: A list of topics or regex patterns the consumer should subsrcribe to.
+ :param apply_function: The function that should be applied to fetched one at a time.
+ name of dag file executing the function and the function name delimited by a `.`
+ :param apply_function_batch: The function that should be applied to a batch of messages fetched. Can not
+ be used with `apply_function`. Intended for transactional workloads where an expensive task might
+ be called before or after operations on the messages are taken.
+ :param apply_function_args: Additional arguments that should be applied to the callable, defaults to None
+ :param apply_function_kwargs: Additional key word arguments that should be applied to the callable
+ , defaults to None
+ :param consumer_config: the config dictionary for the kafka client (additional information available on the
+ confluent-python-kafka documentation), defaults to None, defaults to None
+ :param commit_cadence: When consumers should commit offsets ("never", "end_of_batch","end_of_operator"), defaults to "end_of_operator";
+ if end_of_operator, the commit() is called based on the max_messsages arg. Commits are made after the operator has processed the apply_function method for the maximum messages in the operator.
+ if end_of_batch, the commit() is called based on the max_batch_size arg. Commits are made after each batch has processed by the apply_function method for all messages in the batch.
+ if never, close() is called without calling the commit() method.
+ :param max_messages: The maximum total number of messages an operator should read from Kafka,
+ defaults to None
+ :param max_batch_size: The maximum number of messages a consumer should read when polling, defaults to 1000
+ :param poll_timeout: How long the Kafka consumer should wait before determining no more messages are available,
+ defaults to 60
+ """
+
+ BLUE = "#ffefeb"
+ ui_color = BLUE
+ template_fields = (
+ "topics",
+ "apply_function",
+ "apply_function_args",
+ "apply_function_kwargs",
+ "config",
+ )
+
+ def __init__(
+ self,
+ topics: str | Sequence[str],
+ apply_function: Callable[..., Any] | str | None = None,
+ apply_function_batch: Callable[..., Any] | str | None = None,
+ apply_function_args: Sequence[Any] | None = None,
+ apply_function_kwargs: dict[Any, Any] | None = None,
+ consumer_config: dict[Any, Any] | None = None,
+ commit_cadence: str | None = "end_of_operator",
+ max_messages: int | None = None,
+ max_batch_size: int = 1000,
+ poll_timeout: float | None = 60,
+ **kwargs: Any,
+ ) -> None:
+
+ super().__init__(**kwargs)
+
+ self.topics = topics
+ self.apply_function = apply_function or None
+ self.apply_function_batch = apply_function_batch or None
Review Comment:
```suggestion
self.apply_function = apply_function
self.apply_function_batch = apply_function_batch
```
Since `None` is the default.
##########
airflow/providers/apache/kafka/provider.yaml:
##########
Review Comment:
Should use Apache Kafka rather than Kafka throughout this file.
##########
airflow/providers/apache/kafka/operators/consume_from_topic.py:
##########
@@ -0,0 +1,170 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+from functools import partial
+from typing import Any, Callable, Sequence
+
+from airflow.exceptions import AirflowException
+from airflow.models import BaseOperator
+from airflow.providers.apache.kafka.hooks.consumer import KafkaConsumerHook
+from airflow.providers.apache.kafka.shared_utils import _get_callable
+
+VALID_COMMIT_CADENCE = {"never", "end_of_batch", "end_of_operator"}
+
+
+class ConsumeFromTopicOperator(BaseOperator):
+ """ConsumeFromTopicOperator An operator that consumes from Kafka a topic(s) and processing the messages.
+
+ The operator creates a Kafka consumer that reads a batch of messages from the cluster and processes them
+ using the user supplied callable function. The consumer will continue to read in batches until it reaches
+ the end of the log or reads a maximum number of messages is reached.
+
+ :param topics: A list of topics or regex patterns the consumer should subsrcribe to.
+ :param apply_function: The function that should be applied to fetched one at a time.
+ name of dag file executing the function and the function name delimited by a `.`
+ :param apply_function_batch: The function that should be applied to a batch of messages fetched. Can not
+ be used with `apply_function`. Intended for transactional workloads where an expensive task might
+ be called before or after operations on the messages are taken.
+ :param apply_function_args: Additional arguments that should be applied to the callable, defaults to None
+ :param apply_function_kwargs: Additional key word arguments that should be applied to the callable
+ , defaults to None
+ :param consumer_config: the config dictionary for the kafka client (additional information available on the
+ confluent-python-kafka documentation), defaults to None, defaults to None
+ :param commit_cadence: When consumers should commit offsets ("never", "end_of_batch","end_of_operator"), defaults to "end_of_operator";
+ if end_of_operator, the commit() is called based on the max_messsages arg. Commits are made after the operator has processed the apply_function method for the maximum messages in the operator.
+ if end_of_batch, the commit() is called based on the max_batch_size arg. Commits are made after each batch has processed by the apply_function method for all messages in the batch.
+ if never, close() is called without calling the commit() method.
+ :param max_messages: The maximum total number of messages an operator should read from Kafka,
+ defaults to None
+ :param max_batch_size: The maximum number of messages a consumer should read when polling, defaults to 1000
+ :param poll_timeout: How long the Kafka consumer should wait before determining no more messages are available,
+ defaults to 60
+ """
+
+ BLUE = "#ffefeb"
+ ui_color = BLUE
+ template_fields = (
+ "topics",
+ "apply_function",
+ "apply_function_args",
+ "apply_function_kwargs",
+ "config",
+ )
+
+ def __init__(
+ self,
+ topics: str | Sequence[str],
+ apply_function: Callable[..., Any] | str | None = None,
+ apply_function_batch: Callable[..., Any] | str | None = None,
+ apply_function_args: Sequence[Any] | None = None,
+ apply_function_kwargs: dict[Any, Any] | None = None,
+ consumer_config: dict[Any, Any] | None = None,
+ commit_cadence: str | None = "end_of_operator",
+ max_messages: int | None = None,
+ max_batch_size: int = 1000,
+ poll_timeout: float | None = 60,
+ **kwargs: Any,
+ ) -> None:
+
+ super().__init__(**kwargs)
+
+ self.topics = topics
+ self.apply_function = apply_function or None
+ self.apply_function_batch = apply_function_batch or None
+ self.apply_function_args = apply_function_args or ()
+ self.apply_function_kwargs = apply_function_kwargs or {}
+ self.config = consumer_config or {}
+ self.commit_cadence = commit_cadence
+ self.max_messages = max_messages or True
+ self.max_batch_size = max_batch_size
+ self.poll_timeout = poll_timeout
+
+ if self.commit_cadence not in VALID_COMMIT_CADENCE:
+ raise AirflowException(
+ f"commit_cadence must be one of {VALID_COMMIT_CADENCE}. Got {self.commit_cadence}"
+ )
+
+ if self.max_messages and self.max_batch_size > self.max_messages:
+ self.log.warning(
+ f"max_batch_size ({self.max_batch_size}) > max_messages"
+ f" ({self.max_messages}). Setting max_messages to"
+ f" {self.max_batch_size}"
+ )
Review Comment:
Let's avoid using f-string interpolation in logging statements and use %-formatting instead.
##########
airflow/providers/apache/kafka/provider.yaml:
##########
@@ -0,0 +1,57 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+---
+package-name: apache-airflow-providers-kafka
+name: Kafka
+
+description: |
+ `Kafka <https://kafka.apache.org/>`__
+versions:
+ - 0.3.0
+
+dependencies:
+ - apache-airflow>=2.2.0
Review Comment:
All providers now have a minimum Airflow version of 2.3.0.
```suggestion
- apache-airflow>=2.3.0
```
##########
tests/providers/apache/kafka/hooks/test_admin_client.py:
##########
@@ -0,0 +1,69 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+import logging
+import unittest
+
+import pytest
+
+from airflow import AirflowException
+
+# Import Hook
+from airflow.providers.apache.kafka.hooks.admin_client import KafkaAdminClientHook
+
+log = logging.getLogger(__name__)
+
+
+class TestSampleHook(unittest.TestCase):
Review Comment:
Net-new tests should use `pytest` instead.
##########
airflow/providers/apache/kafka/operators/produce_to_topic.py:
##########
@@ -0,0 +1,127 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+import logging
+from functools import partial
+from typing import Any, Callable, Sequence
+
+from airflow.exceptions import AirflowException
+from airflow.models import BaseOperator
+from airflow.providers.apache.kafka.hooks.producer import KafkaProducerHook
+from airflow.providers.apache.kafka.shared_utils import _get_callable
+
+local_logger = logging.getLogger("airflow")
+
+
+def acked(err, msg):
+ if err is not None:
+ local_logger.error(f"Failed to deliver message: {err}")
+ else:
+ local_logger.info(
+ f"Produced record to topic {msg.topic()} partition [{msg.partition()}] @ offset {msg.offset()}"
+ )
+
+
+class ProduceToTopicOperator(BaseOperator):
+ """ProduceToTopicOperator An operator that produces messages to a Kafka topic
+
+ Registers a producer to a kafka topic and publishes messages to the log.
+
+
+ :param topic: The topic the producer should produce to, defaults to None
+ :param producer_function: The function that generates key/value pairs as messages for production, defaults to None
+ :param producer_function_args: Additional arguments to be applied to the producer callable, defaults to None
+ :param producer_function_kwargs: Additional keyword arguments to be applied to the producer callable,
+ defaults to None
+ :param delivery_callback: The callback to apply after delivery(or failure) of a message, defaults to None
+ :param synchronous: If writes to kafka should be fully synchronous, defaults to True
+ :param kafka_config: the config dictionary for the kafka client (additional information available on the
+ confluent-python-kafka documentation), defaults to None
+ :param poll_timeout: How long of a delay should be applied when calling poll after production to kafka,
+ defaults to 0
+ :raises AirflowException: _description_
+ """
+
+ template_fields = (
+ "topic",
+ "producer_function",
+ "producer_function_args",
+ "producer_function_kwargs",
+ "kafka_config",
+ )
+
+ def __init__(
+ self,
+ topic: str,
+ producer_function: str | Callable[..., Any],
+ producer_function_args: Sequence[Any] | None = None,
+ producer_function_kwargs: dict[Any, Any] | None = None,
+ delivery_callback: str | None = None,
+ synchronous: bool | None = True,
Review Comment:
```suggestion
synchronous: bool = True,
```
##########
airflow/providers/apache/kafka/operators/consume_from_topic.py:
##########
@@ -0,0 +1,170 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+from functools import partial
+from typing import Any, Callable, Sequence
+
+from airflow.exceptions import AirflowException
+from airflow.models import BaseOperator
+from airflow.providers.apache.kafka.hooks.consumer import KafkaConsumerHook
+from airflow.providers.apache.kafka.shared_utils import _get_callable
+
+VALID_COMMIT_CADENCE = {"never", "end_of_batch", "end_of_operator"}
+
+
+class ConsumeFromTopicOperator(BaseOperator):
+ """ConsumeFromTopicOperator An operator that consumes from Kafka a topic(s) and processing the messages.
+
+ The operator creates a Kafka consumer that reads a batch of messages from the cluster and processes them
+ using the user supplied callable function. The consumer will continue to read in batches until it reaches
+ the end of the log or reads a maximum number of messages is reached.
+
+ :param topics: A list of topics or regex patterns the consumer should subsrcribe to.
+ :param apply_function: The function that should be applied to fetched one at a time.
+ name of dag file executing the function and the function name delimited by a `.`
+ :param apply_function_batch: The function that should be applied to a batch of messages fetched. Can not
+ be used with `apply_function`. Intended for transactional workloads where an expensive task might
+ be called before or after operations on the messages are taken.
+ :param apply_function_args: Additional arguments that should be applied to the callable, defaults to None
+ :param apply_function_kwargs: Additional key word arguments that should be applied to the callable
+ , defaults to None
+ :param consumer_config: the config dictionary for the kafka client (additional information available on the
+ confluent-python-kafka documentation), defaults to None, defaults to None
+ :param commit_cadence: When consumers should commit offsets ("never", "end_of_batch","end_of_operator"), defaults to "end_of_operator";
+ if end_of_operator, the commit() is called based on the max_messsages arg. Commits are made after the operator has processed the apply_function method for the maximum messages in the operator.
+ if end_of_batch, the commit() is called based on the max_batch_size arg. Commits are made after each batch has processed by the apply_function method for all messages in the batch.
+ if never, close() is called without calling the commit() method.
+ :param max_messages: The maximum total number of messages an operator should read from Kafka,
+ defaults to None
+ :param max_batch_size: The maximum number of messages a consumer should read when polling, defaults to 1000
+ :param poll_timeout: How long the Kafka consumer should wait before determining no more messages are available,
+ defaults to 60
+ """
+
+ BLUE = "#ffefeb"
+ ui_color = BLUE
+ template_fields = (
+ "topics",
+ "apply_function",
+ "apply_function_args",
+ "apply_function_kwargs",
+ "config",
+ )
+
+ def __init__(
+ self,
+ topics: str | Sequence[str],
+ apply_function: Callable[..., Any] | str | None = None,
+ apply_function_batch: Callable[..., Any] | str | None = None,
+ apply_function_args: Sequence[Any] | None = None,
+ apply_function_kwargs: dict[Any, Any] | None = None,
+ consumer_config: dict[Any, Any] | None = None,
+ commit_cadence: str | None = "end_of_operator",
+ max_messages: int | None = None,
+ max_batch_size: int = 1000,
+ poll_timeout: float | None = 60,
Review Comment:
```suggestion
poll_timeout: float = 60,
```
##########
airflow/providers/apache/kafka/shared_utils.py:
##########
@@ -0,0 +1,53 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+import importlib
+from typing import Any, Callable
+
+
+def _get_callable(function_string: str) -> Callable:
Review Comment:
There is a [native `import_string` util function](https://github.com/apache/airflow/blob/4f3751aab677904f043d3c0657eb8283d93a9bbd/airflow/utils/module_loading.py#L26-L41) that might be used instead?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [airflow] josh-fell commented on pull request #30175: Add provider for Apache Kafka
Posted by "josh-fell (via GitHub)" <gi...@apache.org>.
josh-fell commented on PR #30175:
URL: https://github.com/apache/airflow/pull/30175#issuecomment-1474583598
How does a user define an Airflow Connection to use these hooks?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [airflow] dylanbstorey commented on pull request #30175: Add provider for Apache Kafka
Posted by "dylanbstorey (via GitHub)" <gi...@apache.org>.
dylanbstorey commented on PR #30175:
URL: https://github.com/apache/airflow/pull/30175#issuecomment-1504113380
@potiuk - integration tests are in and ready for your review. I'm struggling with the AMD a bit but i don't expect that will take me more than a few cycles once i sit down with it on and arm64 machine.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [airflow] dylanbstorey commented on a diff in pull request #30175: Add provider for Apache Kafka
Posted by "dylanbstorey (via GitHub)" <gi...@apache.org>.
dylanbstorey commented on code in PR #30175:
URL: https://github.com/apache/airflow/pull/30175#discussion_r1167626115
##########
airflow/providers/apache/kafka/hooks/produce.py:
##########
Review Comment:
same as above.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [airflow] dylanbstorey commented on a diff in pull request #30175: Add provider for Apache Kafka
Posted by "dylanbstorey (via GitHub)" <gi...@apache.org>.
dylanbstorey commented on code in PR #30175:
URL: https://github.com/apache/airflow/pull/30175#discussion_r1168017450
##########
airflow/providers/apache/kafka/operators/consume.py:
##########
@@ -0,0 +1,184 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+from functools import partial
+from typing import Any, Callable, Sequence
+
+from airflow.exceptions import AirflowException
+from airflow.models import BaseOperator
+from airflow.providers.apache.kafka.hooks.consume import KafkaConsumerHook
+from airflow.utils.module_loading import import_string
+
+VALID_COMMIT_CADENCE = {"never", "end_of_batch", "end_of_operator"}
+
+
+class ConsumeFromTopicOperator(BaseOperator):
+ """An operator that consumes from Kafka a topic(s) and processing the messages.
+
+ The operator creates a Kafka consumer that reads a batch of messages from the cluster and processes them
+ using the user supplied callable function. The consumer will continue to read in batches until it reaches
+ the end of the log or reads a maximum number of messages is reached.
+
+ :param kafka_config_id: The connection object to use, defaults to "kafka_default"
+ :param topics: A list of topics or regex patterns the consumer should subscribe to.
+ :param apply_function: The function that should be applied to fetched one at a time.
+ name of dag file executing the function and the function name delimited by a `.`
+ :param apply_function_batch: The function that should be applied to a batch of messages fetched. Can not
+ be used with `apply_function`. Intended for transactional workloads where an expensive task might
+ be called before or after operations on the messages are taken.
+ :param apply_function_args: Additional arguments that should be applied to the callable, defaults to None
+ :param apply_function_kwargs: Additional key word arguments that should be applied to the callable
+ defaults to None
+ :param commit_cadence: When consumers should commit offsets ("never", "end_of_batch","end_of_operator"),
+ defaults to "end_of_operator";
+ if end_of_operator, the commit() is called based on the max_messages arg. Commits are made after the
+ operator has processed the apply_function method for the maximum messages in the operator.
+ if end_of_batch, the commit() is called based on the max_batch_size arg. Commits are made after each
+ batch has processed by the apply_function method for all messages in the batch.
+ if never, close() is called without calling the commit() method.
+ :param max_messages: The maximum total number of messages an operator should read from Kafka,
+ defaults to None
+ :param max_batch_size: The maximum number of messages a consumer should read when polling,
+ defaults to 1000
+ :param poll_timeout: How long the Kafka consumer should wait before determining no more messages are
+ available, defaults to 60
+
+ .. seealso::
+ For more information on how to use this operator, take a look at the guide:
+ :ref:`howto/operator:ConsumeFromTopicOperator`
+ """
+
+ BLUE = "#ffefeb"
+ ui_color = BLUE
+ template_fields = (
+ "topics",
+ "apply_function",
+ "apply_function_args",
+ "apply_function_kwargs",
+ "kafka_config_id",
+ )
+
+ def __init__(
+ self,
+ topics: str | Sequence[str],
+ kafka_config_id: str = "kafka_default",
+ apply_function: Callable[..., Any] | str | None = None,
+ apply_function_batch: Callable[..., Any] | str | None = None,
+ apply_function_args: Sequence[Any] | None = None,
+ apply_function_kwargs: dict[Any, Any] | None = None,
+ commit_cadence: str | None = "end_of_operator",
+ max_messages: int | None = None,
+ max_batch_size: int = 1000,
+ poll_timeout: float = 60,
+ **kwargs: Any,
+ ) -> None:
+
+ super().__init__(**kwargs)
+
+ self.topics = topics
+ self.apply_function = apply_function
+ self.apply_function_batch = apply_function_batch
+ self.apply_function_args = apply_function_args or ()
+ self.apply_function_kwargs = apply_function_kwargs or {}
+ self.kafka_config_id = kafka_config_id
+ self.commit_cadence = commit_cadence
+ self.max_messages = max_messages or True
+ self.max_batch_size = max_batch_size
+ self.poll_timeout = poll_timeout
+
+ if self.commit_cadence not in VALID_COMMIT_CADENCE:
+ raise AirflowException(
+ f"commit_cadence must be one of {VALID_COMMIT_CADENCE}. Got {self.commit_cadence}"
+ )
+
+ if self.max_messages and self.max_batch_size > self.max_messages:
+ self.log.warning(
+ "max_batch_size (%s) > max_messages (%s). Setting max_messages to %s ",
+ self.max_batch_size,
+ self.max_messages,
+ self.max_batch_size,
+ )
+
+ if self.commit_cadence == "never":
+ self.commit_cadence = None
+
+ if apply_function and apply_function_batch:
+ raise AirflowException(
+ "One of apply_function or apply_function_batch must be supplied, not both."
+ )
+
+ def execute(self, context) -> Any:
+
+ consumer = KafkaConsumerHook(topics=self.topics, kafka_config_id=self.kafka_config_id).get_consumer()
+
+ if isinstance(self.apply_function, str):
+ self.apply_function = import_string(self.apply_function)
+
+ if isinstance(self.apply_function_batch, str):
+ self.apply_function_batch = import_string(self.apply_function_batch)
+
+ if self.apply_function:
+ apply_callable = partial(
+ self.apply_function, *self.apply_function_args, **self.apply_function_kwargs # type: ignore
+ )
+
+ if self.apply_function_batch:
+ apply_callable = partial(
+ self.apply_function_batch, # type: ignore
+ *self.apply_function_args,
+ **self.apply_function_kwargs,
+ )
+
+ messages_left = self.max_messages
+ messages_processed = 0
+
+ while (
+ messages_left > 0
+ ): # bool(True > 0) == True in the case where self.max_messages isn't set by the user
+
+ if not isinstance(messages_left, bool):
+ batch_size = self.max_batch_size if messages_left > self.max_batch_size else messages_left
+ else:
+ batch_size = self.max_batch_size
+
+ msgs = consumer.consume(num_messages=batch_size, timeout=self.poll_timeout)
+ messages_left -= len(msgs)
Review Comment:
Sentinel in - TMWYT
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [airflow] potiuk commented on pull request #30175: Add provider for Apache Kafka
Posted by "potiuk (via GitHub)" <gi...@apache.org>.
potiuk commented on PR #30175:
URL: https://github.com/apache/airflow/pull/30175#issuecomment-1509721143
Thanks @dylanbstorey for being so receptive to out expectations!
I love the way it is now nicely integrated into our testing framework, it's going to be really easy to maintain it in the future. I think for any new provider we are going to accept, we will point to this PR with "And here is how you do it" :).
@eladkal - waiting for your re-review.
@ephraimbuddy @jedcunningham @eladkal @kaxil - if we merge it before Monday. then I think we will have a chance to include it in 2.6 constraints and extras (we will need to cherry-pick at least part of it to 2.6.0 branch - easiest it will be to cherry-pick the whole of it - to include Kafka in extras of Airlfow and constraints for 2.6.0.
I also think, taking into account Kafka's popularity, it **might** be a good idea to make it installed as part of the PROD image of ours (it will slightly increase the size of our image with confluent and asgiref added as dependency) WDYT?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [airflow] dylanbstorey commented on a diff in pull request #30175: Add provider for Apache Kafka
Posted by "dylanbstorey (via GitHub)" <gi...@apache.org>.
dylanbstorey commented on code in PR #30175:
URL: https://github.com/apache/airflow/pull/30175#discussion_r1167626176
##########
airflow/providers/apache/kafka/operators/consume.py:
##########
@@ -0,0 +1,184 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+from functools import partial
+from typing import Any, Callable, Sequence
+
+from airflow.exceptions import AirflowException
+from airflow.models import BaseOperator
+from airflow.providers.apache.kafka.hooks.consume import KafkaConsumerHook
+from airflow.utils.module_loading import import_string
+
+VALID_COMMIT_CADENCE = {"never", "end_of_batch", "end_of_operator"}
Review Comment:
Yeah this was on my radar and something we can do it the future.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [airflow] dylanbstorey commented on pull request #30175: Add provider for Apache Kafka
Posted by "dylanbstorey (via GitHub)" <gi...@apache.org>.
dylanbstorey commented on PR #30175:
URL: https://github.com/apache/airflow/pull/30175#issuecomment-1514055555
@hussein-awala - already did it at your request.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [airflow] potiuk commented on pull request #30175: Add provider for Apache Kafka
Posted by "potiuk (via GitHub)" <gi...@apache.org>.
potiuk commented on PR #30175:
URL: https://github.com/apache/airflow/pull/30175#issuecomment-1517896853
And you need to rebuild the breeze images.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [airflow] boring-cyborg[bot] commented on pull request #30175: Kafka Provider
Posted by "boring-cyborg[bot] (via GitHub)" <gi...@apache.org>.
boring-cyborg[bot] commented on PR #30175:
URL: https://github.com/apache/airflow/pull/30175#issuecomment-1474441883
Congratulations on your first Pull Request and welcome to the Apache Airflow community! If you have any issues or are unsure about any anything please check our Contribution Guide (https://github.com/apache/airflow/blob/main/CONTRIBUTING.rst)
Here are some useful points:
- Pay attention to the quality of your code (ruff, mypy and type annotations). Our [pre-commits]( https://github.com/apache/airflow/blob/main/STATIC_CODE_CHECKS.rst#prerequisites-for-pre-commit-hooks) will help you with that.
- In case of a new feature add useful documentation (in docstrings or in `docs/` directory). Adding a new operator? Check this short [guide](https://github.com/apache/airflow/blob/main/docs/apache-airflow/howto/custom-operator.rst) Consider adding an example DAG that shows how users should use it.
- Consider using [Breeze environment](https://github.com/apache/airflow/blob/main/BREEZE.rst) for testing locally, it's a heavy docker but it ships with a working Airflow and a lot of integrations.
- Be patient and persistent. It might take some time to get a review or get the final approval from Committers.
- Please follow [ASF Code of Conduct](https://www.apache.org/foundation/policies/conduct) for all communication including (but not limited to) comments on Pull Requests, Mailing list and Slack.
- Be sure to read the [Airflow Coding style]( https://github.com/apache/airflow/blob/main/CONTRIBUTING.rst#coding-style-and-best-practices).
Apache Airflow is a community-driven project and together we are making it better 🚀.
In case of doubts contact the developers at:
Mailing List: dev@airflow.apache.org
Slack: https://s.apache.org/airflow-slack
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [airflow] dylanbstorey commented on pull request #30175: Add provider for Apache Kafka
Posted by "dylanbstorey (via GitHub)" <gi...@apache.org>.
dylanbstorey commented on PR #30175:
URL: https://github.com/apache/airflow/pull/30175#issuecomment-1474547096
Looks like i need to get `librdkafka` into the ARM container (not terribly surprising). I don't understand the build harness well enough to know where to start.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [airflow] dylanbstorey commented on a diff in pull request #30175: Add provider for Apache Kafka
Posted by "dylanbstorey (via GitHub)" <gi...@apache.org>.
dylanbstorey commented on code in PR #30175:
URL: https://github.com/apache/airflow/pull/30175#discussion_r1140899765
##########
airflow/providers/apache/kafka/hooks/producer.py:
##########
@@ -0,0 +1,55 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+from typing import Any
+
+from confluent_kafka import Producer
+
+from airflow.exceptions import AirflowException
+from airflow.hooks.base import BaseHook
+
+
+class KafkaProducerHook(BaseHook):
+ """KafkaProducerHook
+
+ A hook for creating a Kafka Producer
+ :param config: A config dictionary to use with confluent_kafka library, defaults to None
+ """
+
+ def __init__(
+ self,
+ config: dict[Any, Any] | None = None,
+ no_broker: bool = False,
+ ) -> None:
+ super().__init__()
+
+ self.config: dict[Any, Any] = config or {}
+
+ if not self.config.get("bootstrap.servers", None):
Review Comment:
From the original provider package
> Why is most of the configuration handled in a dict ? Because that's how confluent-kafka does it. I'd rather maintain interfaces that people already using kafka are comfortable with as a starting point - I'm happy to add more options/ interfaces in later but would prefer to be thoughtful about it to ensure that there difference between these operators and the actual client interface are minimal.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [airflow] dylanbstorey commented on a diff in pull request #30175: Add provider for Apache Kafka
Posted by "dylanbstorey (via GitHub)" <gi...@apache.org>.
dylanbstorey commented on code in PR #30175:
URL: https://github.com/apache/airflow/pull/30175#discussion_r1153161426
##########
airflow/providers/apache/kafka/hooks/base.py:
##########
@@ -0,0 +1,71 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+from typing import Any
+
+from confluent_kafka.admin import AdminClient
+
+from airflow.hooks.base import BaseHook
+
+
+class KafkaHook(BaseHook):
+ """A base hook for interacting with Apache Kafka"""
+
+ conn_name_attr = "kafka_config_id"
+ default_conn_name = "kafka_config_default"
+ conn_type = "kafka"
+ hook_name = "Apache Kafka"
+
+ def __init__(self, kafka_config_id=default_conn_name, *args, **kwargs):
+ """Initialize our Base"""
+ super().__init__()
+ self.kafka_config_id = kafka_config_id
+ self.get_conn()
Review Comment:
It's purely for validation.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [airflow] dylanbstorey commented on a diff in pull request #30175: Add provider for Apache Kafka
Posted by "dylanbstorey (via GitHub)" <gi...@apache.org>.
dylanbstorey commented on code in PR #30175:
URL: https://github.com/apache/airflow/pull/30175#discussion_r1167626112
##########
airflow/providers/apache/kafka/hooks/consume.py:
##########
Review Comment:
The consumer is cleaned up on exit automatically and `close()` means a lot more in this context with kafka so I don't think its necessary or potentially a good idea.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [airflow] potiuk commented on a diff in pull request #30175: Add provider for Apache Kafka
Posted by "potiuk (via GitHub)" <gi...@apache.org>.
potiuk commented on code in PR #30175:
URL: https://github.com/apache/airflow/pull/30175#discussion_r1167480506
##########
docs/apache-airflow-providers-apache-kafka/connections/kafka.rst:
##########
@@ -0,0 +1,38 @@
+ .. Licensed to the Apache Software Foundation (ASF) under one
+ or more contributor license agreements. See the NOTICE file
+ distributed with this work for additional information
+ regarding copyright ownership. The ASF licenses this file
+ to you under the Apache License, Version 2.0 (the
+ "License"); you may not use this file except in compliance
+ with the License. You may obtain a copy of the License at
+
+ .. http://www.apache.org/licenses/LICENSE-2.0
+
+ .. Unless required by applicable law or agreed to in writing,
+ software distributed under the License is distributed on an
+ "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ KIND, either express or implied. See the License for the
+ specific language governing permissions and limitations
+ under the License.
+
+.. _howto/connection: kafka
+
+Apache Kafka Connection
+========================
+
+The Apache Kafka connection type configures a connection to Apache Kafka via the ``confluent-kafka`` Python package.
+
+Default Connection IDs
+----------------------
+
+Kafka hooks and operators use ``kafka_default`` by default.
Review Comment:
Good one.That reminds me that is about the last thing where we will have "core" referring to particular "providers" (once we move executors to providers). One more last step to truly decouple providers from the core.
And yeah. I think we should have a default connection here.
##########
docs/apache-airflow-providers-apache-kafka/connections/kafka.rst:
##########
@@ -0,0 +1,38 @@
+ .. Licensed to the Apache Software Foundation (ASF) under one
+ or more contributor license agreements. See the NOTICE file
+ distributed with this work for additional information
+ regarding copyright ownership. The ASF licenses this file
+ to you under the Apache License, Version 2.0 (the
+ "License"); you may not use this file except in compliance
+ with the License. You may obtain a copy of the License at
+
+ .. http://www.apache.org/licenses/LICENSE-2.0
+
+ .. Unless required by applicable law or agreed to in writing,
+ software distributed under the License is distributed on an
+ "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ KIND, either express or implied. See the License for the
+ specific language governing permissions and limitations
+ under the License.
+
+.. _howto/connection: kafka
+
+Apache Kafka Connection
+========================
+
+The Apache Kafka connection type configures a connection to Apache Kafka via the ``confluent-kafka`` Python package.
+
+Default Connection IDs
+----------------------
+
+Kafka hooks and operators use ``kafka_default`` by default.
Review Comment:
Good one. That reminds me that is about the last thing where we will have "core" referring to particular "providers" (once we move executors to providers). One more last step to truly decouple providers from the core.
And yeah. I think we should have a default connection here.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [airflow] hussein-awala commented on a diff in pull request #30175: Add provider for Apache Kafka
Posted by "hussein-awala (via GitHub)" <gi...@apache.org>.
hussein-awala commented on code in PR #30175:
URL: https://github.com/apache/airflow/pull/30175#discussion_r1167634983
##########
airflow/providers/apache/kafka/operators/consume.py:
##########
@@ -0,0 +1,184 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+from functools import partial
+from typing import Any, Callable, Sequence
+
+from airflow.exceptions import AirflowException
+from airflow.models import BaseOperator
+from airflow.providers.apache.kafka.hooks.consume import KafkaConsumerHook
+from airflow.utils.module_loading import import_string
+
+VALID_COMMIT_CADENCE = {"never", "end_of_batch", "end_of_operator"}
+
+
+class ConsumeFromTopicOperator(BaseOperator):
+ """An operator that consumes from Kafka a topic(s) and processing the messages.
+
+ The operator creates a Kafka consumer that reads a batch of messages from the cluster and processes them
+ using the user supplied callable function. The consumer will continue to read in batches until it reaches
+ the end of the log or reads a maximum number of messages is reached.
+
+ :param kafka_config_id: The connection object to use, defaults to "kafka_default"
+ :param topics: A list of topics or regex patterns the consumer should subscribe to.
+ :param apply_function: The function that should be applied to fetched one at a time.
+ name of dag file executing the function and the function name delimited by a `.`
+ :param apply_function_batch: The function that should be applied to a batch of messages fetched. Can not
+ be used with `apply_function`. Intended for transactional workloads where an expensive task might
+ be called before or after operations on the messages are taken.
+ :param apply_function_args: Additional arguments that should be applied to the callable, defaults to None
+ :param apply_function_kwargs: Additional key word arguments that should be applied to the callable
+ defaults to None
+ :param commit_cadence: When consumers should commit offsets ("never", "end_of_batch","end_of_operator"),
+ defaults to "end_of_operator";
+ if end_of_operator, the commit() is called based on the max_messages arg. Commits are made after the
+ operator has processed the apply_function method for the maximum messages in the operator.
+ if end_of_batch, the commit() is called based on the max_batch_size arg. Commits are made after each
+ batch has processed by the apply_function method for all messages in the batch.
+ if never, close() is called without calling the commit() method.
+ :param max_messages: The maximum total number of messages an operator should read from Kafka,
+ defaults to None
+ :param max_batch_size: The maximum number of messages a consumer should read when polling,
+ defaults to 1000
+ :param poll_timeout: How long the Kafka consumer should wait before determining no more messages are
+ available, defaults to 60
+
+ .. seealso::
+ For more information on how to use this operator, take a look at the guide:
+ :ref:`howto/operator:ConsumeFromTopicOperator`
+ """
+
+ BLUE = "#ffefeb"
+ ui_color = BLUE
+ template_fields = (
+ "topics",
+ "apply_function",
+ "apply_function_args",
+ "apply_function_kwargs",
+ "kafka_config_id",
+ )
+
+ def __init__(
+ self,
+ topics: str | Sequence[str],
+ kafka_config_id: str = "kafka_default",
+ apply_function: Callable[..., Any] | str | None = None,
+ apply_function_batch: Callable[..., Any] | str | None = None,
+ apply_function_args: Sequence[Any] | None = None,
+ apply_function_kwargs: dict[Any, Any] | None = None,
+ commit_cadence: str | None = "end_of_operator",
+ max_messages: int | None = None,
+ max_batch_size: int = 1000,
+ poll_timeout: float = 60,
+ **kwargs: Any,
+ ) -> None:
+
+ super().__init__(**kwargs)
+
+ self.topics = topics
+ self.apply_function = apply_function
+ self.apply_function_batch = apply_function_batch
+ self.apply_function_args = apply_function_args or ()
+ self.apply_function_kwargs = apply_function_kwargs or {}
+ self.kafka_config_id = kafka_config_id
+ self.commit_cadence = commit_cadence
+ self.max_messages = max_messages or True
+ self.max_batch_size = max_batch_size
+ self.poll_timeout = poll_timeout
+
+ if self.commit_cadence not in VALID_COMMIT_CADENCE:
+ raise AirflowException(
+ f"commit_cadence must be one of {VALID_COMMIT_CADENCE}. Got {self.commit_cadence}"
+ )
+
+ if self.max_messages and self.max_batch_size > self.max_messages:
+ self.log.warning(
+ "max_batch_size (%s) > max_messages (%s). Setting max_messages to %s ",
+ self.max_batch_size,
+ self.max_messages,
+ self.max_batch_size,
+ )
+
+ if self.commit_cadence == "never":
+ self.commit_cadence = None
+
+ if apply_function and apply_function_batch:
+ raise AirflowException(
+ "One of apply_function or apply_function_batch must be supplied, not both."
+ )
+
+ def execute(self, context) -> Any:
+
+ consumer = KafkaConsumerHook(topics=self.topics, kafka_config_id=self.kafka_config_id).get_consumer()
+
+ if isinstance(self.apply_function, str):
+ self.apply_function = import_string(self.apply_function)
+
+ if isinstance(self.apply_function_batch, str):
+ self.apply_function_batch = import_string(self.apply_function_batch)
+
+ if self.apply_function:
+ apply_callable = partial(
+ self.apply_function, *self.apply_function_args, **self.apply_function_kwargs # type: ignore
+ )
+
+ if self.apply_function_batch:
+ apply_callable = partial(
+ self.apply_function_batch, # type: ignore
+ *self.apply_function_args,
+ **self.apply_function_kwargs,
+ )
+
+ messages_left = self.max_messages
+ messages_processed = 0
+
+ while (
+ messages_left > 0
+ ): # bool(True > 0) == True in the case where self.max_messages isn't set by the user
+
+ if not isinstance(messages_left, bool):
+ batch_size = self.max_batch_size if messages_left > self.max_batch_size else messages_left
+ else:
+ batch_size = self.max_batch_size
+
+ msgs = consumer.consume(num_messages=batch_size, timeout=self.poll_timeout)
+ messages_left -= len(msgs)
Review Comment:
I think my comment was not clear, I will try to rephrase it.
If `max_messages` is None, you replace it by `True`, and in this case, we should read the messages until there is no more messages.
But it is not the case now:
`messages_left = self.max_messages = True`, in the first iteration, you have ` messages_left -= len(msgs)` which is `True -= <positive number>`, and the result of this statement is a negative number because int(True) is 1, so for any len(msgs) > 0, `messages_left` will be <= 0, and the loop will break because `messages_left` <= 0.
WDYT?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [airflow] dylanbstorey commented on pull request #30175: Add provider for Apache Kafka
Posted by "dylanbstorey (via GitHub)" <gi...@apache.org>.
dylanbstorey commented on PR #30175:
URL: https://github.com/apache/airflow/pull/30175#issuecomment-1509996117
Hey man -
I feel like this is probably escalating for little/no reason. I read a few messages too fast and probably should have gotten a better understanding before I sent a terse response.
I appreciate (and value) the feedback of everyone on this so far - so please keep it coming. I'll make sure i'm a bit more rested and thoughtful in my responses (and keep less of a hair trigger on closing threads). Thanks for your grace and patience on this.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [airflow] josh-fell commented on pull request #30175: Add provider for Apache Kafka
Posted by "josh-fell (via GitHub)" <gi...@apache.org>.
josh-fell commented on PR #30175:
URL: https://github.com/apache/airflow/pull/30175#issuecomment-1518127477
Huge. Congrats Dylan 🚀
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [airflow] potiuk commented on pull request #30175: Add provider for Apache Kafka
Posted by "potiuk (via GitHub)" <gi...@apache.org>.
potiuk commented on PR #30175:
URL: https://github.com/apache/airflow/pull/30175#issuecomment-1518119085
:tada: :tada: :tada: :tada: :tada: :tada: :tada: :tada: :tada: :tada: :tada: :tada: :tada: :tada: :tada: :tada: :tada: :tada:
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [airflow] dylanbstorey commented on a diff in pull request #30175: Add provider for Apache Kafka
Posted by "dylanbstorey (via GitHub)" <gi...@apache.org>.
dylanbstorey commented on code in PR #30175:
URL: https://github.com/apache/airflow/pull/30175#discussion_r1140899172
##########
airflow/providers/apache/kafka/hooks/admin_client.py:
##########
@@ -0,0 +1,74 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+from typing import Any, Sequence
+
+from confluent_kafka.admin import AdminClient, NewTopic
+
+from airflow.exceptions import AirflowException
+from airflow.hooks.base import BaseHook
+
+
+class KafkaAdminClientHook(BaseHook):
+ """KafkaAdminClientHook
+
+ A hook for interacting with the Kafka Cluster
+
+ :param config: A config dictionary to use with confluent_kafka library, defaults to None
+ """
+
+ def __init__(
+ self,
+ config: dict[Any, Any] | None = None,
+ ) -> None:
+ super().__init__()
+
+ self.config: dict[Any, Any] = config or {}
+
+ if not (self.config.get("bootstrap.servers", None)):
+ raise AirflowException("config['bootsrap.servers'] must be provided.")
Review Comment:
The design principle is to mirror confluent-kafka config as close as possible so that there's no changes taking this to a common consumer or vice versa.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [airflow] dylanbstorey commented on a diff in pull request #30175: Add provider for Apache Kafka
Posted by "dylanbstorey (via GitHub)" <gi...@apache.org>.
dylanbstorey commented on code in PR #30175:
URL: https://github.com/apache/airflow/pull/30175#discussion_r1140899448
##########
airflow/providers/apache/kafka/hooks/consumer.py:
##########
@@ -0,0 +1,60 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+from typing import Any, Sequence
+
+from confluent_kafka import Consumer
+
+from airflow.exceptions import AirflowException
+from airflow.hooks.base import BaseHook
+
+
+class KafkaConsumerHook(BaseHook):
+ """KafkaConsumerHook
+
+ A hook for creating a Kafka Consumer
+ :param config: A config dictionary to use with confluent_kafka library, defaults to None
+ """
+
+
+ def __init__(
+ self,
+ topics: Sequence[str],
+ config: dict[Any, Any] | None = None,
+ ) -> None:
+ super().__init__()
+
+ self.config: dict[Any, Any] = config or {}
+ self.topics = topics
+
+ if not self.config.get("group.id", None):
+ raise AirflowException(
+ "The 'group.id' parameter must be set in the config dictionary'. Got <None>"
+ )
+
+ if not self.config.get("bootstrap.servers", None):
+ raise AirflowException("config['bootsrap.servers'] must be provided.")
Review Comment:
no.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [airflow] dylanbstorey commented on a diff in pull request #30175: Add provider for Apache Kafka
Posted by "dylanbstorey (via GitHub)" <gi...@apache.org>.
dylanbstorey commented on code in PR #30175:
URL: https://github.com/apache/airflow/pull/30175#discussion_r1140757323
##########
airflow/providers/kafka/CHANGELOG.md:
##########
Review Comment:
lol .... that makes too much sense.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [airflow] potiuk commented on pull request #30175: Add provider for Apache Kafka
Posted by "potiuk (via GitHub)" <gi...@apache.org>.
potiuk commented on PR #30175:
URL: https://github.com/apache/airflow/pull/30175#issuecomment-1508925267
> @potiuk - integration tests are in and ready for your review.
The tests are ready, yeah, but they should also be run as separate step in the ci.yml :) - see https://github.com/apache/airflow/blob/main/.github/workflows/ci.yml#L1150 (tests-integration-postgres/mysql ) . I **hope** it will be as easy as adding the right step there (see the others)(so you can do it while rebasing the change/solving conflicts and generating the images for breeze) .. If there will be any troubles I can have quick look (ping me on slack) or we can skip it and do a follow-up PR.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [airflow] raphaelauv commented on pull request #30175: Add provider for Apache Kafka
Posted by "raphaelauv (via GitHub)" <gi...@apache.org>.
raphaelauv commented on PR #30175:
URL: https://github.com/apache/airflow/pull/30175#issuecomment-1513755494
what do you mean by ?
```
I'm going to need more to go on to address your comment.
```
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [airflow] hussein-awala commented on a diff in pull request #30175: Add provider for Apache Kafka
Posted by "hussein-awala (via GitHub)" <gi...@apache.org>.
hussein-awala commented on code in PR #30175:
URL: https://github.com/apache/airflow/pull/30175#discussion_r1168025486
##########
airflow/providers/apache/kafka/operators/consume.py:
##########
@@ -0,0 +1,184 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+from functools import partial
+from typing import Any, Callable, Sequence
+
+from airflow.exceptions import AirflowException
+from airflow.models import BaseOperator
+from airflow.providers.apache.kafka.hooks.consume import KafkaConsumerHook
+from airflow.utils.module_loading import import_string
+
+VALID_COMMIT_CADENCE = {"never", "end_of_batch", "end_of_operator"}
+
+
+class ConsumeFromTopicOperator(BaseOperator):
+ """An operator that consumes from Kafka a topic(s) and processing the messages.
+
+ The operator creates a Kafka consumer that reads a batch of messages from the cluster and processes them
+ using the user supplied callable function. The consumer will continue to read in batches until it reaches
+ the end of the log or reads a maximum number of messages is reached.
+
+ :param kafka_config_id: The connection object to use, defaults to "kafka_default"
+ :param topics: A list of topics or regex patterns the consumer should subscribe to.
+ :param apply_function: The function that should be applied to fetched one at a time.
+ name of dag file executing the function and the function name delimited by a `.`
+ :param apply_function_batch: The function that should be applied to a batch of messages fetched. Can not
+ be used with `apply_function`. Intended for transactional workloads where an expensive task might
+ be called before or after operations on the messages are taken.
+ :param apply_function_args: Additional arguments that should be applied to the callable, defaults to None
+ :param apply_function_kwargs: Additional key word arguments that should be applied to the callable
+ defaults to None
+ :param commit_cadence: When consumers should commit offsets ("never", "end_of_batch","end_of_operator"),
+ defaults to "end_of_operator";
+ if end_of_operator, the commit() is called based on the max_messages arg. Commits are made after the
+ operator has processed the apply_function method for the maximum messages in the operator.
+ if end_of_batch, the commit() is called based on the max_batch_size arg. Commits are made after each
+ batch has processed by the apply_function method for all messages in the batch.
+ if never, close() is called without calling the commit() method.
+ :param max_messages: The maximum total number of messages an operator should read from Kafka,
+ defaults to None
+ :param max_batch_size: The maximum number of messages a consumer should read when polling,
+ defaults to 1000
+ :param poll_timeout: How long the Kafka consumer should wait before determining no more messages are
+ available, defaults to 60
+
+ .. seealso::
+ For more information on how to use this operator, take a look at the guide:
+ :ref:`howto/operator:ConsumeFromTopicOperator`
+ """
+
+ BLUE = "#ffefeb"
+ ui_color = BLUE
+ template_fields = (
+ "topics",
+ "apply_function",
+ "apply_function_args",
+ "apply_function_kwargs",
+ "kafka_config_id",
+ )
+
+ def __init__(
+ self,
+ topics: str | Sequence[str],
+ kafka_config_id: str = "kafka_default",
+ apply_function: Callable[..., Any] | str | None = None,
+ apply_function_batch: Callable[..., Any] | str | None = None,
+ apply_function_args: Sequence[Any] | None = None,
+ apply_function_kwargs: dict[Any, Any] | None = None,
+ commit_cadence: str | None = "end_of_operator",
+ max_messages: int | None = None,
+ max_batch_size: int = 1000,
+ poll_timeout: float = 60,
+ **kwargs: Any,
+ ) -> None:
+
+ super().__init__(**kwargs)
+
+ self.topics = topics
+ self.apply_function = apply_function
+ self.apply_function_batch = apply_function_batch
+ self.apply_function_args = apply_function_args or ()
+ self.apply_function_kwargs = apply_function_kwargs or {}
+ self.kafka_config_id = kafka_config_id
+ self.commit_cadence = commit_cadence
+ self.max_messages = max_messages or True
+ self.max_batch_size = max_batch_size
+ self.poll_timeout = poll_timeout
+
+ if self.commit_cadence not in VALID_COMMIT_CADENCE:
+ raise AirflowException(
+ f"commit_cadence must be one of {VALID_COMMIT_CADENCE}. Got {self.commit_cadence}"
+ )
+
+ if self.max_messages and self.max_batch_size > self.max_messages:
+ self.log.warning(
+ "max_batch_size (%s) > max_messages (%s). Setting max_messages to %s ",
+ self.max_batch_size,
+ self.max_messages,
+ self.max_batch_size,
+ )
+
+ if self.commit_cadence == "never":
+ self.commit_cadence = None
+
+ if apply_function and apply_function_batch:
+ raise AirflowException(
+ "One of apply_function or apply_function_batch must be supplied, not both."
+ )
+
+ def execute(self, context) -> Any:
+
+ consumer = KafkaConsumerHook(topics=self.topics, kafka_config_id=self.kafka_config_id).get_consumer()
+
+ if isinstance(self.apply_function, str):
+ self.apply_function = import_string(self.apply_function)
+
+ if isinstance(self.apply_function_batch, str):
+ self.apply_function_batch = import_string(self.apply_function_batch)
+
+ if self.apply_function:
+ apply_callable = partial(
+ self.apply_function, *self.apply_function_args, **self.apply_function_kwargs # type: ignore
+ )
+
+ if self.apply_function_batch:
+ apply_callable = partial(
+ self.apply_function_batch, # type: ignore
+ *self.apply_function_args,
+ **self.apply_function_kwargs,
+ )
+
+ messages_left = self.max_messages
+ messages_processed = 0
+
+ while (
+ messages_left > 0
+ ): # bool(True > 0) == True in the case where self.max_messages isn't set by the user
+
+ if not isinstance(messages_left, bool):
+ batch_size = self.max_batch_size if messages_left > self.max_batch_size else messages_left
+ else:
+ batch_size = self.max_batch_size
+
+ msgs = consumer.consume(num_messages=batch_size, timeout=self.poll_timeout)
+ messages_left -= len(msgs)
Review Comment:
looks good for me! But according to PEP naming conventions, we use uppercase names for constants, and `READ_TO_END` is an instance attribute, IMO it should be in lowercase.
##########
airflow/providers/apache/kafka/operators/consume.py:
##########
@@ -0,0 +1,189 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+from functools import partial
+from typing import Any, Callable, Sequence
+
+from airflow.exceptions import AirflowException
+from airflow.models import BaseOperator
+from airflow.providers.apache.kafka.hooks.consume import KafkaConsumerHook
+from airflow.utils.module_loading import import_string
+
+VALID_COMMIT_CADENCE = {"never", "end_of_batch", "end_of_operator"}
+
+
+class ConsumeFromTopicOperator(BaseOperator):
+ """An operator that consumes from Kafka a topic(s) and processing the messages.
+
+ The operator creates a Kafka consumer that reads a batch of messages from the cluster and processes them
+ using the user supplied callable function. The consumer will continue to read in batches until it reaches
+ the end of the log or reads a maximum number of messages is reached.
+
+ :param kafka_config_id: The connection object to use, defaults to "kafka_default"
+ :param topics: A list of topics or regex patterns the consumer should subscribe to.
+ :param apply_function: The function that should be applied to fetched one at a time.
+ name of dag file executing the function and the function name delimited by a `.`
+ :param apply_function_batch: The function that should be applied to a batch of messages fetched. Can not
+ be used with `apply_function`. Intended for transactional workloads where an expensive task might
+ be called before or after operations on the messages are taken.
+ :param apply_function_args: Additional arguments that should be applied to the callable, defaults to None
+ :param apply_function_kwargs: Additional key word arguments that should be applied to the callable
+ defaults to None
+ :param commit_cadence: When consumers should commit offsets ("never", "end_of_batch","end_of_operator"),
+ defaults to "end_of_operator";
+ if end_of_operator, the commit() is called based on the max_messages arg. Commits are made after the
+ operator has processed the apply_function method for the maximum messages in the operator.
+ if end_of_batch, the commit() is called based on the max_batch_size arg. Commits are made after each
+ batch has processed by the apply_function method for all messages in the batch.
+ if never, close() is called without calling the commit() method.
+ :param max_messages: The maximum total number of messages an operator should read from Kafka,
+ defaults to None implying read to the end of the topic.
+ :param max_batch_size: The maximum number of messages a consumer should read when polling,
+ defaults to 1000
+ :param poll_timeout: How long the Kafka consumer should wait before determining no more messages are
+ available, defaults to 60
+
+ .. seealso::
+ For more information on how to use this operator, take a look at the guide:
+ :ref:`howto/operator:ConsumeFromTopicOperator`
+ """
+
+ BLUE = "#ffefeb"
+ ui_color = BLUE
+ template_fields = (
+ "topics",
+ "apply_function",
+ "apply_function_args",
+ "apply_function_kwargs",
+ "kafka_config_id",
+ )
+
+ def __init__(
+ self,
+ topics: str | Sequence[str],
+ kafka_config_id: str = "kafka_default",
+ apply_function: Callable[..., Any] | str | None = None,
+ apply_function_batch: Callable[..., Any] | str | None = None,
+ apply_function_args: Sequence[Any] | None = None,
+ apply_function_kwargs: dict[Any, Any] | None = None,
+ commit_cadence: str | None = "end_of_operator",
+ max_messages: int | None = None,
+ max_batch_size: int = 1000,
+ poll_timeout: float = 60,
+ **kwargs: Any,
+ ) -> None:
+
+ super().__init__(**kwargs)
+
+ self.topics = topics
+ self.apply_function = apply_function
+ self.apply_function_batch = apply_function_batch
+ self.apply_function_args = apply_function_args or ()
+ self.apply_function_kwargs = apply_function_kwargs or {}
+ self.kafka_config_id = kafka_config_id
+ self.commit_cadence = commit_cadence
+ self.max_messages = max_messages or True
+ self.max_batch_size = max_batch_size
+ self.poll_timeout = poll_timeout
+
+ if self.max_messages is True:
+ self.READ_TO_END = True
+ else:
+ self.READ_TO_END = False
+
+ if self.commit_cadence not in VALID_COMMIT_CADENCE:
+ raise AirflowException(
+ f"commit_cadence must be one of {VALID_COMMIT_CADENCE}. Got {self.commit_cadence}"
+ )
+
+ if self.max_messages and self.max_batch_size > self.max_messages:
+ self.log.warning(
+ "max_batch_size (%s) > max_messages (%s). Setting max_messages to %s ",
+ self.max_batch_size,
+ self.max_messages,
+ self.max_batch_size,
+ )
+
+ if self.commit_cadence == "never":
+ self.commit_cadence = None
+
+ if apply_function and apply_function_batch:
+ raise AirflowException(
+ "One of apply_function or apply_function_batch must be supplied, not both."
+ )
+
+ def execute(self, context) -> Any:
+
+ consumer = KafkaConsumerHook(topics=self.topics, kafka_config_id=self.kafka_config_id).get_consumer()
+
+ if isinstance(self.apply_function, str):
+ self.apply_function = import_string(self.apply_function)
+
+ if isinstance(self.apply_function_batch, str):
+ self.apply_function_batch = import_string(self.apply_function_batch)
+
+ if self.apply_function:
+ apply_callable = partial(
+ self.apply_function, *self.apply_function_args, **self.apply_function_kwargs # type: ignore
+ )
+
+ if self.apply_function_batch:
+ apply_callable = partial(
+ self.apply_function_batch, # type: ignore
+ *self.apply_function_args,
+ **self.apply_function_kwargs,
+ )
+
+ messages_left = self.max_messages
+ messages_processed = 0
+
+ while self.READ_TO_END or (
+ messages_left > 0
+ ): # bool(True > 0) == True in the case where self.max_messages isn't set by the user
+
+ if not isinstance(messages_left, bool):
+ batch_size = self.max_batch_size if messages_left > self.max_batch_size else messages_left
+ else:
+ batch_size = self.max_batch_size
+
+ msgs = consumer.consume(num_messages=batch_size, timeout=self.poll_timeout)
+ messages_left -= len(msgs)
+ messages_processed += len(msgs)
Review Comment:
I think `messages_processed` is not used too
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [airflow] dylanbstorey commented on a diff in pull request #30175: Add provider for Apache Kafka
Posted by "dylanbstorey (via GitHub)" <gi...@apache.org>.
dylanbstorey commented on code in PR #30175:
URL: https://github.com/apache/airflow/pull/30175#discussion_r1167626661
##########
airflow/providers/apache/kafka/operators/consume.py:
##########
@@ -0,0 +1,184 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+from functools import partial
+from typing import Any, Callable, Sequence
+
+from airflow.exceptions import AirflowException
+from airflow.models import BaseOperator
+from airflow.providers.apache.kafka.hooks.consume import KafkaConsumerHook
+from airflow.utils.module_loading import import_string
+
+VALID_COMMIT_CADENCE = {"never", "end_of_batch", "end_of_operator"}
+
+
+class ConsumeFromTopicOperator(BaseOperator):
+ """An operator that consumes from Kafka a topic(s) and processing the messages.
+
+ The operator creates a Kafka consumer that reads a batch of messages from the cluster and processes them
+ using the user supplied callable function. The consumer will continue to read in batches until it reaches
+ the end of the log or reads a maximum number of messages is reached.
+
+ :param kafka_config_id: The connection object to use, defaults to "kafka_default"
+ :param topics: A list of topics or regex patterns the consumer should subscribe to.
+ :param apply_function: The function that should be applied to fetched one at a time.
+ name of dag file executing the function and the function name delimited by a `.`
+ :param apply_function_batch: The function that should be applied to a batch of messages fetched. Can not
+ be used with `apply_function`. Intended for transactional workloads where an expensive task might
+ be called before or after operations on the messages are taken.
+ :param apply_function_args: Additional arguments that should be applied to the callable, defaults to None
+ :param apply_function_kwargs: Additional key word arguments that should be applied to the callable
+ defaults to None
+ :param commit_cadence: When consumers should commit offsets ("never", "end_of_batch","end_of_operator"),
+ defaults to "end_of_operator";
+ if end_of_operator, the commit() is called based on the max_messages arg. Commits are made after the
+ operator has processed the apply_function method for the maximum messages in the operator.
+ if end_of_batch, the commit() is called based on the max_batch_size arg. Commits are made after each
+ batch has processed by the apply_function method for all messages in the batch.
+ if never, close() is called without calling the commit() method.
+ :param max_messages: The maximum total number of messages an operator should read from Kafka,
+ defaults to None
+ :param max_batch_size: The maximum number of messages a consumer should read when polling,
+ defaults to 1000
+ :param poll_timeout: How long the Kafka consumer should wait before determining no more messages are
+ available, defaults to 60
+
+ .. seealso::
+ For more information on how to use this operator, take a look at the guide:
+ :ref:`howto/operator:ConsumeFromTopicOperator`
+ """
+
+ BLUE = "#ffefeb"
+ ui_color = BLUE
+ template_fields = (
+ "topics",
+ "apply_function",
+ "apply_function_args",
+ "apply_function_kwargs",
+ "kafka_config_id",
+ )
+
+ def __init__(
+ self,
+ topics: str | Sequence[str],
+ kafka_config_id: str = "kafka_default",
+ apply_function: Callable[..., Any] | str | None = None,
+ apply_function_batch: Callable[..., Any] | str | None = None,
+ apply_function_args: Sequence[Any] | None = None,
+ apply_function_kwargs: dict[Any, Any] | None = None,
+ commit_cadence: str | None = "end_of_operator",
+ max_messages: int | None = None,
+ max_batch_size: int = 1000,
+ poll_timeout: float = 60,
+ **kwargs: Any,
+ ) -> None:
+
+ super().__init__(**kwargs)
+
+ self.topics = topics
+ self.apply_function = apply_function
+ self.apply_function_batch = apply_function_batch
+ self.apply_function_args = apply_function_args or ()
+ self.apply_function_kwargs = apply_function_kwargs or {}
+ self.kafka_config_id = kafka_config_id
+ self.commit_cadence = commit_cadence
+ self.max_messages = max_messages or True
+ self.max_batch_size = max_batch_size
+ self.poll_timeout = poll_timeout
+
+ if self.commit_cadence not in VALID_COMMIT_CADENCE:
+ raise AirflowException(
+ f"commit_cadence must be one of {VALID_COMMIT_CADENCE}. Got {self.commit_cadence}"
+ )
+
+ if self.max_messages and self.max_batch_size > self.max_messages:
+ self.log.warning(
+ "max_batch_size (%s) > max_messages (%s). Setting max_messages to %s ",
+ self.max_batch_size,
+ self.max_messages,
+ self.max_batch_size,
+ )
+
+ if self.commit_cadence == "never":
+ self.commit_cadence = None
+
+ if apply_function and apply_function_batch:
+ raise AirflowException(
+ "One of apply_function or apply_function_batch must be supplied, not both."
+ )
+
+ def execute(self, context) -> Any:
+
+ consumer = KafkaConsumerHook(topics=self.topics, kafka_config_id=self.kafka_config_id).get_consumer()
+
+ if isinstance(self.apply_function, str):
+ self.apply_function = import_string(self.apply_function)
+
+ if isinstance(self.apply_function_batch, str):
+ self.apply_function_batch = import_string(self.apply_function_batch)
+
+ if self.apply_function:
+ apply_callable = partial(
+ self.apply_function, *self.apply_function_args, **self.apply_function_kwargs # type: ignore
+ )
+
+ if self.apply_function_batch:
+ apply_callable = partial(
+ self.apply_function_batch, # type: ignore
+ *self.apply_function_args,
+ **self.apply_function_kwargs,
+ )
+
+ messages_left = self.max_messages
+ messages_processed = 0
+
+ while (
+ messages_left > 0
+ ): # bool(True > 0) == True in the case where self.max_messages isn't set by the user
+
+ if not isinstance(messages_left, bool):
+ batch_size = self.max_batch_size if messages_left > self.max_batch_size else messages_left
+ else:
+ batch_size = self.max_batch_size
+
+ msgs = consumer.consume(num_messages=batch_size, timeout=self.poll_timeout)
+ messages_left -= len(msgs)
Review Comment:
https://github.com/dylanbstorey/airflow/blob/92d93e9be7371efe4203005cd30981de1a2243d8/airflow/providers/apache/kafka/operators/consume.py#L152
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [airflow] potiuk merged pull request #30175: Add provider for Apache Kafka
Posted by "potiuk (via GitHub)" <gi...@apache.org>.
potiuk merged PR #30175:
URL: https://github.com/apache/airflow/pull/30175
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [airflow] potiuk commented on pull request #30175: Add provider for Apache Kafka
Posted by "potiuk (via GitHub)" <gi...@apache.org>.
potiuk commented on PR #30175:
URL: https://github.com/apache/airflow/pull/30175#issuecomment-1489923632
FYI: I just started discussion on our devlist about formalising policies about approving new providers https://lists.apache.org/thread/xf4wn1hw08cs725fp962252y7mmptk3b
My proposal is to be pretty receptive to providers like Kafka (I think we need it). One question there which I have a preference but I will see what others think is to expect from providers like Kafka to also have integration tests implemented (we already have them for mongo, pinot. trino, cassandra, celery and kerberos - so it seems appropriate to also have them for Kafka).
@dylanbstorey -> what do you think about adding integration tests for the Kafka provider.
In short:
* they are described here: https://github.com/apache/airflow/blob/main/TESTING.rst#airflow-integration-tests
* basically what the tests are about is that you use a docker image with the "software" being tested (kafka in this case) and integrate with Breezes "docker compose" framework in the way to allow breeze to start such an image via `--integration` switch
* they have a dedicated "integration" (pytest driven) test cases that instead of mocking the software in question perform a few basic operations with the "actually running" software
* they have a separate job in our CI - for example https://github.com/apache/airflow/actions/runs/4561458457/jobs/8047548089 - where the integrations are started (via breeze) and the tests are run
The nice thing about it is that it also makes it easier for the author to be sure that things work as expected - with real services, and when there are regressions, it makes it much easier for anyone to start such integration and fix problems, without having to set up everything from scratch (also it allows to test if the integration works with newever version of Kafka for example - usually by just changing a version of the integration docker image used and making a PR).
Ultimately, this is a way also for the author to be able to perform a hand-over of such provider to the community and make it easier for others to step-in and fix problems.
I am happy to help with implementing those tests. So @dylanbstorey - WDYT? Would that be too much to ask?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [airflow] kaxil commented on a diff in pull request #30175: Add provider for Apache Kafka
Posted by "kaxil (via GitHub)" <gi...@apache.org>.
kaxil commented on code in PR #30175:
URL: https://github.com/apache/airflow/pull/30175#discussion_r1153153162
##########
airflow/providers/apache/kafka/operators/consume.py:
##########
@@ -0,0 +1,182 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+from functools import partial
+from typing import Any, Callable, Sequence
+
+from airflow.exceptions import AirflowException
+from airflow.models import BaseOperator
+from airflow.providers.apache.kafka.hooks.consume import KafkaConsumerHook
+from airflow.utils.module_loading import import_string
+
+VALID_COMMIT_CADENCE = {"never", "end_of_batch", "end_of_operator"}
+
+
+class ConsumeFromTopicOperator(BaseOperator):
+ """An operator that consumes from Kafka a topic(s) and processing the messages.
+
+ The operator creates a Kafka consumer that reads a batch of messages from the cluster and processes them
+ using the user supplied callable function. The consumer will continue to read in batches until it reaches
+ the end of the log or reads a maximum number of messages is reached.
+
+ :param kafka_config_id: The connection object to use, defaults to "kafka_config_default"
+ :param topics: A list of topics or regex patterns the consumer should subscribe to.
+ :param apply_function: The function that should be applied to fetched one at a time.
+ name of dag file executing the function and the function name delimited by a `.`
+ :param apply_function_batch: The function that should be applied to a batch of messages fetched. Can not
+ be used with `apply_function`. Intended for transactional workloads where an expensive task might
+ be called before or after operations on the messages are taken.
+ :param apply_function_args: Additional arguments that should be applied to the callable, defaults to None
+ :param apply_function_kwargs: Additional key word arguments that should be applied to the callable
+ defaults to None
+ :param commit_cadence: When consumers should commit offsets ("never", "end_of_batch","end_of_operator"),
+ defaults to "end_of_operator";
+ if end_of_operator, the commit() is called based on the max_messages arg. Commits are made after the
+ operator has processed the apply_function method for the maximum messages in the operator.
+ if end_of_batch, the commit() is called based on the max_batch_size arg. Commits are made after each
+ batch has processed by the apply_function method for all messages in the batch.
+ if never, close() is called without calling the commit() method.
+ :param max_messages: The maximum total number of messages an operator should read from Kafka,
+ defaults to None
+ :param max_batch_size: The maximum number of messages a consumer should read when polling,
+ defaults to 1000
+ :param poll_timeout: How long the Kafka consumer should wait before determining no more messages are
+ available, defaults to 60
+
+
Review Comment:
```suggestion
```
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [airflow] kaxil commented on a diff in pull request #30175: Add provider for Apache Kafka
Posted by "kaxil (via GitHub)" <gi...@apache.org>.
kaxil commented on code in PR #30175:
URL: https://github.com/apache/airflow/pull/30175#discussion_r1153152863
##########
airflow/providers/apache/kafka/hooks/consume.py:
##########
@@ -0,0 +1,50 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+from typing import Sequence
+
+from confluent_kafka import Consumer
+
+from airflow.providers.apache.kafka.hooks.base import KafkaHook
+
+
+class KafkaConsumerHook(KafkaHook):
+ """
+ A hook for creating a Kafka Consumer
+
+ :param kafka_config_id: The connection object to use, defaults to "kafka_config_default"
+ :param topics: A list of topics to subscribe to.
+ """
+
+ def __init__(self, topics: Sequence[str], kafka_config_id=KafkaHook.default_conn_name) -> None:
+
+ super().__init__(kafka_config_id=kafka_config_id)
+ self.topics = topics
+ if not self.get_conn().get("group.id", None):
+ raise ValueError("The 'group.id' parameter must be set in the config dictionary'. Got <None>")
+
+ if not self.get_conn().get("bootstrap.servers", None):
Review Comment:
I wonder if we want to make `self.get_conn` a cached_property instead
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [airflow] dylanbstorey commented on a diff in pull request #30175: Add provider for Apache Kafka
Posted by "dylanbstorey (via GitHub)" <gi...@apache.org>.
dylanbstorey commented on code in PR #30175:
URL: https://github.com/apache/airflow/pull/30175#discussion_r1163370600
##########
airflow/providers/apache/kafka/example_dags/__init__.py:
##########
Review Comment:
👍
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [airflow] hussein-awala commented on a diff in pull request #30175: Add provider for Apache Kafka
Posted by "hussein-awala (via GitHub)" <gi...@apache.org>.
hussein-awala commented on code in PR #30175:
URL: https://github.com/apache/airflow/pull/30175#discussion_r1167633935
##########
airflow/providers/apache/kafka/operators/consume.py:
##########
@@ -0,0 +1,184 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+from functools import partial
+from typing import Any, Callable, Sequence
+
+from airflow.exceptions import AirflowException
+from airflow.models import BaseOperator
+from airflow.providers.apache.kafka.hooks.consume import KafkaConsumerHook
+from airflow.utils.module_loading import import_string
+
+VALID_COMMIT_CADENCE = {"never", "end_of_batch", "end_of_operator"}
Review Comment:
> As there's no such thing as a private variable in python i'm not 100% sure what you're asking.
If this variable is intended for internal use only, it's a good practice to add a single underscore before its name. This convention signals to the developers that the variable is intended for internal use and should not be accessed from outside the module. Additionally, when using the `from module import *` syntax to import all variables from the module, variables with a single underscore prefix will not be imported.
> re: Enum. Yeah this was on my radar and something we can do it the future.
okay
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [airflow] dylanbstorey commented on a diff in pull request #30175: Add provider for Apache Kafka
Posted by "dylanbstorey (via GitHub)" <gi...@apache.org>.
dylanbstorey commented on code in PR #30175:
URL: https://github.com/apache/airflow/pull/30175#discussion_r1167626974
##########
airflow/providers/apache/kafka/sensors/kafka.py:
##########
Review Comment:
I've had these patterns around since ~1.10.2 and chose not to use them even in limited settings until deferrable existed because they were catastrophically bad. Kafka and Airflow already have plenty of footguns available - I'd prefer not to hand out more w/o need.
We can revisit this conversation as KIP-500 reaches more widespread adoption and I get some time to test some how Airflow 2.0 does with this pattern though.
Also in over a year in operation not a single user has asked for this feature so I'm happy to wait until we see a demand.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [airflow] potiuk commented on pull request #30175: Add provider for Apache Kafka
Posted by "potiuk (via GitHub)" <gi...@apache.org>.
potiuk commented on PR #30175:
URL: https://github.com/apache/airflow/pull/30175#issuecomment-1509719094
FANTASTIC!
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [airflow] eladkal commented on a diff in pull request #30175: Add provider for Apache Kafka
Posted by "eladkal (via GitHub)" <gi...@apache.org>.
eladkal commented on code in PR #30175:
URL: https://github.com/apache/airflow/pull/30175#discussion_r1167478810
##########
docs/apache-airflow-providers-apache-kafka/connections/kafka.rst:
##########
@@ -0,0 +1,38 @@
+ .. Licensed to the Apache Software Foundation (ASF) under one
+ or more contributor license agreements. See the NOTICE file
+ distributed with this work for additional information
+ regarding copyright ownership. The ASF licenses this file
+ to you under the Apache License, Version 2.0 (the
+ "License"); you may not use this file except in compliance
+ with the License. You may obtain a copy of the License at
+
+ .. http://www.apache.org/licenses/LICENSE-2.0
+
+ .. Unless required by applicable law or agreed to in writing,
+ software distributed under the License is distributed on an
+ "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ KIND, either express or implied. See the License for the
+ specific language governing permissions and limitations
+ under the License.
+
+.. _howto/connection: kafka
+
+Apache Kafka Connection
+========================
+
+The Apache Kafka connection type configures a connection to Apache Kafka via the ``confluent-kafka`` Python package.
+
+Default Connection IDs
+----------------------
+
+Kafka hooks and operators use ``kafka_default`` by default.
Review Comment:
Are we missing definition for `kafka_default` in
https://github.com/apache/airflow/blob/ffa2d479e93e452faca1eb1b25c49098f2a98d63/airflow/utils/db.py#L113
or is it intentional not to have it?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [airflow] kaxil commented on a diff in pull request #30175: Add provider for Apache Kafka
Posted by "kaxil (via GitHub)" <gi...@apache.org>.
kaxil commented on code in PR #30175:
URL: https://github.com/apache/airflow/pull/30175#discussion_r1166758697
##########
airflow/providers/apache/kafka/hooks/client.py:
##########
@@ -0,0 +1,67 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+from typing import Any, Sequence
+
+from confluent_kafka.admin import AdminClient, NewTopic
+
+from airflow.providers.apache.kafka.hooks.base import KafkaHook
+
+
+class KafkaAdminClientHook(KafkaHook):
+ """
+ A hook for interacting with the Kafka Cluster
+
+ :param kafka_config_id: The connection object to use, defaults to "kafka_default"
+ """
+
+ def __init__(self, kafka_config_id=KafkaHook.default_conn_name) -> None:
+ super().__init__(kafka_config_id=kafka_config_id)
+
+ def get_admin_client(self) -> AdminClient:
+ """returns an AdminClient for communicating with the cluster
+
+ :return: an interactive admin client for the Kafka cluster
+ :rtype: AdminClient
+ """
+ client = AdminClient(self.get_conn)
+
+ self.log.info("Client %s", client)
Review Comment:
```suggestion
client = AdminClient(self.get_conn)
```
##########
airflow/providers/apache/kafka/hooks/client.py:
##########
@@ -0,0 +1,67 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+from typing import Any, Sequence
+
+from confluent_kafka.admin import AdminClient, NewTopic
+
+from airflow.providers.apache.kafka.hooks.base import KafkaHook
+
+
+class KafkaAdminClientHook(KafkaHook):
+ """
+ A hook for interacting with the Kafka Cluster
+
+ :param kafka_config_id: The connection object to use, defaults to "kafka_default"
+ """
+
+ def __init__(self, kafka_config_id=KafkaHook.default_conn_name) -> None:
+ super().__init__(kafka_config_id=kafka_config_id)
+
+ def get_admin_client(self) -> AdminClient:
+ """returns an AdminClient for communicating with the cluster
+
+ :return: an interactive admin client for the Kafka cluster
+ :rtype: AdminClient
+ """
+ client = AdminClient(self.get_conn)
+
+ self.log.info("Client %s", client)
Review Comment:
Do we want to log this?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [airflow] dylanbstorey commented on a diff in pull request #30175: Add provider for Apache Kafka
Posted by "dylanbstorey (via GitHub)" <gi...@apache.org>.
dylanbstorey commented on code in PR #30175:
URL: https://github.com/apache/airflow/pull/30175#discussion_r1167626176
##########
airflow/providers/apache/kafka/operators/consume.py:
##########
@@ -0,0 +1,184 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+from functools import partial
+from typing import Any, Callable, Sequence
+
+from airflow.exceptions import AirflowException
+from airflow.models import BaseOperator
+from airflow.providers.apache.kafka.hooks.consume import KafkaConsumerHook
+from airflow.utils.module_loading import import_string
+
+VALID_COMMIT_CADENCE = {"never", "end_of_batch", "end_of_operator"}
Review Comment:
As there's no such thing as a private variable in python i'm not 100% sure what you're asking.
re: Enum. Yeah this was on my radar and something we can do it the future.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [airflow] hussein-awala commented on pull request #30175: Add provider for Apache Kafka
Posted by "hussein-awala (via GitHub)" <gi...@apache.org>.
hussein-awala commented on PR #30175:
URL: https://github.com/apache/airflow/pull/30175#issuecomment-1509990359
> By all means but your inital comment boiled down to "i prefer this" and as you point out this is to your knowledge at least an un documented convention.
Just to clarify, my initial comment was not based on a personal preference, but rather on a convention or pattern used in other hooks. From what I understand, the `get_conn` method should be responsible for establishing a (cached) connection with the external service. Also, recent added hooks have included a `__enter__` and `__exit__` method to use the hooks with the with statement:
```python
def __enter__(self) -> SomeHook:
return self.get_conn()
def __exit__(self, exc_type, exc_val, exc_tb):
self.client.close()
```
> It also doesn't impact the external API or require a breaking change so I think "we can revisit this in the future" is an appropriate response.
Regarding the breaking changes, once the provider is released, we cannot delete the different methods that create the clients when we want to refacto the method `get_conn`. Instead, we should deprecate them and remove them in the next major version of the provider, as is done in other providers (here is an ex from KPO https://github.com/apache/airflow/pull/29498#discussion_r1167624973). So I'm just trying to avoid deprecating methods in the first releases of the provider by raising a small discussion.
The current blocker is the bug explained in https://github.com/apache/airflow/pull/30175#discussion_r1167634983 (we can discuss it if you think that there is no bug), which should be fixed before merging this PR. However, merging the PR does not mean releasing the provider, so even if we merge the PR we can still update the `get_conn` method (of course there is a convention) in a separate PR before releasing the provider, but there must be an agreement on this, hence the need for a small discussion in this PR.
Sorry to insist on this subject, but I try to apply the rules explained in this #30657:
> it might be enough to have a good and complete PR of the provider
your PR is great, I'm trying to make sure it's complete too.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [airflow] dylanbstorey commented on a diff in pull request #30175: Add provider for Apache Kafka
Posted by "dylanbstorey (via GitHub)" <gi...@apache.org>.
dylanbstorey commented on code in PR #30175:
URL: https://github.com/apache/airflow/pull/30175#discussion_r1173807830
##########
airflow/providers/apache/kafka/hooks/base.py:
##########
@@ -0,0 +1,74 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+from typing import Any
+
+from confluent_kafka.admin import AdminClient
+
+from airflow.compat.functools import cached_property
+from airflow.hooks.base import BaseHook
+
+
+class KafkaBaseHook(BaseHook):
+ """A base hook for interacting with Apache Kafka"""
+
+ conn_name_attr = "kafka_config_id"
+ default_conn_name = "kafka_default"
+ conn_type = "kafka"
+ hook_name = "Apache Kafka"
+
+ def __init__(self, kafka_config_id=default_conn_name, *args, **kwargs):
+ """Initialize our Base"""
+ super().__init__()
+ self.kafka_config_id = kafka_config_id
+ self.get_conn
+
+ @staticmethod
+ def get_ui_field_behaviour() -> dict[str, Any]:
+ """Returns custom field behaviour"""
+ return {
+ "hidden_fields": ["schema", "login", "password", "port", "host"],
+ "relabeling": {"extra": "Config Dict"},
+ "placeholders": {
+ "extra": '{"bootstrap.servers": "localhost:9092"}',
+ },
+ }
+
+ def _get_client(self, config):
+ raise NotImplementedError
+
+ @cached_property
+ def get_conn(self) -> Any:
+ """get the configuration object"""
+ config = self.get_connection(self.kafka_config_id).extra_dejson
+
+ if not (config.get("bootstrap.servers", None)):
+ raise ValueError("config['bootstrap.servers'] must be provided.")
+
+ return self._get_client(config)
+
+ def test_connection(self) -> tuple[bool, str]:
+ """Test Connectivity from the UI"""
+ try:
+ t = AdminClient(self.get_conn(), timeout=10).list_topics()
Review Comment:
sure wouldn't.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [airflow] dylanbstorey commented on pull request #30175: Add provider for Apache Kafka
Posted by "dylanbstorey (via GitHub)" <gi...@apache.org>.
dylanbstorey commented on PR #30175:
URL: https://github.com/apache/airflow/pull/30175#issuecomment-1513752757
I THINK we're done with the substantive feedback. @hussein-awala , can you review and clear as appropriate. (And ping me with anything that needs further discussion please.)
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [airflow] dylanbstorey commented on a diff in pull request #30175: Add provider for Apache Kafka
Posted by "dylanbstorey (via GitHub)" <gi...@apache.org>.
dylanbstorey commented on code in PR #30175:
URL: https://github.com/apache/airflow/pull/30175#discussion_r1168071433
##########
airflow/providers/apache/kafka/operators/consume.py:
##########
@@ -0,0 +1,189 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+from functools import partial
+from typing import Any, Callable, Sequence
+
+from airflow.exceptions import AirflowException
+from airflow.models import BaseOperator
+from airflow.providers.apache.kafka.hooks.consume import KafkaConsumerHook
+from airflow.utils.module_loading import import_string
+
+VALID_COMMIT_CADENCE = {"never", "end_of_batch", "end_of_operator"}
+
+
+class ConsumeFromTopicOperator(BaseOperator):
+ """An operator that consumes from Kafka a topic(s) and processing the messages.
+
+ The operator creates a Kafka consumer that reads a batch of messages from the cluster and processes them
+ using the user supplied callable function. The consumer will continue to read in batches until it reaches
+ the end of the log or reads a maximum number of messages is reached.
+
+ :param kafka_config_id: The connection object to use, defaults to "kafka_default"
+ :param topics: A list of topics or regex patterns the consumer should subscribe to.
+ :param apply_function: The function that should be applied to fetched one at a time.
+ name of dag file executing the function and the function name delimited by a `.`
+ :param apply_function_batch: The function that should be applied to a batch of messages fetched. Can not
+ be used with `apply_function`. Intended for transactional workloads where an expensive task might
+ be called before or after operations on the messages are taken.
+ :param apply_function_args: Additional arguments that should be applied to the callable, defaults to None
+ :param apply_function_kwargs: Additional key word arguments that should be applied to the callable
+ defaults to None
+ :param commit_cadence: When consumers should commit offsets ("never", "end_of_batch","end_of_operator"),
+ defaults to "end_of_operator";
+ if end_of_operator, the commit() is called based on the max_messages arg. Commits are made after the
+ operator has processed the apply_function method for the maximum messages in the operator.
+ if end_of_batch, the commit() is called based on the max_batch_size arg. Commits are made after each
+ batch has processed by the apply_function method for all messages in the batch.
+ if never, close() is called without calling the commit() method.
+ :param max_messages: The maximum total number of messages an operator should read from Kafka,
+ defaults to None implying read to the end of the topic.
+ :param max_batch_size: The maximum number of messages a consumer should read when polling,
+ defaults to 1000
+ :param poll_timeout: How long the Kafka consumer should wait before determining no more messages are
+ available, defaults to 60
+
+ .. seealso::
+ For more information on how to use this operator, take a look at the guide:
+ :ref:`howto/operator:ConsumeFromTopicOperator`
+ """
+
+ BLUE = "#ffefeb"
+ ui_color = BLUE
+ template_fields = (
+ "topics",
+ "apply_function",
+ "apply_function_args",
+ "apply_function_kwargs",
+ "kafka_config_id",
+ )
+
+ def __init__(
+ self,
+ topics: str | Sequence[str],
+ kafka_config_id: str = "kafka_default",
+ apply_function: Callable[..., Any] | str | None = None,
+ apply_function_batch: Callable[..., Any] | str | None = None,
+ apply_function_args: Sequence[Any] | None = None,
+ apply_function_kwargs: dict[Any, Any] | None = None,
+ commit_cadence: str | None = "end_of_operator",
+ max_messages: int | None = None,
+ max_batch_size: int = 1000,
+ poll_timeout: float = 60,
+ **kwargs: Any,
+ ) -> None:
+
+ super().__init__(**kwargs)
+
+ self.topics = topics
+ self.apply_function = apply_function
+ self.apply_function_batch = apply_function_batch
+ self.apply_function_args = apply_function_args or ()
+ self.apply_function_kwargs = apply_function_kwargs or {}
+ self.kafka_config_id = kafka_config_id
+ self.commit_cadence = commit_cadence
+ self.max_messages = max_messages or True
+ self.max_batch_size = max_batch_size
+ self.poll_timeout = poll_timeout
+
+ if self.max_messages is True:
+ self.READ_TO_END = True
+ else:
+ self.READ_TO_END = False
+
+ if self.commit_cadence not in VALID_COMMIT_CADENCE:
+ raise AirflowException(
+ f"commit_cadence must be one of {VALID_COMMIT_CADENCE}. Got {self.commit_cadence}"
+ )
+
+ if self.max_messages and self.max_batch_size > self.max_messages:
+ self.log.warning(
+ "max_batch_size (%s) > max_messages (%s). Setting max_messages to %s ",
+ self.max_batch_size,
+ self.max_messages,
+ self.max_batch_size,
+ )
+
+ if self.commit_cadence == "never":
+ self.commit_cadence = None
+
+ if apply_function and apply_function_batch:
+ raise AirflowException(
+ "One of apply_function or apply_function_batch must be supplied, not both."
+ )
+
+ def execute(self, context) -> Any:
+
+ consumer = KafkaConsumerHook(topics=self.topics, kafka_config_id=self.kafka_config_id).get_consumer()
+
+ if isinstance(self.apply_function, str):
+ self.apply_function = import_string(self.apply_function)
+
+ if isinstance(self.apply_function_batch, str):
+ self.apply_function_batch = import_string(self.apply_function_batch)
+
+ if self.apply_function:
+ apply_callable = partial(
+ self.apply_function, *self.apply_function_args, **self.apply_function_kwargs # type: ignore
+ )
+
+ if self.apply_function_batch:
+ apply_callable = partial(
+ self.apply_function_batch, # type: ignore
+ *self.apply_function_args,
+ **self.apply_function_kwargs,
+ )
+
+ messages_left = self.max_messages
+ messages_processed = 0
+
+ while self.READ_TO_END or (
+ messages_left > 0
+ ): # bool(True > 0) == True in the case where self.max_messages isn't set by the user
+
+ if not isinstance(messages_left, bool):
+ batch_size = self.max_batch_size if messages_left > self.max_batch_size else messages_left
+ else:
+ batch_size = self.max_batch_size
+
+ msgs = consumer.consume(num_messages=batch_size, timeout=self.poll_timeout)
+ messages_left -= len(msgs)
+ messages_processed += len(msgs)
Review Comment:
good catch, dead code from earlier versions
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [airflow] potiuk commented on pull request #30175: Add provider for Apache Kafka
Posted by "potiuk (via GitHub)" <gi...@apache.org>.
potiuk commented on PR #30175:
URL: https://github.com/apache/airflow/pull/30175#issuecomment-1509100487
@dylanbstorey -> comment to last commit.
There is no need to create separate JOB.
Theree are two jobs: postgres and mysql.
* The "postgres" one needs a new step for "kafka".
* The "mysql" one needs "kafka" to be added to TESTABLE_INTEGRATIONS in `src/airflow_breeze/global_constants.py`
The difference between the two is that "postgres" is run on "public" runners and runs integration by integration, where mysql runs all of them on self-hosted (big ones) .
There is no need for separate "job".
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [airflow] kaxil commented on pull request #30175: Add provider for Apache Kafka
Posted by "kaxil (via GitHub)" <gi...@apache.org>.
kaxil commented on PR #30175:
URL: https://github.com/apache/airflow/pull/30175#issuecomment-1490183294
For other reviews on why this is needed: This Provider gets ~180k monthly downloads.
https://pypistats.org/packages/airflow-provider-kafka
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [airflow] eladkal commented on a diff in pull request #30175: Add provider for Apache Kafka
Posted by "eladkal (via GitHub)" <gi...@apache.org>.
eladkal commented on code in PR #30175:
URL: https://github.com/apache/airflow/pull/30175#discussion_r1141023776
##########
airflow/providers/apache/kafka/CHANGELOG.md:
##########
Review Comment:
Please clean up.
This is change log for Airflow provider. Historical change log from the astronomer package is not something that should be here.
##########
airflow/providers/apache/kafka/provider.yaml:
##########
@@ -0,0 +1,57 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+---
+package-name: apache-airflow-providers-kafka
+name: Apache Kafka
+
+description: |
+ `Apache Kafka <https://kafka.apache.org/>`__
+versions:
+ - 0.3.0
Review Comment:
```suggestion
- 1.0.0
```
We start releases by 1.0.0
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [airflow] dylanbstorey commented on pull request #30175: Add provider for Apache Kafka
Posted by "dylanbstorey (via GitHub)" <gi...@apache.org>.
dylanbstorey commented on PR #30175:
URL: https://github.com/apache/airflow/pull/30175#issuecomment-1509428450
🤦 , yeah that makes way more sense in retrospect. It's been a long week.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [airflow] dylanbstorey commented on pull request #30175: Add provider for Apache Kafka
Posted by "dylanbstorey (via GitHub)" <gi...@apache.org>.
dylanbstorey commented on PR #30175:
URL: https://github.com/apache/airflow/pull/30175#issuecomment-1514015868
@raphaelauv , what changes are you looking for. Based on your comment - i don't understand.
> Producer default setup "could/should" be at least once and preserve order with
There is no "default producer" to setup , so I don't understand what you're suggesting I change.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [airflow] dylanbstorey commented on a diff in pull request #30175: Add provider for Apache Kafka
Posted by "dylanbstorey (via GitHub)" <gi...@apache.org>.
dylanbstorey commented on code in PR #30175:
URL: https://github.com/apache/airflow/pull/30175#discussion_r1168071268
##########
airflow/providers/apache/kafka/operators/consume.py:
##########
@@ -0,0 +1,184 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+from functools import partial
+from typing import Any, Callable, Sequence
+
+from airflow.exceptions import AirflowException
+from airflow.models import BaseOperator
+from airflow.providers.apache.kafka.hooks.consume import KafkaConsumerHook
+from airflow.utils.module_loading import import_string
+
+VALID_COMMIT_CADENCE = {"never", "end_of_batch", "end_of_operator"}
+
+
+class ConsumeFromTopicOperator(BaseOperator):
+ """An operator that consumes from Kafka a topic(s) and processing the messages.
+
+ The operator creates a Kafka consumer that reads a batch of messages from the cluster and processes them
+ using the user supplied callable function. The consumer will continue to read in batches until it reaches
+ the end of the log or reads a maximum number of messages is reached.
+
+ :param kafka_config_id: The connection object to use, defaults to "kafka_default"
+ :param topics: A list of topics or regex patterns the consumer should subscribe to.
+ :param apply_function: The function that should be applied to fetched one at a time.
+ name of dag file executing the function and the function name delimited by a `.`
+ :param apply_function_batch: The function that should be applied to a batch of messages fetched. Can not
+ be used with `apply_function`. Intended for transactional workloads where an expensive task might
+ be called before or after operations on the messages are taken.
+ :param apply_function_args: Additional arguments that should be applied to the callable, defaults to None
+ :param apply_function_kwargs: Additional key word arguments that should be applied to the callable
+ defaults to None
+ :param commit_cadence: When consumers should commit offsets ("never", "end_of_batch","end_of_operator"),
+ defaults to "end_of_operator";
+ if end_of_operator, the commit() is called based on the max_messages arg. Commits are made after the
+ operator has processed the apply_function method for the maximum messages in the operator.
+ if end_of_batch, the commit() is called based on the max_batch_size arg. Commits are made after each
+ batch has processed by the apply_function method for all messages in the batch.
+ if never, close() is called without calling the commit() method.
+ :param max_messages: The maximum total number of messages an operator should read from Kafka,
+ defaults to None
+ :param max_batch_size: The maximum number of messages a consumer should read when polling,
+ defaults to 1000
+ :param poll_timeout: How long the Kafka consumer should wait before determining no more messages are
+ available, defaults to 60
+
+ .. seealso::
+ For more information on how to use this operator, take a look at the guide:
+ :ref:`howto/operator:ConsumeFromTopicOperator`
+ """
+
+ BLUE = "#ffefeb"
+ ui_color = BLUE
+ template_fields = (
+ "topics",
+ "apply_function",
+ "apply_function_args",
+ "apply_function_kwargs",
+ "kafka_config_id",
+ )
+
+ def __init__(
+ self,
+ topics: str | Sequence[str],
+ kafka_config_id: str = "kafka_default",
+ apply_function: Callable[..., Any] | str | None = None,
+ apply_function_batch: Callable[..., Any] | str | None = None,
+ apply_function_args: Sequence[Any] | None = None,
+ apply_function_kwargs: dict[Any, Any] | None = None,
+ commit_cadence: str | None = "end_of_operator",
+ max_messages: int | None = None,
+ max_batch_size: int = 1000,
+ poll_timeout: float = 60,
+ **kwargs: Any,
+ ) -> None:
+
+ super().__init__(**kwargs)
+
+ self.topics = topics
+ self.apply_function = apply_function
+ self.apply_function_batch = apply_function_batch
+ self.apply_function_args = apply_function_args or ()
+ self.apply_function_kwargs = apply_function_kwargs or {}
+ self.kafka_config_id = kafka_config_id
+ self.commit_cadence = commit_cadence
+ self.max_messages = max_messages or True
+ self.max_batch_size = max_batch_size
+ self.poll_timeout = poll_timeout
+
+ if self.commit_cadence not in VALID_COMMIT_CADENCE:
+ raise AirflowException(
+ f"commit_cadence must be one of {VALID_COMMIT_CADENCE}. Got {self.commit_cadence}"
+ )
+
+ if self.max_messages and self.max_batch_size > self.max_messages:
+ self.log.warning(
+ "max_batch_size (%s) > max_messages (%s). Setting max_messages to %s ",
+ self.max_batch_size,
+ self.max_messages,
+ self.max_batch_size,
+ )
+
+ if self.commit_cadence == "never":
+ self.commit_cadence = None
+
+ if apply_function and apply_function_batch:
+ raise AirflowException(
+ "One of apply_function or apply_function_batch must be supplied, not both."
+ )
+
+ def execute(self, context) -> Any:
+
+ consumer = KafkaConsumerHook(topics=self.topics, kafka_config_id=self.kafka_config_id).get_consumer()
+
+ if isinstance(self.apply_function, str):
+ self.apply_function = import_string(self.apply_function)
+
+ if isinstance(self.apply_function_batch, str):
+ self.apply_function_batch = import_string(self.apply_function_batch)
+
+ if self.apply_function:
+ apply_callable = partial(
+ self.apply_function, *self.apply_function_args, **self.apply_function_kwargs # type: ignore
+ )
+
+ if self.apply_function_batch:
+ apply_callable = partial(
+ self.apply_function_batch, # type: ignore
+ *self.apply_function_args,
+ **self.apply_function_kwargs,
+ )
+
+ messages_left = self.max_messages
+ messages_processed = 0
+
+ while (
+ messages_left > 0
+ ): # bool(True > 0) == True in the case where self.max_messages isn't set by the user
+
+ if not isinstance(messages_left, bool):
+ batch_size = self.max_batch_size if messages_left > self.max_batch_size else messages_left
+ else:
+ batch_size = self.max_batch_size
+
+ msgs = consumer.consume(num_messages=batch_size, timeout=self.poll_timeout)
+ messages_left -= len(msgs)
Review Comment:
:facepalm: you're 100% correct.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [airflow] kaxil commented on a diff in pull request #30175: Kafka Provider
Posted by "kaxil (via GitHub)" <gi...@apache.org>.
kaxil commented on code in PR #30175:
URL: https://github.com/apache/airflow/pull/30175#discussion_r1140752289
##########
.github/ISSUE_TEMPLATE/airflow_providers_bug_report.yml:
##########
@@ -66,6 +66,7 @@ body:
- influxdb
- jdbc
- jenkins
+ - kafka
Review Comment:
```suggestion
- apache-kafka
```
to follow the same convention as Apache beam, spark etc
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [airflow] kaxil commented on a diff in pull request #30175: Kafka Provider
Posted by "kaxil (via GitHub)" <gi...@apache.org>.
kaxil commented on code in PR #30175:
URL: https://github.com/apache/airflow/pull/30175#discussion_r1140753329
##########
.github/ISSUE_TEMPLATE/airflow_providers_bug_report.yml:
##########
@@ -66,6 +66,7 @@ body:
- influxdb
- jdbc
- jenkins
+ - apache-kafka
Review Comment:
This will now need to go on L37 since it is alphabetically sorted
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [airflow] kaxil commented on a diff in pull request #30175: Kafka Provider
Posted by "kaxil (via GitHub)" <gi...@apache.org>.
kaxil commented on code in PR #30175:
URL: https://github.com/apache/airflow/pull/30175#discussion_r1140752880
##########
airflow/providers/kafka/CHANGELOG.md:
##########
Review Comment:
You should move this under https://github.com/apache/airflow/tree/main/airflow/providers/apache
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [airflow] dylanbstorey commented on a diff in pull request #30175: Add provider for Apache Kafka
Posted by "dylanbstorey (via GitHub)" <gi...@apache.org>.
dylanbstorey commented on code in PR #30175:
URL: https://github.com/apache/airflow/pull/30175#discussion_r1140900144
##########
airflow/providers/apache/kafka/operators/consume_from_topic.py:
##########
@@ -0,0 +1,170 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+from functools import partial
+from typing import Any, Callable, Sequence
+
+from airflow.exceptions import AirflowException
+from airflow.models import BaseOperator
+from airflow.providers.apache.kafka.hooks.consumer import KafkaConsumerHook
+from airflow.providers.apache.kafka.shared_utils import _get_callable
+
+VALID_COMMIT_CADENCE = {"never", "end_of_batch", "end_of_operator"}
+
+
+class ConsumeFromTopicOperator(BaseOperator):
+ """ConsumeFromTopicOperator An operator that consumes from Kafka a topic(s) and processing the messages.
+
+ The operator creates a Kafka consumer that reads a batch of messages from the cluster and processes them
+ using the user supplied callable function. The consumer will continue to read in batches until it reaches
+ the end of the log or reads a maximum number of messages is reached.
+
+ :param topics: A list of topics or regex patterns the consumer should subsrcribe to.
+ :param apply_function: The function that should be applied to fetched one at a time.
+ name of dag file executing the function and the function name delimited by a `.`
+ :param apply_function_batch: The function that should be applied to a batch of messages fetched. Can not
+ be used with `apply_function`. Intended for transactional workloads where an expensive task might
+ be called before or after operations on the messages are taken.
+ :param apply_function_args: Additional arguments that should be applied to the callable, defaults to None
+ :param apply_function_kwargs: Additional key word arguments that should be applied to the callable
+ , defaults to None
+ :param consumer_config: the config dictionary for the kafka client (additional information available on the
+ confluent-python-kafka documentation), defaults to None, defaults to None
+ :param commit_cadence: When consumers should commit offsets ("never", "end_of_batch","end_of_operator"), defaults to "end_of_operator";
+ if end_of_operator, the commit() is called based on the max_messsages arg. Commits are made after the operator has processed the apply_function method for the maximum messages in the operator.
+ if end_of_batch, the commit() is called based on the max_batch_size arg. Commits are made after each batch has processed by the apply_function method for all messages in the batch.
+ if never, close() is called without calling the commit() method.
+ :param max_messages: The maximum total number of messages an operator should read from Kafka,
+ defaults to None
+ :param max_batch_size: The maximum number of messages a consumer should read when polling, defaults to 1000
+ :param poll_timeout: How long the Kafka consumer should wait before determining no more messages are available,
+ defaults to 60
+ """
+
+ BLUE = "#ffefeb"
+ ui_color = BLUE
+ template_fields = (
+ "topics",
+ "apply_function",
+ "apply_function_args",
+ "apply_function_kwargs",
+ "config",
+ )
+
+ def __init__(
+ self,
+ topics: str | Sequence[str],
+ apply_function: Callable[..., Any] | str | None = None,
+ apply_function_batch: Callable[..., Any] | str | None = None,
+ apply_function_args: Sequence[Any] | None = None,
+ apply_function_kwargs: dict[Any, Any] | None = None,
+ consumer_config: dict[Any, Any] | None = None,
+ commit_cadence: str | None = "end_of_operator",
+ max_messages: int | None = None,
+ max_batch_size: int = 1000,
+ poll_timeout: float | None = 60,
+ **kwargs: Any,
+ ) -> None:
+
+ super().__init__(**kwargs)
+
+ self.topics = topics
+ self.apply_function = apply_function or None
+ self.apply_function_batch = apply_function_batch or None
+ self.apply_function_args = apply_function_args or ()
+ self.apply_function_kwargs = apply_function_kwargs or {}
+ self.config = consumer_config or {}
+ self.commit_cadence = commit_cadence
+ self.max_messages = max_messages or True
+ self.max_batch_size = max_batch_size
+ self.poll_timeout = poll_timeout
+
+ if self.commit_cadence not in VALID_COMMIT_CADENCE:
+ raise AirflowException(
Review Comment:
No strong feelings.
##########
airflow/providers/apache/kafka/operators/consume_from_topic.py:
##########
@@ -0,0 +1,170 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+from functools import partial
+from typing import Any, Callable, Sequence
+
+from airflow.exceptions import AirflowException
+from airflow.models import BaseOperator
+from airflow.providers.apache.kafka.hooks.consumer import KafkaConsumerHook
+from airflow.providers.apache.kafka.shared_utils import _get_callable
+
+VALID_COMMIT_CADENCE = {"never", "end_of_batch", "end_of_operator"}
+
+
+class ConsumeFromTopicOperator(BaseOperator):
+ """ConsumeFromTopicOperator An operator that consumes from Kafka a topic(s) and processing the messages.
+
+ The operator creates a Kafka consumer that reads a batch of messages from the cluster and processes them
+ using the user supplied callable function. The consumer will continue to read in batches until it reaches
+ the end of the log or reads a maximum number of messages is reached.
+
+ :param topics: A list of topics or regex patterns the consumer should subsrcribe to.
+ :param apply_function: The function that should be applied to fetched one at a time.
+ name of dag file executing the function and the function name delimited by a `.`
+ :param apply_function_batch: The function that should be applied to a batch of messages fetched. Can not
+ be used with `apply_function`. Intended for transactional workloads where an expensive task might
+ be called before or after operations on the messages are taken.
+ :param apply_function_args: Additional arguments that should be applied to the callable, defaults to None
+ :param apply_function_kwargs: Additional key word arguments that should be applied to the callable
+ , defaults to None
+ :param consumer_config: the config dictionary for the kafka client (additional information available on the
+ confluent-python-kafka documentation), defaults to None, defaults to None
+ :param commit_cadence: When consumers should commit offsets ("never", "end_of_batch","end_of_operator"), defaults to "end_of_operator";
+ if end_of_operator, the commit() is called based on the max_messsages arg. Commits are made after the operator has processed the apply_function method for the maximum messages in the operator.
+ if end_of_batch, the commit() is called based on the max_batch_size arg. Commits are made after each batch has processed by the apply_function method for all messages in the batch.
+ if never, close() is called without calling the commit() method.
+ :param max_messages: The maximum total number of messages an operator should read from Kafka,
+ defaults to None
+ :param max_batch_size: The maximum number of messages a consumer should read when polling, defaults to 1000
+ :param poll_timeout: How long the Kafka consumer should wait before determining no more messages are available,
+ defaults to 60
+ """
+
+ BLUE = "#ffefeb"
+ ui_color = BLUE
+ template_fields = (
+ "topics",
+ "apply_function",
+ "apply_function_args",
+ "apply_function_kwargs",
+ "config",
+ )
+
+ def __init__(
+ self,
+ topics: str | Sequence[str],
+ apply_function: Callable[..., Any] | str | None = None,
+ apply_function_batch: Callable[..., Any] | str | None = None,
+ apply_function_args: Sequence[Any] | None = None,
+ apply_function_kwargs: dict[Any, Any] | None = None,
+ consumer_config: dict[Any, Any] | None = None,
+ commit_cadence: str | None = "end_of_operator",
+ max_messages: int | None = None,
+ max_batch_size: int = 1000,
+ poll_timeout: float | None = 60,
+ **kwargs: Any,
+ ) -> None:
+
+ super().__init__(**kwargs)
+
+ self.topics = topics
+ self.apply_function = apply_function or None
+ self.apply_function_batch = apply_function_batch or None
+ self.apply_function_args = apply_function_args or ()
+ self.apply_function_kwargs = apply_function_kwargs or {}
+ self.config = consumer_config or {}
+ self.commit_cadence = commit_cadence
+ self.max_messages = max_messages or True
Review Comment:
If you go further down you'll see the while loop. When `max_messages` == True, it will read to the end of the log as a default behavior.
##########
airflow/providers/apache/kafka/provider.yaml:
##########
Review Comment:
Missed this after Kaxil's feedback earlier. Thanks !
##########
airflow/providers/apache/kafka/shared_utils.py:
##########
@@ -0,0 +1,53 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+import importlib
+from typing import Any, Callable
+
+
+def _get_callable(function_string: str) -> Callable:
Review Comment:
there is.
The doc string indicates "attr/class", while true it is pretty much the same code - i'm not sure if the intent will stay consistent in the future.
##########
airflow/providers/apache/kafka/operators/consume_from_topic.py:
##########
@@ -0,0 +1,170 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+from functools import partial
+from typing import Any, Callable, Sequence
+
+from airflow.exceptions import AirflowException
+from airflow.models import BaseOperator
+from airflow.providers.apache.kafka.hooks.consumer import KafkaConsumerHook
+from airflow.providers.apache.kafka.shared_utils import _get_callable
+
+VALID_COMMIT_CADENCE = {"never", "end_of_batch", "end_of_operator"}
+
+
+class ConsumeFromTopicOperator(BaseOperator):
+ """ConsumeFromTopicOperator An operator that consumes from Kafka a topic(s) and processing the messages.
+
+ The operator creates a Kafka consumer that reads a batch of messages from the cluster and processes them
+ using the user supplied callable function. The consumer will continue to read in batches until it reaches
+ the end of the log or reads a maximum number of messages is reached.
+
+ :param topics: A list of topics or regex patterns the consumer should subsrcribe to.
+ :param apply_function: The function that should be applied to fetched one at a time.
+ name of dag file executing the function and the function name delimited by a `.`
+ :param apply_function_batch: The function that should be applied to a batch of messages fetched. Can not
+ be used with `apply_function`. Intended for transactional workloads where an expensive task might
+ be called before or after operations on the messages are taken.
+ :param apply_function_args: Additional arguments that should be applied to the callable, defaults to None
+ :param apply_function_kwargs: Additional key word arguments that should be applied to the callable
+ , defaults to None
+ :param consumer_config: the config dictionary for the kafka client (additional information available on the
+ confluent-python-kafka documentation), defaults to None, defaults to None
+ :param commit_cadence: When consumers should commit offsets ("never", "end_of_batch","end_of_operator"), defaults to "end_of_operator";
+ if end_of_operator, the commit() is called based on the max_messsages arg. Commits are made after the operator has processed the apply_function method for the maximum messages in the operator.
+ if end_of_batch, the commit() is called based on the max_batch_size arg. Commits are made after each batch has processed by the apply_function method for all messages in the batch.
+ if never, close() is called without calling the commit() method.
+ :param max_messages: The maximum total number of messages an operator should read from Kafka,
+ defaults to None
+ :param max_batch_size: The maximum number of messages a consumer should read when polling, defaults to 1000
+ :param poll_timeout: How long the Kafka consumer should wait before determining no more messages are available,
+ defaults to 60
+ """
+
+ BLUE = "#ffefeb"
+ ui_color = BLUE
+ template_fields = (
+ "topics",
+ "apply_function",
+ "apply_function_args",
+ "apply_function_kwargs",
+ "config",
+ )
+
+ def __init__(
+ self,
+ topics: str | Sequence[str],
+ apply_function: Callable[..., Any] | str | None = None,
+ apply_function_batch: Callable[..., Any] | str | None = None,
+ apply_function_args: Sequence[Any] | None = None,
+ apply_function_kwargs: dict[Any, Any] | None = None,
+ consumer_config: dict[Any, Any] | None = None,
+ commit_cadence: str | None = "end_of_operator",
+ max_messages: int | None = None,
+ max_batch_size: int = 1000,
+ poll_timeout: float | None = 60,
+ **kwargs: Any,
+ ) -> None:
+
+ super().__init__(**kwargs)
+
+ self.topics = topics
+ self.apply_function = apply_function or None
+ self.apply_function_batch = apply_function_batch or None
+ self.apply_function_args = apply_function_args or ()
+ self.apply_function_kwargs = apply_function_kwargs or {}
+ self.config = consumer_config or {}
+ self.commit_cadence = commit_cadence
+ self.max_messages = max_messages or True
+ self.max_batch_size = max_batch_size
+ self.poll_timeout = poll_timeout
+
+ if self.commit_cadence not in VALID_COMMIT_CADENCE:
+ raise AirflowException(
+ f"commit_cadence must be one of {VALID_COMMIT_CADENCE}. Got {self.commit_cadence}"
+ )
+
+ if self.max_messages and self.max_batch_size > self.max_messages:
+ self.log.warning(
+ f"max_batch_size ({self.max_batch_size}) > max_messages"
+ f" ({self.max_messages}). Setting max_messages to"
+ f" {self.max_batch_size}"
+ )
Review Comment:
looks like the pre-commit hook didn't catch this one.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [airflow] dylanbstorey commented on a diff in pull request #30175: Add provider for Apache Kafka
Posted by "dylanbstorey (via GitHub)" <gi...@apache.org>.
dylanbstorey commented on code in PR #30175:
URL: https://github.com/apache/airflow/pull/30175#discussion_r1140899915
##########
airflow/providers/apache/kafka/operators/consume_from_topic.py:
##########
@@ -0,0 +1,170 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+from functools import partial
+from typing import Any, Callable, Sequence
+
+from airflow.exceptions import AirflowException
+from airflow.models import BaseOperator
+from airflow.providers.apache.kafka.hooks.consumer import KafkaConsumerHook
+from airflow.providers.apache.kafka.shared_utils import _get_callable
+
+VALID_COMMIT_CADENCE = {"never", "end_of_batch", "end_of_operator"}
+
+
+class ConsumeFromTopicOperator(BaseOperator):
+ """ConsumeFromTopicOperator An operator that consumes from Kafka a topic(s) and processing the messages.
+
+ The operator creates a Kafka consumer that reads a batch of messages from the cluster and processes them
+ using the user supplied callable function. The consumer will continue to read in batches until it reaches
+ the end of the log or reads a maximum number of messages is reached.
+
+ :param topics: A list of topics or regex patterns the consumer should subsrcribe to.
+ :param apply_function: The function that should be applied to fetched one at a time.
+ name of dag file executing the function and the function name delimited by a `.`
+ :param apply_function_batch: The function that should be applied to a batch of messages fetched. Can not
+ be used with `apply_function`. Intended for transactional workloads where an expensive task might
+ be called before or after operations on the messages are taken.
+ :param apply_function_args: Additional arguments that should be applied to the callable, defaults to None
+ :param apply_function_kwargs: Additional key word arguments that should be applied to the callable
+ , defaults to None
+ :param consumer_config: the config dictionary for the kafka client (additional information available on the
+ confluent-python-kafka documentation), defaults to None, defaults to None
+ :param commit_cadence: When consumers should commit offsets ("never", "end_of_batch","end_of_operator"), defaults to "end_of_operator";
+ if end_of_operator, the commit() is called based on the max_messsages arg. Commits are made after the operator has processed the apply_function method for the maximum messages in the operator.
+ if end_of_batch, the commit() is called based on the max_batch_size arg. Commits are made after each batch has processed by the apply_function method for all messages in the batch.
+ if never, close() is called without calling the commit() method.
+ :param max_messages: The maximum total number of messages an operator should read from Kafka,
+ defaults to None
+ :param max_batch_size: The maximum number of messages a consumer should read when polling, defaults to 1000
+ :param poll_timeout: How long the Kafka consumer should wait before determining no more messages are available,
+ defaults to 60
+ """
+
+ BLUE = "#ffefeb"
+ ui_color = BLUE
+ template_fields = (
+ "topics",
+ "apply_function",
+ "apply_function_args",
+ "apply_function_kwargs",
+ "config",
+ )
+
+ def __init__(
+ self,
+ topics: str | Sequence[str],
+ apply_function: Callable[..., Any] | str | None = None,
+ apply_function_batch: Callable[..., Any] | str | None = None,
+ apply_function_args: Sequence[Any] | None = None,
+ apply_function_kwargs: dict[Any, Any] | None = None,
+ consumer_config: dict[Any, Any] | None = None,
+ commit_cadence: str | None = "end_of_operator",
Review Comment:
All of these `| None` were added by pre-commit hooks.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [airflow] dylanbstorey commented on a diff in pull request #30175: Add provider for Apache Kafka
Posted by "dylanbstorey (via GitHub)" <gi...@apache.org>.
dylanbstorey commented on code in PR #30175:
URL: https://github.com/apache/airflow/pull/30175#discussion_r1140899642
##########
airflow/providers/apache/kafka/hooks/consumer.py:
##########
@@ -0,0 +1,60 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+from typing import Any, Sequence
+
+from confluent_kafka import Consumer
+
+from airflow.exceptions import AirflowException
+from airflow.hooks.base import BaseHook
+
+
+class KafkaConsumerHook(BaseHook):
+ """KafkaConsumerHook
+
+ A hook for creating a Kafka Consumer
+ :param config: A config dictionary to use with confluent_kafka library, defaults to None
+ """
+
+
+ def __init__(
+ self,
+ topics: Sequence[str],
+ config: dict[Any, Any] | None = None,
+ ) -> None:
+ super().__init__()
+
+ self.config: dict[Any, Any] = config or {}
+ self.topics = topics
+
+ if not self.config.get("group.id", None):
+ raise AirflowException(
+ "The 'group.id' parameter must be set in the config dictionary'. Got <None>"
+ )
+
+ if not self.config.get("bootstrap.servers", None):
+ raise AirflowException("config['bootsrap.servers'] must be provided.")
+
+ def get_consumer(self) -> Consumer:
Review Comment:
Consumers that aren't subscribed to topics are a no-op so I don't think there's really a reason to do that.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [airflow] dylanbstorey commented on pull request #30175: Add provider for Apache Kafka
Posted by "dylanbstorey (via GitHub)" <gi...@apache.org>.
dylanbstorey commented on PR #30175:
URL: https://github.com/apache/airflow/pull/30175#issuecomment-1490191057
@potiuk
Your note gave me a nice chuckle, I had already started down this path a bit but got "distracted" with some of the failing tests instead.
So my short answer (and i'll add it to the dev mailing list is), I think that requiring integration tests will raise the bar on entry, require additional work/support from the community, and ultimately will support a higher quality code base.
I'd glady accept your help here. Thanks.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [airflow] potiuk commented on pull request #30175: Add provider for Apache Kafka
Posted by "potiuk (via GitHub)" <gi...@apache.org>.
potiuk commented on PR #30175:
URL: https://github.com/apache/airflow/pull/30175#issuecomment-1518118638
I restarted the last failing MSSQL test (flaky) but for all practical purposes, it's green. Merging.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [airflow] dylanbstorey commented on a diff in pull request #30175: Add provider for Apache Kafka
Posted by "dylanbstorey (via GitHub)" <gi...@apache.org>.
dylanbstorey commented on code in PR #30175:
URL: https://github.com/apache/airflow/pull/30175#discussion_r1167626974
##########
airflow/providers/apache/kafka/sensors/kafka.py:
##########
Review Comment:
I've had these patterns around since ~1.10.2 and chose not to use them even in limited settings until deferrable existed because they were catastrophically bad. Kafka and Airflow already have plenty of footguns available - I'd prefer not to hand out more.
We can revisit this conversation as KIP-500 reaches more widespread adoption though.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [airflow] hussein-awala commented on a diff in pull request #30175: Add provider for Apache Kafka
Posted by "hussein-awala (via GitHub)" <gi...@apache.org>.
hussein-awala commented on code in PR #30175:
URL: https://github.com/apache/airflow/pull/30175#discussion_r1167634983
##########
airflow/providers/apache/kafka/operators/consume.py:
##########
@@ -0,0 +1,184 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+from functools import partial
+from typing import Any, Callable, Sequence
+
+from airflow.exceptions import AirflowException
+from airflow.models import BaseOperator
+from airflow.providers.apache.kafka.hooks.consume import KafkaConsumerHook
+from airflow.utils.module_loading import import_string
+
+VALID_COMMIT_CADENCE = {"never", "end_of_batch", "end_of_operator"}
+
+
+class ConsumeFromTopicOperator(BaseOperator):
+ """An operator that consumes from Kafka a topic(s) and processing the messages.
+
+ The operator creates a Kafka consumer that reads a batch of messages from the cluster and processes them
+ using the user supplied callable function. The consumer will continue to read in batches until it reaches
+ the end of the log or reads a maximum number of messages is reached.
+
+ :param kafka_config_id: The connection object to use, defaults to "kafka_default"
+ :param topics: A list of topics or regex patterns the consumer should subscribe to.
+ :param apply_function: The function that should be applied to fetched one at a time.
+ name of dag file executing the function and the function name delimited by a `.`
+ :param apply_function_batch: The function that should be applied to a batch of messages fetched. Can not
+ be used with `apply_function`. Intended for transactional workloads where an expensive task might
+ be called before or after operations on the messages are taken.
+ :param apply_function_args: Additional arguments that should be applied to the callable, defaults to None
+ :param apply_function_kwargs: Additional key word arguments that should be applied to the callable
+ defaults to None
+ :param commit_cadence: When consumers should commit offsets ("never", "end_of_batch","end_of_operator"),
+ defaults to "end_of_operator";
+ if end_of_operator, the commit() is called based on the max_messages arg. Commits are made after the
+ operator has processed the apply_function method for the maximum messages in the operator.
+ if end_of_batch, the commit() is called based on the max_batch_size arg. Commits are made after each
+ batch has processed by the apply_function method for all messages in the batch.
+ if never, close() is called without calling the commit() method.
+ :param max_messages: The maximum total number of messages an operator should read from Kafka,
+ defaults to None
+ :param max_batch_size: The maximum number of messages a consumer should read when polling,
+ defaults to 1000
+ :param poll_timeout: How long the Kafka consumer should wait before determining no more messages are
+ available, defaults to 60
+
+ .. seealso::
+ For more information on how to use this operator, take a look at the guide:
+ :ref:`howto/operator:ConsumeFromTopicOperator`
+ """
+
+ BLUE = "#ffefeb"
+ ui_color = BLUE
+ template_fields = (
+ "topics",
+ "apply_function",
+ "apply_function_args",
+ "apply_function_kwargs",
+ "kafka_config_id",
+ )
+
+ def __init__(
+ self,
+ topics: str | Sequence[str],
+ kafka_config_id: str = "kafka_default",
+ apply_function: Callable[..., Any] | str | None = None,
+ apply_function_batch: Callable[..., Any] | str | None = None,
+ apply_function_args: Sequence[Any] | None = None,
+ apply_function_kwargs: dict[Any, Any] | None = None,
+ commit_cadence: str | None = "end_of_operator",
+ max_messages: int | None = None,
+ max_batch_size: int = 1000,
+ poll_timeout: float = 60,
+ **kwargs: Any,
+ ) -> None:
+
+ super().__init__(**kwargs)
+
+ self.topics = topics
+ self.apply_function = apply_function
+ self.apply_function_batch = apply_function_batch
+ self.apply_function_args = apply_function_args or ()
+ self.apply_function_kwargs = apply_function_kwargs or {}
+ self.kafka_config_id = kafka_config_id
+ self.commit_cadence = commit_cadence
+ self.max_messages = max_messages or True
+ self.max_batch_size = max_batch_size
+ self.poll_timeout = poll_timeout
+
+ if self.commit_cadence not in VALID_COMMIT_CADENCE:
+ raise AirflowException(
+ f"commit_cadence must be one of {VALID_COMMIT_CADENCE}. Got {self.commit_cadence}"
+ )
+
+ if self.max_messages and self.max_batch_size > self.max_messages:
+ self.log.warning(
+ "max_batch_size (%s) > max_messages (%s). Setting max_messages to %s ",
+ self.max_batch_size,
+ self.max_messages,
+ self.max_batch_size,
+ )
+
+ if self.commit_cadence == "never":
+ self.commit_cadence = None
+
+ if apply_function and apply_function_batch:
+ raise AirflowException(
+ "One of apply_function or apply_function_batch must be supplied, not both."
+ )
+
+ def execute(self, context) -> Any:
+
+ consumer = KafkaConsumerHook(topics=self.topics, kafka_config_id=self.kafka_config_id).get_consumer()
+
+ if isinstance(self.apply_function, str):
+ self.apply_function = import_string(self.apply_function)
+
+ if isinstance(self.apply_function_batch, str):
+ self.apply_function_batch = import_string(self.apply_function_batch)
+
+ if self.apply_function:
+ apply_callable = partial(
+ self.apply_function, *self.apply_function_args, **self.apply_function_kwargs # type: ignore
+ )
+
+ if self.apply_function_batch:
+ apply_callable = partial(
+ self.apply_function_batch, # type: ignore
+ *self.apply_function_args,
+ **self.apply_function_kwargs,
+ )
+
+ messages_left = self.max_messages
+ messages_processed = 0
+
+ while (
+ messages_left > 0
+ ): # bool(True > 0) == True in the case where self.max_messages isn't set by the user
+
+ if not isinstance(messages_left, bool):
+ batch_size = self.max_batch_size if messages_left > self.max_batch_size else messages_left
+ else:
+ batch_size = self.max_batch_size
+
+ msgs = consumer.consume(num_messages=batch_size, timeout=self.poll_timeout)
+ messages_left -= len(msgs)
Review Comment:
I think my comment was not clear, I will try to rephrase it.
If `max_messages` is None, you replace it by `True`, and in this case, we should read the message until there is no more messages.
But it is not the case now:
`messages_left = self.max_messages = True`, in the first iteration, you have ` messages_left -= len(msgs)` which is `True -= <positive number>`, and the result of this statement is a negative number because int(True) is 1, so for any len(msgs) > 0, `messages_left` will be <= 0, and the loop will break because `messages_left` <= 0.
WDYT?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [airflow] hussein-awala commented on a diff in pull request #30175: Add provider for Apache Kafka
Posted by "hussein-awala (via GitHub)" <gi...@apache.org>.
hussein-awala commented on code in PR #30175:
URL: https://github.com/apache/airflow/pull/30175#discussion_r1167632694
##########
airflow/providers/apache/kafka/hooks/consume.py:
##########
Review Comment:
Actually it depends on the method `get_conn`, if it returns the client as I suggested, and you use it in the hook methods, no need to initialize and close the client in each method, we close it once at the end of the hook, and it is useful when we use the hook with the `with` statement.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [airflow] dylanbstorey commented on a diff in pull request #30175: Add provider for Apache Kafka
Posted by "dylanbstorey (via GitHub)" <gi...@apache.org>.
dylanbstorey commented on code in PR #30175:
URL: https://github.com/apache/airflow/pull/30175#discussion_r1167647290
##########
airflow/providers/apache/kafka/operators/consume.py:
##########
@@ -0,0 +1,184 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+from functools import partial
+from typing import Any, Callable, Sequence
+
+from airflow.exceptions import AirflowException
+from airflow.models import BaseOperator
+from airflow.providers.apache.kafka.hooks.consume import KafkaConsumerHook
+from airflow.utils.module_loading import import_string
+
+VALID_COMMIT_CADENCE = {"never", "end_of_batch", "end_of_operator"}
+
+
+class ConsumeFromTopicOperator(BaseOperator):
+ """An operator that consumes from Kafka a topic(s) and processing the messages.
+
+ The operator creates a Kafka consumer that reads a batch of messages from the cluster and processes them
+ using the user supplied callable function. The consumer will continue to read in batches until it reaches
+ the end of the log or reads a maximum number of messages is reached.
+
+ :param kafka_config_id: The connection object to use, defaults to "kafka_default"
+ :param topics: A list of topics or regex patterns the consumer should subscribe to.
+ :param apply_function: The function that should be applied to fetched one at a time.
+ name of dag file executing the function and the function name delimited by a `.`
+ :param apply_function_batch: The function that should be applied to a batch of messages fetched. Can not
+ be used with `apply_function`. Intended for transactional workloads where an expensive task might
+ be called before or after operations on the messages are taken.
+ :param apply_function_args: Additional arguments that should be applied to the callable, defaults to None
+ :param apply_function_kwargs: Additional key word arguments that should be applied to the callable
+ defaults to None
+ :param commit_cadence: When consumers should commit offsets ("never", "end_of_batch","end_of_operator"),
+ defaults to "end_of_operator";
+ if end_of_operator, the commit() is called based on the max_messages arg. Commits are made after the
+ operator has processed the apply_function method for the maximum messages in the operator.
+ if end_of_batch, the commit() is called based on the max_batch_size arg. Commits are made after each
+ batch has processed by the apply_function method for all messages in the batch.
+ if never, close() is called without calling the commit() method.
+ :param max_messages: The maximum total number of messages an operator should read from Kafka,
+ defaults to None
+ :param max_batch_size: The maximum number of messages a consumer should read when polling,
+ defaults to 1000
+ :param poll_timeout: How long the Kafka consumer should wait before determining no more messages are
+ available, defaults to 60
+
+ .. seealso::
+ For more information on how to use this operator, take a look at the guide:
+ :ref:`howto/operator:ConsumeFromTopicOperator`
+ """
+
+ BLUE = "#ffefeb"
+ ui_color = BLUE
+ template_fields = (
+ "topics",
+ "apply_function",
+ "apply_function_args",
+ "apply_function_kwargs",
+ "kafka_config_id",
+ )
+
+ def __init__(
+ self,
+ topics: str | Sequence[str],
+ kafka_config_id: str = "kafka_default",
+ apply_function: Callable[..., Any] | str | None = None,
+ apply_function_batch: Callable[..., Any] | str | None = None,
+ apply_function_args: Sequence[Any] | None = None,
+ apply_function_kwargs: dict[Any, Any] | None = None,
+ commit_cadence: str | None = "end_of_operator",
+ max_messages: int | None = None,
+ max_batch_size: int = 1000,
+ poll_timeout: float = 60,
+ **kwargs: Any,
+ ) -> None:
+
+ super().__init__(**kwargs)
+
+ self.topics = topics
+ self.apply_function = apply_function
+ self.apply_function_batch = apply_function_batch
+ self.apply_function_args = apply_function_args or ()
+ self.apply_function_kwargs = apply_function_kwargs or {}
+ self.kafka_config_id = kafka_config_id
+ self.commit_cadence = commit_cadence
+ self.max_messages = max_messages or True
+ self.max_batch_size = max_batch_size
+ self.poll_timeout = poll_timeout
+
+ if self.commit_cadence not in VALID_COMMIT_CADENCE:
+ raise AirflowException(
+ f"commit_cadence must be one of {VALID_COMMIT_CADENCE}. Got {self.commit_cadence}"
+ )
+
+ if self.max_messages and self.max_batch_size > self.max_messages:
+ self.log.warning(
+ "max_batch_size (%s) > max_messages (%s). Setting max_messages to %s ",
+ self.max_batch_size,
+ self.max_messages,
+ self.max_batch_size,
+ )
+
+ if self.commit_cadence == "never":
+ self.commit_cadence = None
+
+ if apply_function and apply_function_batch:
+ raise AirflowException(
+ "One of apply_function or apply_function_batch must be supplied, not both."
+ )
+
+ def execute(self, context) -> Any:
+
+ consumer = KafkaConsumerHook(topics=self.topics, kafka_config_id=self.kafka_config_id).get_consumer()
+
+ if isinstance(self.apply_function, str):
+ self.apply_function = import_string(self.apply_function)
+
+ if isinstance(self.apply_function_batch, str):
+ self.apply_function_batch = import_string(self.apply_function_batch)
+
+ if self.apply_function:
+ apply_callable = partial(
+ self.apply_function, *self.apply_function_args, **self.apply_function_kwargs # type: ignore
+ )
+
+ if self.apply_function_batch:
+ apply_callable = partial(
+ self.apply_function_batch, # type: ignore
+ *self.apply_function_args,
+ **self.apply_function_kwargs,
+ )
+
+ messages_left = self.max_messages
+ messages_processed = 0
+
+ while (
+ messages_left > 0
+ ): # bool(True > 0) == True in the case where self.max_messages isn't set by the user
+
+ if not isinstance(messages_left, bool):
+ batch_size = self.max_batch_size if messages_left > self.max_batch_size else messages_left
+ else:
+ batch_size = self.max_batch_size
+
+ msgs = consumer.consume(num_messages=batch_size, timeout=self.poll_timeout)
+ messages_left -= len(msgs)
Review Comment:
Ah i see the case i'm missing. I'll add a sentinel to manage it.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [airflow] josh-fell commented on a diff in pull request #30175: Add provider for Apache Kafka
Posted by "josh-fell (via GitHub)" <gi...@apache.org>.
josh-fell commented on code in PR #30175:
URL: https://github.com/apache/airflow/pull/30175#discussion_r1173210530
##########
airflow/providers/apache/kafka/hooks/base.py:
##########
@@ -0,0 +1,74 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+from typing import Any
+
+from confluent_kafka.admin import AdminClient
+
+from airflow.compat.functools import cached_property
+from airflow.hooks.base import BaseHook
+
+
+class KafkaBaseHook(BaseHook):
+ """A base hook for interacting with Apache Kafka"""
+
+ conn_name_attr = "kafka_config_id"
+ default_conn_name = "kafka_default"
+ conn_type = "kafka"
+ hook_name = "Apache Kafka"
+
+ def __init__(self, kafka_config_id=default_conn_name, *args, **kwargs):
Review Comment:
Can you add `kafka_config_id` to the docstring here too?
##########
airflow/providers/apache/kafka/sensors/kafka.py:
##########
@@ -0,0 +1,218 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+from typing import Any, Callable, Sequence
+
+from airflow.models import BaseOperator
+from airflow.providers.apache.kafka.triggers.await_message import AwaitMessageTrigger
+
+VALID_COMMIT_CADENCE = {"never", "end_of_batch", "end_of_operator"}
+
+
+class AwaitMessageSensor(BaseOperator):
+ """An Airflow operator that defers until a specific message is published to Kafka.
+
+ The operator creates a consumer that reads the Kafka log until it encounters a positive event.
+
+ The behavior of the consumer for this trigger is as follows:
+ - poll the Kafka topics for a message
+ - if no message returned, sleep
+ - process the message with provided callable and commit the message offset
+ - if callable returns any data, raise a TriggerEvent with the return data
+ - else continue to next message
+ - return event (as default xcom or specific xcom key)
+
+
+ :param kafka_config_id: The connection object to use, defaults to "kafka_default"
+ :param topics: Topics (or topic regex) to use for reading from
+ :param apply_function: The function to apply to messages to determine if an event occurred. As a dot
+ notation string.
+ :param apply_function_args: Arguments to be applied to the processing function,
+ defaults to None
+ :param apply_function_kwargs: Key word arguments to be applied to the processing function,
+ defaults to None
+ :param poll_timeout: How long the kafka consumer should wait for a message to arrive from the kafka
+ cluster,defaults to 1
+ :param poll_interval: How long the kafka consumer should sleep after reaching the end of the Kafka log,
+ defaults to 5
+ :param xcom_push_key: the name of a key to push the returned message to, defaults to None
+
+
+ """
+
+ BLUE = "#ffefeb"
+ ui_color = BLUE
+
+ template_fields = (
+ "topics",
+ "apply_function",
+ "apply_function_args",
+ "apply_function_kwargs",
+ "kafka_config_id",
+ )
+
+ def __init__(
+ self,
+ topics: Sequence[str],
+ apply_function: str,
+ kafka_config_id: str = "kafka_default",
+ apply_function_args: Sequence[Any] | None = None,
+ apply_function_kwargs: dict[Any, Any] | None = None,
+ poll_timeout: float = 1,
+ poll_interval: float = 5,
+ xcom_push_key=None,
+ **kwargs: Any,
+ ) -> None:
+
+ super().__init__(**kwargs)
+
+ self.topics = topics
+ self.apply_function = apply_function
+ self.apply_function_args = apply_function_args
+ self.apply_function_kwargs = apply_function_kwargs
+ self.kafka_config_id = kafka_config_id
+ self.poll_timeout = poll_timeout
+ self.poll_interval = poll_interval
+ self.xcom_push_key = xcom_push_key
+
+ def execute(self, context) -> Any:
+
+ self.defer(
+ trigger=AwaitMessageTrigger(
+ topics=self.topics,
+ apply_function=self.apply_function,
+ apply_function_args=self.apply_function_args,
+ apply_function_kwargs=self.apply_function_kwargs,
+ kafka_config_id=self.kafka_config_id,
+ poll_timeout=self.poll_timeout,
+ poll_interval=self.poll_interval,
+ ),
+ method_name="execute_complete",
+ )
+
+ def execute_complete(self, context, event=None):
+ if self.xcom_push_key:
+ self.xcom_push(context, key=self.xcom_push_key, value=event)
+ return event
+
+
+class AwaitMessageTriggerFunctionSensor(BaseOperator):
+ """An Airflow operator that defers until a specific message is published to
+ Kafka, then triggers a registered function
+
+
+ The behavior of the consumer for this trigger is as follows:
+ - poll the Kafka topics for a message
+ - if no message returned, sleep
+ - process the message with provided callable and commit the message offset
+ - if callable returns any data, raise a TriggerEvent with the return data
+ - else continue to next message
+ - return event (as default xcom or specific xcom key)
+
+
+ :param kafka_config_id: The connection object to use, defaults to "kafka_default"
+ :param topics: Topics (or topic regex) to use for reading from
+ :param apply_function: The function to apply to messages to determine if an event occurred. As a dot
+ notation string.
+ :param event_triggered_function: The callable to trigger once the apply_function encounters a
+ positive event.
+ :param apply_function_args: Arguments to be applied to the processing function, defaults to None
+ :param apply_function_kwargs: Key word arguments to be applied to the processing function,
+ defaults to None
+ :param poll_timeout: How long the kafka consumer should wait for a message to arrive from the kafka
+ cluster, defaults to 1
+ :param poll_interval: How long the kafka consumer should sleep after reaching the end of the Kafka log,
+ defaults to 5
+
+
+ """
+
+ BLUE = "#ffefeb"
+ ui_color = BLUE
+
+ template_fields = (
+ "topics",
+ "apply_function",
+ "apply_function_args",
+ "apply_function_kwargs",
+ "kafka_config_id",
+ )
+
+ def __init__(
+ self,
+ topics: Sequence[str],
+ apply_function: str,
+ event_triggered_function: Callable,
+ kafka_config_id: str = "kafka_default",
+ apply_function_args: Sequence[Any] | None = None,
+ apply_function_kwargs: dict[Any, Any] | None = None,
+ poll_timeout: float = 1,
+ poll_interval: float = 5,
+ **kwargs: Any,
+ ) -> None:
+
+ super().__init__(**kwargs)
+
+ self.topics = topics
+ self.apply_function = apply_function
+ self.apply_function_args = apply_function_args
+ self.apply_function_kwargs = apply_function_kwargs
+ self.kafka_config_id = kafka_config_id
+ self.poll_timeout = poll_timeout
+ self.poll_interval = poll_interval
+ self.event_triggered_function = event_triggered_function
+
+ if not callable(self.event_triggered_function):
+ raise TypeError(
+ "parameter event_triggered_function is expected to be of type callable,"
+ "got {type(event_triggered_function)}"
Review Comment:
```suggestion
f"got {type(event_triggered_function)}"
```
##########
airflow/providers/apache/kafka/sensors/kafka.py:
##########
@@ -0,0 +1,218 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+from typing import Any, Callable, Sequence
+
+from airflow.models import BaseOperator
+from airflow.providers.apache.kafka.triggers.await_message import AwaitMessageTrigger
+
+VALID_COMMIT_CADENCE = {"never", "end_of_batch", "end_of_operator"}
+
+
+class AwaitMessageSensor(BaseOperator):
+ """An Airflow operator that defers until a specific message is published to Kafka.
+
+ The operator creates a consumer that reads the Kafka log until it encounters a positive event.
Review Comment:
```suggestion
The sensor creates a consumer that reads the Kafka log until it encounters a positive event.
```
##########
airflow/providers/apache/kafka/hooks/base.py:
##########
@@ -0,0 +1,74 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+from typing import Any
+
+from confluent_kafka.admin import AdminClient
+
+from airflow.compat.functools import cached_property
+from airflow.hooks.base import BaseHook
+
+
+class KafkaBaseHook(BaseHook):
+ """A base hook for interacting with Apache Kafka"""
+
+ conn_name_attr = "kafka_config_id"
+ default_conn_name = "kafka_default"
+ conn_type = "kafka"
+ hook_name = "Apache Kafka"
+
+ def __init__(self, kafka_config_id=default_conn_name, *args, **kwargs):
+ """Initialize our Base"""
+ super().__init__()
+ self.kafka_config_id = kafka_config_id
+ self.get_conn
+
+ @staticmethod
+ def get_ui_field_behaviour() -> dict[str, Any]:
+ """Returns custom field behaviour"""
+ return {
+ "hidden_fields": ["schema", "login", "password", "port", "host"],
+ "relabeling": {"extra": "Config Dict"},
+ "placeholders": {
+ "extra": '{"bootstrap.servers": "localhost:9092"}',
+ },
+ }
+
+ def _get_client(self, config):
+ raise NotImplementedError
+
+ @cached_property
+ def get_conn(self) -> Any:
+ """get the configuration object"""
+ config = self.get_connection(self.kafka_config_id).extra_dejson
+
+ if not (config.get("bootstrap.servers", None)):
+ raise ValueError("config['bootstrap.servers'] must be provided.")
+
+ return self._get_client(config)
+
+ def test_connection(self) -> tuple[bool, str]:
+ """Test Connectivity from the UI"""
+ try:
+ t = AdminClient(self.get_conn(), timeout=10).list_topics()
Review Comment:
Does this still work given the recent changes to `get_conn`?
##########
airflow/providers/apache/kafka/sensors/kafka.py:
##########
@@ -0,0 +1,218 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+from typing import Any, Callable, Sequence
+
+from airflow.models import BaseOperator
+from airflow.providers.apache.kafka.triggers.await_message import AwaitMessageTrigger
+
+VALID_COMMIT_CADENCE = {"never", "end_of_batch", "end_of_operator"}
+
+
+class AwaitMessageSensor(BaseOperator):
+ """An Airflow operator that defers until a specific message is published to Kafka.
Review Comment:
```suggestion
"""An Airflow sensor that defers until a specific message is published to Kafka.
```
##########
airflow/providers/apache/kafka/sensors/kafka.py:
##########
@@ -0,0 +1,218 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+from typing import Any, Callable, Sequence
+
+from airflow.models import BaseOperator
+from airflow.providers.apache.kafka.triggers.await_message import AwaitMessageTrigger
+
+VALID_COMMIT_CADENCE = {"never", "end_of_batch", "end_of_operator"}
+
+
+class AwaitMessageSensor(BaseOperator):
+ """An Airflow operator that defers until a specific message is published to Kafka.
+
+ The operator creates a consumer that reads the Kafka log until it encounters a positive event.
+
+ The behavior of the consumer for this trigger is as follows:
+ - poll the Kafka topics for a message
+ - if no message returned, sleep
+ - process the message with provided callable and commit the message offset
+ - if callable returns any data, raise a TriggerEvent with the return data
+ - else continue to next message
+ - return event (as default xcom or specific xcom key)
+
+
+ :param kafka_config_id: The connection object to use, defaults to "kafka_default"
+ :param topics: Topics (or topic regex) to use for reading from
+ :param apply_function: The function to apply to messages to determine if an event occurred. As a dot
+ notation string.
+ :param apply_function_args: Arguments to be applied to the processing function,
+ defaults to None
+ :param apply_function_kwargs: Key word arguments to be applied to the processing function,
+ defaults to None
+ :param poll_timeout: How long the kafka consumer should wait for a message to arrive from the kafka
+ cluster,defaults to 1
+ :param poll_interval: How long the kafka consumer should sleep after reaching the end of the Kafka log,
+ defaults to 5
+ :param xcom_push_key: the name of a key to push the returned message to, defaults to None
+
+
+ """
+
+ BLUE = "#ffefeb"
+ ui_color = BLUE
+
+ template_fields = (
+ "topics",
+ "apply_function",
+ "apply_function_args",
+ "apply_function_kwargs",
+ "kafka_config_id",
+ )
+
+ def __init__(
+ self,
+ topics: Sequence[str],
+ apply_function: str,
+ kafka_config_id: str = "kafka_default",
+ apply_function_args: Sequence[Any] | None = None,
+ apply_function_kwargs: dict[Any, Any] | None = None,
+ poll_timeout: float = 1,
+ poll_interval: float = 5,
+ xcom_push_key=None,
+ **kwargs: Any,
+ ) -> None:
+
+ super().__init__(**kwargs)
+
+ self.topics = topics
+ self.apply_function = apply_function
+ self.apply_function_args = apply_function_args
+ self.apply_function_kwargs = apply_function_kwargs
+ self.kafka_config_id = kafka_config_id
+ self.poll_timeout = poll_timeout
+ self.poll_interval = poll_interval
+ self.xcom_push_key = xcom_push_key
+
+ def execute(self, context) -> Any:
+
+ self.defer(
+ trigger=AwaitMessageTrigger(
+ topics=self.topics,
+ apply_function=self.apply_function,
+ apply_function_args=self.apply_function_args,
+ apply_function_kwargs=self.apply_function_kwargs,
+ kafka_config_id=self.kafka_config_id,
+ poll_timeout=self.poll_timeout,
+ poll_interval=self.poll_interval,
+ ),
+ method_name="execute_complete",
+ )
+
+ def execute_complete(self, context, event=None):
+ if self.xcom_push_key:
+ self.xcom_push(context, key=self.xcom_push_key, value=event)
+ return event
+
+
+class AwaitMessageTriggerFunctionSensor(BaseOperator):
+ """An Airflow operator that defers until a specific message is published to
Review Comment:
```suggestion
"""An Airflow sensor that defers until a specific message is published to
```
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org