spark-framework 0.2.0__py3-none-any.whl

spark_framework/__init__.py

@@ -0,0 +1,6 @@
+from spark_framework.core.config import PipelineConfig
+from spark_framework.core.pipeline import Pipeline, PipelineResult
+from spark_framework.framework import SparkFramework
+
+__version__ = "0.2.0"
+__all__ = ["SparkFramework", "Pipeline", "PipelineResult", "PipelineConfig"]

spark_framework/cli.py

@@ -0,0 +1,52 @@
+"""CLI entry point para o SparkFramework (spark-framework <config.json>)."""
+
+import argparse
+import sys
+
+from spark_framework import SparkFramework
+
+
+def build_parser() -> argparse.ArgumentParser:
+    parser = argparse.ArgumentParser(
+        prog="spark-framework",
+        description="SparkFramework — Motor de pipelines Spark orientado a JSON",
+        formatter_class=argparse.RawDescriptionHelpFormatter,
+        epilog="""
+Exemplos:
+  spark-framework examples/basic_parquet.json
+  spark-framework examples/iceberg_upsert.json --stop-spark
+  spark-framework tests/ingestion_csv_to_parquet.json
+        """,
+    )
+    parser.add_argument("config", help="Caminho para o arquivo JSON de configuracao")
+    parser.add_argument(
+        "--stop-spark",
+        action="store_true",
+        help="Encerra a SparkSession ao final",
+    )
+    return parser
+
+
+def main() -> None:
+    parser = build_parser()
+    args = parser.parse_args()
+
+    fw = SparkFramework()
+    result = fw.run(args.config)
+
+    print(result.summary())
+
+    if result.validation_results:
+        print("\nValidacoes:")
+        for r in result.validation_results:
+            print(f"  {r}")
+
+    if args.stop_spark:
+        fw.stop()
+
+    if not result.success:
+        sys.exit(1)
+
+
+if __name__ == "__main__":
+    main()

spark_framework/core/config.py

