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

"""This module defines the IBM watsonx.data Presto stage."""

from ibm_watsonx_data_integration.cpd_models.connections_model import Connection
from ibm_watsonx_data_integration.services.datastage.models.connections.watsonx_data_connection import WatsonxDataConn
from ibm_watsonx_data_integration.services.datastage.models.flow.dag import StageNode
from ibm_watsonx_data_integration.services.datastage.models.stage_models.watsonx_data import watsonx_data
from typing import TYPE_CHECKING

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


class IBMWatsonxDataPrestoStage(StageNode):
    """This class defines the IBM watsonx.data Presto stage."""

    def __init__(self, dag: "DAG", configuration: watsonx_data = None, *, label: str = None) -> None:
        """Initializes an instance of a IBM watsonx.data Presto stage."""
        super().__init__(dag, configuration=configuration, label=label)
        self.configuration: watsonx_data = configuration or watsonx_data()

    def use_connection(self, connection: Connection) -> None:
        """Adds a project-level connection to this connector stage."""
        self.configuration.connection = WatsonxDataConn().from_connection(connection)
