from typing import Optional

from pyspark.sql import DataFrame as SparkDataFrame
from pyspark.sql import SparkSession

from mlpype.base.data.data_source import DataSource


class SparkSqlSource(DataSource[SparkDataFrame]):
    def __init__(
        self,
        query: str,
        spark_session: Optional[SparkSession] = None,
    ) -> None:
        """Use Spark SQL as an input data source.

        Args:
            spark_session (SparkSession): The current SparkSession.
            query (str): The query to run to get the data.
        """
        super().__init__()
        self.query = query
        if spark_session is None:
            spark_session = SparkSession.getActiveSession()
        self.spark_session = spark_session

    def read(self) -> SparkDataFrame:
        """Execute the query stored in this object.

        Returns:
            SparkDataFrame: The DataFrame generated by the query.
        """
        return self.spark_session.sql(self.query)