@@ -0,0 +1,154 @@
+from __future__ import annotations
+
+import json
+from dataclasses import dataclass, field
+from pathlib import Path
+from typing import Any, Dict, List, Optional
+
+
+@dataclass
+class SparkConfig:
+    app_name: str = "SparkFramework"
+    master: str = "local[*]"
+    configs: Dict[str, str] = field(default_factory=dict)
+
+    @classmethod
+    def from_dict(cls, data: Dict[str, Any]) -> SparkConfig:
+        return cls(
+            app_name=data.get("app_name", "SparkFramework"),
+            master=data.get("master", "local[*]"),
+            configs=data.get("configs", {}),
+        )
+
+
+@dataclass
+class InputConfig:
+    """Configuração da fonte de dados principal do pipeline."""
+
+    format: str
+    path: str
+    options: Dict[str, Any] = field(default_factory=dict)
+
+    @classmethod
+    def from_dict(cls, data: Dict[str, Any]) -> InputConfig:
+        return cls(
+            format=data["format"].lower(),
+            path=data["path"],
+            options=data.get("options", {}),
+        )
+
+
+_TRANSFORMATION_META_KEYS = {"type", "skip_if_false"}
+
+
+@dataclass
+class TransformationConfig:
+    type: str
+    params: Dict[str, Any] = field(default_factory=dict)
+    skip_if_false: Optional[str] = None
+
+    @classmethod
+    def from_dict(cls, data: Dict[str, Any]) -> TransformationConfig:
+        return cls(
+            type=data["type"],
+            skip_if_false=data.get("skip_if_false"),
+            params={k: v for k, v in data.items() if k not in _TRANSFORMATION_META_KEYS},
+        )
+
+
+@dataclass
+class ValidationRule:
+    type: str
+    params: Dict[str, Any] = field(default_factory=dict)
+
+    @classmethod
+    def from_dict(cls, data: Dict[str, Any]) -> ValidationRule:
+        return cls(
+            type=data["type"],
+            params={k: v for k, v in data.items() if k != "type"},
+        )
+
+
+@dataclass
+class ValidationConfig:
+    on_failure: str = "fail"  # fail | warn | skip
+    rules: List[ValidationRule] = field(default_factory=list)
+
+    @classmethod
+    def from_dict(cls, data: Dict[str, Any]) -> ValidationConfig:
+        return cls(
+            on_failure=data.get("on_failure", "fail"),
+            rules=[ValidationRule.from_dict(r) for r in data.get("rules", [])],
+        )
+
+
+@dataclass
+class OutputConfig:
+    """Configuração de um destino de escrita do pipeline.
+
+    O campo `columns` permite selecionar quais colunas serão escritas
+    neste destino específico, sem alterar o DataFrame das demais saídas.
+    Se omitido, todas as colunas são escritas.
+    """
+
+    format: str
+    path: str
+    mode: str = "overwrite"  # append | overwrite | merge
+    partition_by: List[str] = field(default_factory=list)
+    columns: Optional[List[str]] = None
+    options: Dict[str, Any] = field(default_factory=dict)
+
+    @classmethod
+    def from_dict(cls, data: Dict[str, Any]) -> OutputConfig:
+        return cls(
+            format=data["format"].lower(),
+            path=data["path"],
+            mode=data.get("mode", "overwrite"),
+            partition_by=data.get("partition_by", []),
+            columns=data.get("columns"),
+            options=data.get("options", {}),
+        )
+
+
+@dataclass
+class PipelineConfig:
+    name: str
+    input: InputConfig
+    outputs: List[OutputConfig]
+    description: str = ""
+    spark: SparkConfig = field(default_factory=SparkConfig)
+    transformations: List[TransformationConfig] = field(default_factory=list)
+    validations: ValidationConfig = field(default_factory=ValidationConfig)
+
+    @classmethod
+    def from_file(cls, path: str) -> PipelineConfig:
+        from spark_framework.utils.includes import resolve_includes
+        content = Path(path).read_text(encoding="utf-8")
+        data = json.loads(content)
+        data = resolve_includes(data, Path(path).parent)
+        return cls.from_dict(data)
+
+    @classmethod
+    def from_dict(cls, data: Dict[str, Any]) -> PipelineConfig:
+        # Aceita "output" (objeto único) ou "outputs" (lista)
+        if "outputs" in data:
+            outputs = [OutputConfig.from_dict(o) for o in data["outputs"]]
+        elif "output" in data:
+            outputs = [OutputConfig.from_dict(data["output"])]
+        else:
+            raise ValueError(
+                "O JSON do pipeline precisa ter 'output' (objeto) ou 'outputs' (lista)."
+            )
+
+        return cls(
+            name=data["name"],
+            description=data.get("description", ""),
+            spark=SparkConfig.from_dict(data.get("spark", {})),
+            input=InputConfig.from_dict(data["input"]),
+            transformations=[
+                TransformationConfig.from_dict(t)
+                for t in data.get("transformations", [])
+            ],
+            validations=ValidationConfig.from_dict(data.get("validations", {})),
+            outputs=outputs,
+        )

spark_framework/core/context.py

