# -------------------------------------------------------------------
# This file was auto-generated. Do not edit directly.
# Generated by: StageGeneratorV2
# -------------------------------------------------------------------

from typing import TYPE_CHECKING

from ibm_watsonx_data_integration.services.datastage.models.flow.dag import StageNode
from ibm_watsonx_data_integration.services.datastage.models.stage_models.apache_kafka import apache_kafka

from ibm_watsonx_data_integration.cpd_models.connections_model import Connection
from ibm_watsonx_data_integration.services.datastage.models.connections.apache_kafka_connection import ApacheKafkaConn


if TYPE_CHECKING:
    from ibm_watsonx_data_integration.services.datastage.models.flow.dag import DAG



class ApacheKafkaStage(StageNode):
    def __init__(self, dag: "DAG", configuration: apache_kafka=None, *, label: str = None):
        super().__init__(dag, configuration=configuration, label=label)
        self.configuration: apache_kafka = configuration or apache_kafka()

    def use_connection(self, connection: Connection):
        self.configuration.connection = ApacheKafkaConn().from_connection(connection)