@@ -0,0 +1,74 @@
+from __future__ import annotations
+
+import os
+
+from pyspark.sql import SparkSession
+
+from spark_framework.core.config import SparkConfig
+
+
+def _detect_environment() -> str:
+    """Detecta o ambiente de execução atual.
+
+    Returns:
+        "databricks" | "emr" | "dataproc" | "synapse" | "local"
+    """
+    if "DATABRICKS_RUNTIME_VERSION" in os.environ:
+        return "databricks"
+    if os.environ.get("EMR_CLUSTER_ID") or os.path.exists("/mnt/var/lib/info/job-flow.json"):
+        return "emr"
+    if os.environ.get("DATAPROC_IMAGE_VERSION") or os.environ.get("DATAPROC_CLUSTER_NAME"):
+        return "dataproc"
+    if os.environ.get("SYNAPSE_WORKSPACE_NAME") or os.environ.get("AZURE_DATABRICKS_ORG_ID"):
+        return "synapse"
+    return "local"
+
+
+class SparkContextManager:
+    """Gerencia um singleton de SparkSession para toda a vida do pipeline.
+
+    Em Databricks reutiliza a sessão ativa do runtime — nunca cria uma nova.
+    Em outros ambientes (EMR, Dataproc, Synapse, local) cria/reutiliza via builder.
+    """
+
+    _session: SparkSession | None = None
+
+    @classmethod
+    def get_or_create(cls, config: SparkConfig) -> SparkSession:
+        if cls._session is not None:
+            return cls._session
+
+        env = _detect_environment()
+
+        if env == "databricks":
+            # No Databricks a sessão já existe; criar uma nova causaria erro.
+            cls._session = SparkSession.getActiveSession()
+            if cls._session is None:
+                # Fallback improvável, mas seguro
+                cls._session = SparkSession.builder.getOrCreate()
+        else:
+            builder = SparkSession.builder.appName(config.app_name)
+
+            # Master só faz sentido fora de clusters gerenciados
+            if env == "local":
+                builder = builder.master(config.master)
+
+            for key, value in config.configs.items():
+                builder = builder.config(key, value)
+
+            cls._session = builder.getOrCreate()
+            cls._session.sparkContext.setLogLevel("WARN")
+
+        return cls._session
+
+    @classmethod
+    def stop(cls) -> None:
+        env = _detect_environment()
+        if cls._session is not None and env != "databricks":
+            cls._session.stop()
+        cls._session = None
+
+    @classmethod
+    def current_environment(cls) -> str:
+        """Retorna o ambiente detectado (útil para logs e diagnóstico)."""
+        return _detect_environment()

spark_framework/core/pipeline.py

@@ -0,0 +1,145 @@
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from typing import Any, Dict, List, Optional
+
+from pyspark.sql import DataFrame, SparkSession
+from pyspark.sql import functions as F
+
+from spark_framework.core.config import OutputConfig, PipelineConfig
+from spark_framework.core.context import SparkContextManager
+from spark_framework.io.factory import ReaderFactory, WriterFactory
+from spark_framework.transform.engine import TransformationEngine
+from spark_framework.validation.base import ValidationResult
+from spark_framework.validation.engine import ValidationEngine
+from spark_framework.utils.logger import logger
+
+
+@dataclass
+class PipelineResult:
+    pipeline_name: str
+    success: bool
+    rows_read: int = 0
+    rows_written: int = 0
+    validation_results: List[ValidationResult] = field(default_factory=list)
+    error: Optional[str] = None
+    output_df: Optional[DataFrame] = None  # df após transformações; disponível quando input_df é injetado
+
+    def summary(self) -> str:
+        if not self.success:
+            return f"[FAIL] '{self.pipeline_name}': {self.error}"
+        passed = sum(1 for r in self.validation_results if r.passed)
+        total = len(self.validation_results)
+        return (
+            f"[OK] '{self.pipeline_name}' | "
+            f"lidos={self.rows_read} | "
+            f"escritos={self.rows_written} | "
+            f"validacoes={passed}/{total}"
+        )
+
+
+class Pipeline:
+    """Orquestra o fluxo: leitura → transformação → validação → escrita."""
+
+    def __init__(
+        self,
+        config: PipelineConfig,
+        transform_engine: Optional[TransformationEngine] = None,
+        validation_engine: Optional[ValidationEngine] = None,
+        input_df: Optional[DataFrame] = None,
+        columns: Optional[Dict[str, Any]] = None,
+    ) -> None:
+        self.config = config
+        self._transform_engine = transform_engine or TransformationEngine()
+        self._validation_engine = validation_engine or ValidationEngine()
+        self._input_df = input_df
+        self._columns: Dict[str, Any] = columns or {}
+
+    @classmethod
+    def from_file(cls, path: str) -> Pipeline:
+        return cls(PipelineConfig.from_file(path))
+
+    @classmethod
+    def from_dict(cls, data: dict) -> Pipeline:
+        return cls(PipelineConfig.from_dict(data))
+
+    def run(self) -> PipelineResult:
+        log = logger.bind(pipeline=self.config.name)
+        log.info("Pipeline iniciado")
+
+        try:
+            spark = SparkContextManager.get_or_create(self.config.spark)
+
+            if self._input_df is not None:
+                df = self._input_df
+                rows_read = 0
+                log.info("Input df injetado externamente", colunas=len(df.columns))
+            else:
+                df = ReaderFactory.create(spark, self.config.input).read()
+                df = df.withColumn("ingestion_ts", F.current_timestamp())
+                rows_read = df.count()
+                log.info(
+                    "Leitura concluida",
+                    linhas=rows_read,
+                    formato=self.config.input.format,
+                )
+
+            for col_name, value in self._columns.items():
+                df = df.withColumn(col_name, F.lit(value))
+            if self._columns:
+                log.info("Colunas injetadas", colunas=list(self._columns))
+
+            df = self._transform_engine.apply(df, self.config.transformations)
+            log.info("Transformacoes aplicadas")
+
+            validation_results = self._validation_engine.validate(
+                df, self.config.validations
+            )
+
+            rows_written = df.count()
+            self._write_outputs(spark, df, log)
+            log.info("Pipeline concluido", linhas_escritas=rows_written)
+
+            return PipelineResult(
+                pipeline_name=self.config.name,
+                success=True,
+                rows_read=rows_read,
+                rows_written=rows_written,
+                validation_results=validation_results,
+                output_df=df,
+            )
+
+        except Exception as exc:
+            log.error("Pipeline falhou", error=str(exc))
+            return PipelineResult(
+                pipeline_name=self.config.name,
+                success=False,
+                error=str(exc),
+            )
+
+    def _write_outputs(
+        self, spark: SparkSession, df: DataFrame, log
+    ) -> None:
+        for output in self.config.outputs:
+            output_df = self._project_columns(df, output)
+            log.info(
+                "Escrevendo output",
+                formato=output.format,
+                path=output.path,
+                modo=output.mode,
+                colunas=output.columns or "todas",
+            )
+            WriterFactory.create(spark, output).write(output_df)
+
+    @staticmethod
+    def _project_columns(df: DataFrame, output: OutputConfig) -> DataFrame:
+        """Aplica seleção de colunas se o output tiver 'columns' definido."""
+        if not output.columns:
+            return df
+        missing = [c for c in output.columns if c not in df.columns]
+        if missing:
+            raise ValueError(
+                f"Colunas inexistentes no output '{output.path}': {missing}. "
+                f"Colunas disponiveis: {df.columns}"
+            )
+        return df.select(*output.columns)

spark_framework/framework.py

@@ -0,0 +1,164 @@
+from __future__ import annotations
+
+import json
+from pathlib import Path
+from typing import Any, Dict, Optional
+
+from pyspark.sql import DataFrame
+
+from spark_framework.core.config import PipelineConfig, SparkConfig
+from spark_framework.utils.template import apply_template
+from spark_framework.utils.includes import resolve_includes
+from spark_framework.core.context import SparkContextManager
+from spark_framework.core.pipeline import Pipeline, PipelineResult
+from spark_framework.io.factory import ReaderFactory, WriterFactory
+from spark_framework.io.base import BaseReader, BaseWriter
+from spark_framework.transform.base import BaseTransformation
+from spark_framework.transform.engine import TransformationEngine
+from spark_framework.validation.base import BaseValidator
+from spark_framework.validation.engine import ValidationEngine
+
+
+class SparkFramework:
+    """Ponto de entrada principal para uso do framework como biblioteca.
+
+    Gerencia uma única SparkSession compartilhada entre múltiplas execuções.
+    Permite registrar formatos, transformações e validators customizados que
+    ficam disponíveis em todos os pipelines executados pela mesma instância.
+
+    Uso básico:
+        fw = SparkFramework()
+        fw.run("pipeline_clientes.json")
+        fw.run("pipeline_pedidos.json")
+        fw.stop()
+
+    Com configurações Spark personalizadas:
+        fw = SparkFramework(spark={"master": "yarn", "app_name": "MeuJob"})
+
+    Com extensões customizadas:
+        fw = SparkFramework()
+        fw.register_reader("delta", DeltaReader)
+        fw.register_transformation("normalize", NormalizeTransformation)
+        fw.run("config.json")
+    """
+
+    def __init__(self, spark: Optional[Dict[str, Any]] = None) -> None:
+        self._spark_config = SparkConfig.from_dict(spark or {})
+        self._transform_engine = TransformationEngine()
+        self._validation_engine = ValidationEngine()
+        SparkContextManager.get_or_create(self._spark_config)
+
+    # ------------------------------------------------------------------
+    # Execução de pipelines
+    # ------------------------------------------------------------------
+
+    def run(
+        self,
+        config_path: str,
+        input_df: Optional[DataFrame] = None,
+        columns: Optional[Dict[str, Any]] = None,
+        params: Optional[Dict[str, Any]] = None,
+    ) -> PipelineResult:
+        """Executa um pipeline a partir de um arquivo JSON.
+
+        Args:
+            config_path: caminho para o JSON de configuração.
+            input_df:    DataFrame de entrada; quando fornecido substitui o 'input'
+                         declarado no JSON e não adiciona ingestion_ts automaticamente.
+            columns:     Colunas literais a injetar no df antes das transformações,
+                         ex: {"param_tipo_ativo": "NC", "param_registradora": "CERC"}.
+            params:      Valores de runtime substituídos como {chave} no JSON antes do
+                         parse. Listas viram SQL IN (ex: 'a', 'b'); booleanos viram
+                         "true"/"" (vazio = falsy dispara skip_if_false).
+        """
+        config = self._load_config(config_path, params)
+        self._apply_spark_override(config)
+        return self._execute(config, input_df=input_df, columns=columns)
+
+    def run_from_dict(
+        self,
+        config: Dict[str, Any],
+        input_df: Optional[DataFrame] = None,
+        columns: Optional[Dict[str, Any]] = None,
+        params: Optional[Dict[str, Any]] = None,
+    ) -> PipelineResult:
+        """Executa um pipeline a partir de um dicionário Python."""
+        pipeline_config = self._load_config_from_dict(config, params)
+        self._apply_spark_override(pipeline_config)
+        return self._execute(pipeline_config, input_df=input_df, columns=columns)
+
+    # ------------------------------------------------------------------
+    # Registro de extensões
+    # ------------------------------------------------------------------
+
+    def register_reader(self, format_name: str, reader_cls: type[BaseReader]) -> None:
+        """Registra um leitor customizado para um novo formato."""
+        ReaderFactory.register(format_name, reader_cls)
+
+    def register_writer(self, format_name: str, writer_cls: type[BaseWriter]) -> None:
+        """Registra um escritor customizado para um novo formato."""
+        WriterFactory.register(format_name, writer_cls)
+
+    def register_transformation(
+        self, name: str, transformation_cls: type[BaseTransformation]
+    ) -> None:
+        """Registra uma transformação customizada disponível via JSON."""
+        self._transform_engine.register(name, transformation_cls)
+
+    def register_validator(
+        self, name: str, validator_cls: type[BaseValidator]
+    ) -> None:
+        """Registra um validator customizado disponível via JSON."""
+        self._validation_engine.register(name, validator_cls)
+
+    # ------------------------------------------------------------------
+    # Ciclo de vida
+    # ------------------------------------------------------------------
+
+    def stop(self) -> None:
+        """Encerra a SparkSession."""
+        SparkContextManager.stop()
+
+    # ------------------------------------------------------------------
+    # Interno
+    # ------------------------------------------------------------------
+
+    def _execute(
+        self,
+        config: PipelineConfig,
+        input_df: Optional[DataFrame] = None,
+        columns: Optional[Dict[str, Any]] = None,
+    ) -> PipelineResult:
+        pipeline = Pipeline(
+            config,
+            transform_engine=self._transform_engine,
+            validation_engine=self._validation_engine,
+            input_df=input_df,
+            columns=columns,
+        )
+        return pipeline.run()
+
+    def _load_config(self, path: str, params: Optional[Dict[str, Any]]) -> PipelineConfig:
+        raw = Path(path).read_text(encoding="utf-8")
+        if params:
+            raw = apply_template(raw, params)
+        data = json.loads(raw)
+        data = resolve_includes(data, Path(path).parent, params)
+        return PipelineConfig.from_dict(data)
+
+    def _load_config_from_dict(
+        self, config: Dict[str, Any], params: Optional[Dict[str, Any]]
+    ) -> PipelineConfig:
+        if params:
+            raw = apply_template(json.dumps(config), params)
+            config = json.loads(raw)
+        config = resolve_includes(config, Path.cwd(), params)
+        return PipelineConfig.from_dict(config)
+
+    def _apply_spark_override(self, config: PipelineConfig) -> None:
+        """Garante que configs Spark do framework prevalecem sobre o JSON."""
+        if self._spark_config.app_name != "SparkFramework":
+            config.spark.app_name = self._spark_config.app_name
+        if self._spark_config.master != "local[*]":
+            config.spark.master = self._spark_config.master
+        config.spark.configs.update(self._spark_config.configs)

spark_framework/io/base.py

@@ -0,0 +1,25 @@
+from abc import ABC, abstractmethod
+
+from pyspark.sql import DataFrame, SparkSession
+
+from spark_framework.core.config import InputConfig, OutputConfig
+
+
+class BaseReader(ABC):
+    def __init__(self, spark: SparkSession, config: InputConfig) -> None:
+        self.spark = spark
+        self.config = config
+
+    @abstractmethod
+    def read(self) -> DataFrame:
+        ...
+
+
+class BaseWriter(ABC):
+    def __init__(self, spark: SparkSession, config: OutputConfig) -> None:
+        self.spark = spark
+        self.config = config
+
+    @abstractmethod
+    def write(self, df: DataFrame) -> None:
+        ...

spark_framework/io/csv.py

@@ -0,0 +1,29 @@
+from pyspark.sql import DataFrame
+
+from spark_framework.io.base import BaseReader, BaseWriter
+
+_CSV_READ_DEFAULTS = {
+    "header": "true",
+    "inferSchema": "true",
+    "encoding": "UTF-8",
+}
+
+_CSV_WRITE_DEFAULTS = {
+    "header": "true",
+    "encoding": "UTF-8",
+}
+
+
+class CsvReader(BaseReader):
+    def read(self) -> DataFrame:
+        options = {**_CSV_READ_DEFAULTS, **self.config.options}
+        return self.spark.read.options(**options).csv(self.config.path)
+
+
+class CsvWriter(BaseWriter):
+    def write(self, df: DataFrame) -> None:
+        options = {**_CSV_WRITE_DEFAULTS, **self.config.options}
+        writer = df.write.mode(self.config.mode).options(**options)
+        if self.config.partition_by:
+            writer = writer.partitionBy(*self.config.partition_by)
+        writer.csv(self.config.path)