Airflow NiFi Pipeline Utils

⚠️ Dlytica fork (2.1.x line). This is the dlytica-owned fork of dlytica-gcp/airflow-nifi-pipeline-utils, maintained on main of this repo (dlytica-owned). Pre-fork upstream is tagged v2.0.2-upstream; that line is frozen at 2.0.2; the 2.1.x line adds the historical (bounded-range, append-only, watermark-invisible) and full_incremental (backfill + incremental) load modes for the DataNature platform. It is not published to PyPI/TestPyPI — it ships as a wheel baked into the Airflow worker image (see the DataNature dags repo). Nabil production remains on upstream 2.0.2 and is unaffected.

A production-grade Python framework for building reliable, scalable data ingestion pipelines using Apache Airflow and Apache NiFi.

Built by a data engineer, for data engineers. Eliminates thousands of lines of boilerplate code while providing enterprise-grade reliability, observability, and maintainability.

---

Table of Contents

• Overview
• Key Features
• Architecture
• Installation
• Quick Start
• Core Concepts
• Load Patterns
• Components Reference
• Complete Examples
• Best Practices
• Migration Guide

---

Overview

What is this?

Airflow NiFi Pipeline Utils is a comprehensive framework that standardizes and simplifies the integration between Apache Airflow (workflow orchestration) and Apache NiFi(data flow automation). It provides everything needed to build production-ready data ingestion pipelines.

• Configuration Management: Type-safe configuration with validation
• Authentication: Support for both basic auth and enterprise OIDC/Keycloak
• Load Patterns: Three built-in load types (incremental, full, reconcile)
• State Tracking: PostgreSQL-based pipeline state with watermark management
• Error Handling: Intelligent classification of Retryable vs Non-Retryable failures
• Connection Management: Thread-safe HTTP session pooling
• Batch Processing: Parallel batch execution with automatic retry logic

When should you use this?

This library is designed for teams who:

• Build data pipelines that move data from source systems to data lakes/warehouses
• Use Apache Airflow for orchestration and Apache NiFi for data flow
• Need reliable incremental loading with watermark tracking
• Require state management and observability across pipeline runs
• Want to eliminate code duplication across multiple pipelines
• Need enterprise authentication (OIDC/Keycloak) support
• Run concurrent batch processing workloads

What does it replace?

Before (Manual Implementation):

# Repeated in every DAG - 200+ lines per pipeline
def authenticate_nifi():
    # 30 lines of auth logic
    pass

def generate_batches():
    # 50 lines of time windowing logic
    pass

def send_to_nifi():
    # 80 lines of retry/error handling
    pass

def track_state():
    # 40 lines of PostgreSQL operations
    pass

After (Using This Library):

from airflow_nifi_pipeline_utils import (
    NiFiConfig, SourceConfig, NiFiProcessor, 
    PipelineStateManager, NiFiConnectionManager
)

# 20 lines of configuration + execution
# All complexity handled by the library

Impact: Reduces pipeline code by 80-90% while improving reliability and maintainability.

---

Key Features

Enterprise Authentication

• Dual Mode Support: NiFi native authentication or OIDC/Keycloak
• Automatic Token Management: Caching, refresh, and thread-safe token handling
• Session Pooling: Connection reuse across batch operations
• Security Best Practices: No hardcoded credentials, supports Airflow Variables/Secrets

Three Load Patterns

1. Incremental Loads (Most Common)

• Time-windowed data processing based on bookmark columns
• Automatic watermark tracking and updates
• Handles late-arriving data with configurable lookback windows
• Partitioned execution for large datasets

2. Full Loads (Complete Refreshes)

• Entire table reloads without time filtering
• No watermark updates (doesn't affect incremental loads)
• Ideal for dimension tables and monthly snapshots
• Can run in parallel with incremental loads

3. Reconcile Loads (Backfills)

• Reprocess specific date ranges for data corrections
• No watermark updates (doesn't interfere with ongoing incrementals)
• Partition large backfills into manageable batches
• Perfect for fixing data quality issues or gap filling

4. Historical Loads (2.1.x — append-only bounded ranges)

• Load a bounded [historical_start, historical_end) range append-only: batches carry load_mode='historical', which the NiFi flow's route guards use to send them down the append path — never the truncate path a plain full load takes
• update_watermark is forced False (not just defaulted): historical rows always record high_watermark NULL, so an entity that is also cursor-tracked by incremental runs never has its watermark disturbed — safe to run in parallel with ongoing incrementals
• Chunk with SourceConfigResult.generate_time_partitions() (partition_hours-sized, half-open, contiguous)
• Composes into "full → incremental": append-load history up to a cutoff, seed the incremental watermark at that cutoff, run incremental onward
• ⚠️ Existing pipeline_info databases must apply queries/migrations/2.1.0-widen-load-type-check.sql before first use (the pipeline_run.load_type CHECK predates the new value)

State Management & Observability

PostgreSQL-Based Pipeline Tracking:

• Pipeline Registry: Central catalog of all pipelines
• Run Tracking: Every execution tracked with load_type, status, timestamps
• Batch-Level State: Per-partition success/failure tracking
• Watermark Management: High watermark tracking for incremental loads
• Medallion Architecture Support: Bronze/Silver/Gold layer tracking
• Performance Metrics: Execution times, retry counts, failure analysis

Intelligent Error Handling

Error Classification System:

• Retryable Errors: Network timeouts, rate limiting (429), server errors (502/503/504)
• Non-Retryable Errors: Auth failures (401/403), not found (404), bad requests (400)
• Exponential Backoff: 2s → 4s → 8s → 16s with ±20% jitter
• Circuit Breaker Pattern: Prevents cascading failures
• Detailed Logging: Full context for debugging and alerting

Performance & Scalability

• Thread-Safe Operations: Concurrent batch processing without conflicts
• Connection Pooling: 5 persistent connections, max 10
• Batch Partitioning: Split large time ranges into parallel-processable chunks
• Resource Efficiency: Minimal memory footprint, optimized SQL generation
• Health Checks: Pre-flight validation before sending data

---

Architecture

System Overview

┌──────────────────────────────────────────────────────────────────────────┐
│                      APACHE AIRFLOW (Orchestration Layer)                 │
│                                                                            │
│  ┌────────────────────────────────────────────────────────────────────┐  │
│  │                          Your DAG Code                              │  │
│  │                                                                     │  │
│  │  • Define pipeline configuration                                  │  │
│  │  • Set load type and schedule                                     │  │
│  │  • Configure source tables and watermarks                         │  │
│  └────────────────────────────────────────────────────────────────────┘  │
│                                   ↓                                       │
│  ┌────────────────────────────────────────────────────────────────────┐  │
│  │              Airflow NiFi Pipeline Utils (This Library)            │  │
│  │                                                                     │  │
│  │  ┌──────────────────┐  ┌──────────────────┐  ┌─────────────────┐ │  │
│  │  │  Configuration   │  │  State Manager   │  │   Connection    │ │  │
│  │  │  • NiFiConfig    │  │  • Register      │  │   Manager       │ │  │
│  │  │  • SourceConfig  │  │  • Track State   │  │  • Auth         │ │  │
│  │  │  • Validation    │  │  • Watermarks    │  │  • Sessions     │ │  │
│  │  └──────────────────┘  └──────────────────┘  └─────────────────┘ │  │
│  │                                                                     │  │
│  │  ┌──────────────────┐  ┌──────────────────┐  ┌─────────────────┐ │  │
│  │  │  Batch Generator │  │   NiFi Processor │  │  Error Handler  │ │  │
│  │  │  • Time Windows  │  │  • Send Data     │  │  • Classify     │ │  │
│  │  │  • Partitioning  │  │  • Retry Logic   │  │  • Backoff      │ │  │
│  │  │  • SQL Gen       │  │  • Validation    │  │  • Recovery     │ │  │
│  │  └──────────────────┘  └──────────────────┘  └─────────────────┘ │  │
│  └────────────────────────────────────────────────────────────────────┘  │
└──────────────────────────────────────────────────────────────────────────┘
                                   ↓ HTTP POST
┌──────────────────────────────────────────────────────────────────────────┐
│                      APACHE NIFI (Data Flow Layer)                        │
│                                                                            │
│  ┌────────────────────────────────────────────────────────────────────┐  │
│  │  ListenHTTP Processor (Port per schema: 7000, 7001, 7002...)       │  │
│  │  • Receives batch metadata + SQL query from Airflow                │  │
│  │  • Routes to appropriate data flow based on schema                 │  │
│  └────────────────────────────────────────────────────────────────────┘  │
│                                   ↓                                       │
│  ┌────────────────────────────────────────────────────────────────────┐  │
│  │                    Your NiFi Data Flow                              │  │
│  │                                                                     │  │
│  │  ExecuteSQL → ConvertRecord → CompressContent → PutHDFS/PutS3     │  │
│  │  • Extract data from source database                               │  │
│  │  • Transform to target format (Avro/Parquet/JSON)                 │  │
│  │  • Write to data lake/warehouse                                    │  │
│  └────────────────────────────────────────────────────────────────────┘  │
└──────────────────────────────────────────────────────────────────────────┘
                                   ↓ Status Updates
┌──────────────────────────────────────────────────────────────────────────┐
│                  POSTGRESQL (State Persistence Layer)                     │
│                                                                            │
│  ┌────────────────────────────────────────────────────────────────────┐  │
│  │  state.pipeline_data (Pipeline Registry)                           │  │
│  │  • pipeline_id, pipeline_name, description, created_at             │  │
│  └────────────────────────────────────────────────────────────────────┘  │
│                                                                            │
│  ┌────────────────────────────────────────────────────────────────────┐  │
│  │  state.pipeline_run (Execution Tracking)                           │  │
│  │  • run_id, pipeline_id, load_type, started_at, bronze_ended_at    │  │
│  │  • pipeline_status, ingestion_status, failed_batch_count          │  │
│  └────────────────────────────────────────────────────────────────────┘  │
│                                                                            │
│  ┌────────────────────────────────────────────────────────────────────┐  │
│  │  state.entity_state (Batch-Level Tracking)                         │  │
│  │  • batch_id, run_id, entity_name, partition_key                   │  │
│  │  • starting_watermark, end_watermark, high_watermark              │  │
│  │  • bronze_state_status, started_at, bronze_completed_at           │  │
│  │  • Enables incremental load resume and failure recovery           │  │
│  └────────────────────────────────────────────────────────────────────┘  │
└──────────────────────────────────────────────────────────────────────────┘

Data Flow Sequence

┌────────────┐
│  DAG Start │
└─────┬──────┘
      │
      ▼
┌─────────────────────────────────────────────────────────────┐
│ 1. CONFIGURATION PHASE                                       │
│    • Load pipeline config (table, schema, load_type)        │
│    • Initialize NiFiConfig with auth credentials            │
│    • Create PipelineStateManager                            │
└─────┬───────────────────────────────────────────────────────┘
      │
      ▼
┌─────────────────────────────────────────────────────────────┐
│ 2. REGISTRATION PHASE                                        │
│    • Register pipeline in pipeline_data table               │
│    • Start pipeline_run with load_type                      │
│    • Get last high_watermark (for incremental loads)        │
└─────┬───────────────────────────────────────────────────────┘
      │
      ▼
┌─────────────────────────────────────────────────────────────┐
│ 3. BATCH GENERATION PHASE                                    │
│    • Create SourceConfig with load_type                     │
│    • Calculate time_range:                                   │
│      - Incremental: last_watermark → last_watermark+window  │
│      - Full: None (no time filter)                          │
│      - Reconcile: reconcile_start → reconcile_end           │
│    • Partition into batches based on partition_hours        │
│    • Generate SQL query for each batch                      │
└─────┬───────────────────────────────────────────────────────┘
      │
      ▼
┌─────────────────────────────────────────────────────────────┐
│ 4. AUTHENTICATION PHASE                                      │
│    • NiFiConnectionManager authenticates:                   │
│      - NiFi native: POST /nifi-api/access/token            │
│      - OIDC: POST keycloak/token with client credentials   │
│    • Cache token in thread-local storage                    │
│    • Create HTTP session with connection pooling            │
└─────┬───────────────────────────────────────────────────────┘
      │
      ▼
┌─────────────────────────────────────────────────────────────┐
│ 5. BATCH PROCESSING PHASE (per batch)                       │
│    • Create BatchDataResult with batch metadata             │
│    • Initialize NiFiProcessor                               │
│    • Health check NiFi endpoint                             │
│    • Send batch via HTTP POST to ListenHTTP                 │
│    • Retry on failure with exponential backoff:             │
│      - Attempt 1: immediate                                 │
│      - Attempt 2: wait 2s                                   │
│      - Attempt 3: wait 4s                                   │
│      - Attempt 4: wait 8s                                   │
│    • Log entity_state_start on send                         │
│    • Update high_watermark on success (if enabled)          │
└─────┬───────────────────────────────────────────────────────┘
      │
      ▼
┌─────────────────────────────────────────────────────────────┐
│ 6. ERROR HANDLING PHASE                                      │
│    • Classify HTTP response:                                 │
│      - 200: Success → update watermark, log success         │
│      - 429, 502-504: Retryable → exponential backoff        │
│      - 401, 403, 404, 400: Non-retryable → fail fast       │
│    • Log failures to entity_state with error_message        │
│    • Raise appropriate exception (Retryable/NonRetryable)   │
└─────┬───────────────────────────────────────────────────────┘
      │
      ▼
┌─────────────────────────────────────────────────────────────┐
│ 7. COMPLETION PHASE                                          │
│    • Check all_batches_successful(run_id)                   │
│    • Update failed_batch_count                              │
│    • Complete pipeline_run:                                  │
│      - Set bronze_ended_at timestamp                        │
│      - Update pipeline_status (success/failed/partial)      │
│    • Close NiFi sessions                                     │
└─────┬───────────────────────────────────────────────────────┘
      │
      ▼
┌────────────┐
│  DAG End   │
└────────────┘

Component Interaction Diagram

┌─────────────────────────────────────────────────────────────────────┐
│                         Your Airflow DAG                             │
└───────────────────────────┬─────────────────────────────────────────┘
                            │
                            │ Creates & Configures
                            ▼
            ┌───────────────────────────────┐
            │      NiFiConfig               │
            │  • url, username, password    │◄────────┐
            │  • auth_mode (nifi/oidc)      │         │
            │  • retry settings             │         │ Uses
            └───────────────────────────────┘         │
                            │                         │
                            │ Passed to               │
                            ▼                         │
            ┌───────────────────────────────┐         │
            │  NiFiConnectionManager        │         │
            │  • get_access_token()         │         │
            │  • create_session()           │         │
            │  • health_check()             │         │
            │  • Thread-local sessions      │         │
            └───────────────┬───────────────┘         │
                            │                         │
                            │ Provides Session        │
                            ▼                         │
            ┌───────────────────────────────┐         │
            │       NiFiProcessor           │         │
            │  • process()                  │─────────┘
            │  • _send_http_request()       │
            │  • _handle_response()         │◄────────┐
            │  • _log_state()               │         │
            └───────────────┬───────────────┘         │
                            │                         │
                            │ Updates State           │
                            ▼                         │
            ┌───────────────────────────────┐         │
            │  PipelineStateManager         │         │ Reads State
            │  • register_pipeline()        │         │
            │  • start_pipeline_run()       │         │
            │  • log_entity_state_start()   │         │
            │  • update_high_watermark()    │         │
            │  • complete_pipeline_run()    │         │
            └───────────────┬───────────────┘         │
                            │                         │
                            │ Persists to             │
                            ▼                         │
            ┌───────────────────────────────┐         │
            │      PostgreSQL               │         │
            │  • pipeline_data              │         │
            │  • pipeline_run               │         │
            │  • entity_state               │         │
            └───────────────────────────────┘         │
                                                      │
            ┌───────────────────────────────┐         │
            │      SourceConfig             │         │
            │  • create_source_config()     │         │
            │  • calculate time_range       │─────────┘
            │  • generate SQL queries       │
            │  • partition batches          │
            └───────────────────────────────┘

---

Installation

Requirements

• Python: 3.8 or higher
• Apache Airflow: 2.5 or higher
• PostgreSQL: Any recent version (for state management)
• Apache NiFi: 1.15 or higher

Install via pip

pip install airflow-nifi-pipeline-utils

Install from source

git clone https://github.com/tansandil/airflow-nifi-pipeline-utils.git
cd airflow-nifi-pipeline-utils
pip install -e .

Dependencies

The library automatically installs:

apache-airflow>=2.5.0
requests>=2.20.0
nipyapi==0.22.0
psycopg2-binary==2.9.10
urllib3>=1.26.0

---

Quick Start

1. Database Setup

Run the provided DDL to create state tables:

# Download the schema
wget https://raw.githubusercontent.com/tansandil/airflow-nifi-pipeline-utils/main/pipeline_info_ddl.sql

# Apply to your PostgreSQL database
psql -h your-host -U your-user -d your-database -f pipeline_info_ddl.sql

This creates:

• state.pipeline_data - Pipeline registry
• state.pipeline_run - Run tracking with load_type
• state.entity_state - Batch-level state with watermarks

2. Configure Airflow Connection

Add PostgreSQL connection in Airflow:

# Via Airflow UI: Admin → Connections → Add
Connection Id: postgres_pipeline_state
Connection Type: Postgres
Host: your-postgres-host
Schema: your-database
Login: your-username
Password: your-password
Port: 5432

3. Basic Incremental Load Example

from datetime import datetime
from airflow import DAG
from airflow.operators.python import PythonOperator
from airflow_nifi_pipeline_utils import (
    NiFiConfig,
    SourceConfig,
    NiFiConnectionManager,
    NiFiProcessor,
    PipelineStateManager
)

def run_pipeline(**context):
    # 1. Configure NiFi connection
    nifi_config = NiFiConfig(
        url="https://nifi.example.com:8443",
        username="admin",
        password="your-password",
        auth_mode="nifi"
    )
    
    # 2. Initialize state manager
    state_mgr = PipelineStateManager(
        postgres_conn_id='postgres_pipeline_state',
        state_schema='state'
    )
    
    # 3. Register pipeline and start run
    pipeline_name = state_mgr.register_pipeline('orders')
    pipeline_id, run_id = state_mgr.start_pipeline_run(
        pipeline_name=pipeline_name,
        load_type='incremental'
    )
    
    # 4. Get last watermark
    last_watermark = state_mgr.get_latest_high_watermark(pipeline_name)
    if not last_watermark:
        last_watermark = '2024-01-01 00:00:00'
    
    # 5. Configure source table
    source_config = SourceConfig.create_source_config({
        'table_name': 'orders',
        'schema_name': 'sales',
        'pipeline_name': pipeline_name,
        'run_id': run_id,
        'load_type': 'incremental',
        'bookmark_column': 'updated_at',
        'last_watermark': last_watermark,
        'partition_hours': 8,
        'incremental_window': {'days': 1, 'hours': 0}
    })
    
    # 6. Generate batches (simplified - you'll add logic here)
    batches = generate_batches(source_config)  # Your implementation
    
    # 7. Process all batches
    with NiFiConnectionManager(nifi_config) as nifi_manager:
        for batch_dict in batches:
            processor = NiFiProcessor(
                batch_data_dict=batch_dict,
                config=nifi_config,
                postgres_conn_id='postgres_pipeline_state',
                state_schema='state',
                nifi_manager=nifi_manager
            )
            result = processor.process()
            print(f"Batch {batch_dict['partition_key']}: {result['status_code']}")
    
    # 8. Complete pipeline run
    state_mgr.complete_pipeline_run(pipeline_id, run_id, 'bronze', 'success')

# Create DAG
with DAG(
    dag_id='orders_incremental_load',
    start_date=datetime(2024, 1, 1),
    schedule_interval='@daily',
    catchup=False,
    tags=['data-ingestion', 'incremental']
) as dag:
    
    load_task = PythonOperator(
        task_id='load_orders',
        python_callable=run_pipeline
    )

4. Run Your DAG

# Test locally
airflow dags test orders_incremental_load 2024-01-01

# Deploy and trigger
airflow dags trigger orders_incremental_load

---

Core Concepts

1. Load Types

The library supports four fundamental data loading patterns:

Load Type	Purpose	Watermark Updates	Time Filtering	Use Case
incremental	Regular scheduled loads	Yes	Based on bookmark_column	Daily/hourly ingestion
full	Complete table reload (truncates)	No	None	Dimension tables, snapshots
reconcile	Backfill specific ranges	Yes/No	reconcile_start to reconcile_end	Data fixes, gap filling
historical (2.1.x)	Append-only bounded history — batches carry load_mode='historical' so NiFi routes them past the truncate	No (forced)	historical_start to historical_end	History backfill behind live incrementals; "full → incremental" composition

2. Watermark Management

High Watermark is the timestamp that marks the boundary of successfully processed data.

Timeline: ──────────────────────────────────────────────────►
          2024-01-01        2024-01-02        2024-01-03
          
Run 1:    [=========]
          Start: 2024-01-01 00:00:00
          End:   2024-01-01 08:00:00
          Success → high_watermark = 2024-01-01 08:00:00

Run 2:                [=========]
          Start: 2024-01-01 08:00:00  ← Picks up from last watermark
          End:   2024-01-01 16:00:00
          Success → high_watermark = 2024-01-01 16:00:00

Run 3 (Reconcile):    [=====]  ← Doesn't update watermark
          Start: 2024-01-01 02:00:00
          End:   2024-01-01 06:00:00
          Success → high_watermark unchanged (still 2024-01-01 16:00:00)

Key Rules:

• Watermark updates only on success when update_watermark=True
• Incremental loads: update_watermark=True (default)
• Full/Reconcile loads: update_watermark=False (default)
• Failed batches never update watermarks

3. Batch Partitioning

Large time ranges are split into manageable batches for parallel processing:

# Example: Process 3 days of data in 8-hour batches
source_config = SourceConfig.create_source_config({
    'last_watermark': '2024-01-01 00:00:00',
    'incremental_window': {'days': 3, 'hours': 0},
    'partition_hours': 8
})

# Generates 9 batches:
# Batch 1: 2024-01-01 00:00:00 → 2024-01-01 08:00:00
# Batch 2: 2024-01-01 08:00:00 → 2024-01-01 16:00:00
# Batch 3: 2024-01-01 16:00:00 → 2024-01-02 00:00:00
# ... (9 total batches)

Benefits:

• Enables parallel processing
• Reduces memory footprint
• Allows granular retry on failure
• Improves monitoring visibility

4. State Tracking Hierarchy

pipeline_data (Pipeline Registry)
    ├── pipeline_id: 1
    ├── pipeline_name: "orders_ingestion"
    └── description: "Daily order ingestion from Oracle"
            │
            └── pipeline_run (Execution Tracking)
                    ├── run_id: 123
                    ├── pipeline_id: 1
                    ├── load_type: "incremental"
                    ├── started_at: 2024-01-10 08:00:00
                    └── pipeline_status: "running"
                            │
                            └── entity_state (Batch Tracking)
                                    ├── batch_id: 456
                                    ├── run_id: 123
                                    ├── partition_key: "2024-01-10_00:00"
                                    ├── starting_watermark: 2024-01-10 00:00:00
                                    ├── end_watermark: 2024-01-10 08:00:00
                                    ├── high_watermark: 2024-01-10 08:00:00
                                    └── bronze_state_status: "success"

5. Thread Safety

The library uses thread-local storage for HTTP sessions:

# Each thread gets its own session
_local = threading.local()

def get_session():
    if not hasattr(_local, 'session'):
        _local.session = create_new_session()
    return _local.session

Why? Enables concurrent batch processing without session conflicts.

6. Error Classification

HTTP Response
     │
     ├─► 200 OK
     │       └─► Success: Update watermark, log success
     │
     ├─► 429, 502, 503, 504
     │       └─► RetryableNiFiError: Exponential backoff, retry
     │
     └─► 401, 403, 404, 400
             └─► NonRetryableNiFiError: Fail fast, alert ops team

---

Load Patterns

Pattern 1: Incremental Load

Use Case: Regular scheduled ingestion of new/updated records.

How It Works:

1. Retrieves last high_watermark from entity_state
2. Calculates end time: last_watermark + incremental_window
3. Generates SQL: WHERE bookmark_column >= last_watermark AND bookmark_column < end_time
4. Partitions time range into batches
5. On success, updates high_watermark to end_watermark

Configuration:

source_config = SourceConfig.create_source_config({
    'table_name': 'transactions',
    'schema_name': 'finance',
    'pipeline_name': 'transactions_ingestion',
    'run_id': 123,
    'load_type': 'incremental',
    
    # Required for incremental
    'bookmark_column': 'transaction_date',
    'last_watermark': '2024-01-10 00:00:00',
    'partition_hours': 8,
    'incremental_window': {'days': 1, 'hours': 0}
})

# Automatically sets update_watermark=True

Generated SQL:

SELECT t.*, 'transactions' AS table_name, 'finance' AS schema_name 
FROM finance.transactions t 
WHERE t.transaction_date >= '2024-01-10 00:00:00' 
  AND t.transaction_date < '2024-01-11 00:00:00'

Batch Example:

Time Range: 2024-01-10 00:00:00 → 2024-01-11 00:00:00 (24 hours)
Partition: 8 hours

Batch 1: 2024-01-10 00:00:00 → 2024-01-10 08:00:00
Batch 2: 2024-01-10 08:00:00 → 2024-01-10 16:00:00
Batch 3: 2024-01-10 16:00:00 → 2024-01-11 00:00:00

Each batch processed independently, watermark updated per batch

Pattern 2: Full Load

Use Case: Complete table refresh, monthly snapshots, dimension tables.

How It Works:

1. No watermark lookup
2. No time filtering in SQL
3. Loads entire table
4. Does NOT update high_watermark
5. Can run in parallel with incremental loads

Configuration:

source_config = SourceConfig.create_source_config({
    'table_name': 'customers',
    'schema_name': 'crm',
    'pipeline_name': 'customers_full_load',
    'run_id': 124,
    'load_type': 'full',
    
    # Optional: partition for large tables
    'partition_hours': None  # Single batch
})

# Automatically sets update_watermark=False

Generated SQL:

SELECT t.*, 'customers' AS table_name, 'crm' AS schema_name 
FROM crm.customers t

Use Cases:

• Monthly dimension table snapshots
• Initial load of lookup tables
• Complete refresh after schema changes
• Parallel full loads for testing

Pattern 3: Reconcile Load

Use Case: Backfill missing data, fix data quality issues, fill gaps.

How It Works:

1. No watermark lookup
2. Time filtering based on reconcile_start and reconcile_end
3. Generates SQL with specified date range
4. Partitions large backfills
5. Does NOT update high_watermark (doesn't interfere with incremental)

Configuration:

source_config = SourceConfig.create_source_config({
    'table_name': 'orders',
    'schema_name': 'sales',
    'pipeline_name': 'orders_reconcile',
    'run_id': 125,
    'load_type': 'reconcile',
    
    # Required for reconcile
    'bookmark_column': 'order_date',
    'reconcile_start': '2024-01-01 00:00:00',
    'reconcile_end': '2024-01-07 23:59:59',
    'partition_hours': 12
})

# Automatically sets update_watermark=False

Generated SQL:

SELECT t.*, 'orders' AS table_name, 'sales' AS schema_name 
FROM sales.orders t 
WHERE t.order_date >= '2024-01-01 00:00:00' 
  AND t.order_date < '2024-01-07 23:59:59'

Real-World Scenario:

Problem: Discovered missing data in January due to source system issue

Solution:
  Day 1: Run reconcile load for January
    ├─► Does NOT affect ongoing incremental loads
    ├─► high_watermark remains at current date (e.g., Feb 10)
    └─► Incremental continues processing new data

  Day 2: Incremental load continues normally
    ├─► Picks up from Feb 10 watermark
    └─► Not affected by yesterday's reconcile

Batch Example:

Time Range: 2024-01-01 00:00:00 → 2024-01-07 23:59:59 (7 days)
Partition: 12 hours

Batch 1:  2024-01-01 00:00:00 → 2024-01-01 12:00:00
Batch 2:  2024-01-01 12:00:00 → 2024-01-02 00:00:00
Batch 3:  2024-01-02 00:00:00 → 2024-01-02 12:00:00
...
Batch 14: 2024-01-07 12:00:00 → 2024-01-07 23:59:59

Total: 14 batches, none update high_watermark

---

🔧 Components Reference

NiFiConfig

Purpose: Centralized NiFi connection configuration with validation.

Parameters:

Parameter	Type	Required	Default	Description
url	str	Yes	-	NiFi base URL (e.g., https://nifi:8443)
username	str	Yes	-	Username for authentication
password	str	Yes	-	Password for authentication
auth_mode	str	No	'nifi'	Auth mode: 'nifi' or 'oidc'
keycloak_token_url	str	No	None	Keycloak token endpoint (required if auth_mode='oidc')
keycloak_client_id	str	No	None	Keycloak client ID (required if auth_mode='oidc')
keycloak_client_secret	str	No	None	Keycloak client secret (required if auth_mode='oidc')
max_attempts	int	No	4	Maximum retry attempts per request
base_delay	int	No	2	Base delay in seconds for exponential backoff
connection_timeout	int	No	30	Connection timeout in seconds
retryable_status_codes	set	No	{429, 502, 503, 504}	HTTP codes that trigger retries
non_retryable_status_codes	set	No	{401, 403, 404, 400, 409}	HTTP codes that fail fast

Properties:

config.api_url          # Full API URL: https://nifi:8443/nifi-api
config.auth_url         # Token endpoint: https://nifi:8443/nifi-api/access/token
config.is_oidc          # Boolean: True if auth_mode='oidc'
config.is_nifi_auth     # Boolean: True if auth_mode='nifi'

Example:

# Basic NiFi auth
config = NiFiConfig(
    url="https://nifi.prod.example.com:8443",
    username="nifi_service_account",
    password="secure_password"
)

# OIDC auth
config = NiFiConfig(
    url="https://nifi.prod.example.com:8443",
    username="user@company.com",
    password="user_password",
    auth_mode="oidc",
    keycloak_token_url="https://keycloak.company.com/realms/master/protocol/openid-connect/token",
    keycloak_client_id="nifi-client",
    keycloak_client_secret="client_secret"
)

SourceConfig

Purpose: Table configuration with automatic SQL generation and batch partitioning.

Parameters:

Parameter	Type	Required	Default	Description
table_name	str	Yes	-	Source table name
schema_name	str	Yes	-	Database schema name
pipeline_name	str	Yes	-	Pipeline identifier
run_id	int	Yes	-	Current run ID
load_type	str	No	'incremental'	Load type: 'incremental', 'full', 'reconcile'
bookmark_column	str	Conditional	None	Timestamp column (required for incremental/reconcile)
last_watermark	str	Conditional	None	Starting watermark (required for incremental)
partition_hours	int	Conditional	None	Hours per batch (required for incremental/reconcile)
incremental_window	dict	No	{'days': 1, 'hours': 0}	Time window for incremental loads
reconcile_start	str	Conditional	None	Start timestamp (required for reconcile)
reconcile_end	str	Conditional	None	End timestamp (required for reconcile)
update_watermark	bool	No	Auto	Whether to update watermark (auto-set based on load_type)
custom_sql	str	No	None	Override generated SQL

Properties:

config.time_range              # (start_datetime, end_datetime) or None
config.cleaned_last_watermark  # Parsed datetime object

Methods:

config.create_sql_query(start_time, end_time)  # Generate SQL for batch
config.format_timestamp(dt)                     # Format datetime for SQL

Example:

# Incremental
config = SourceConfig.create_source_config({
    'table_name': 'orders',
    'schema_name': 'sales',
    'pipeline_name': 'orders_ingestion',
    'run_id': 123,
    'load_type': 'incremental',
    'bookmark_column': 'updated_at',
    'last_watermark': '2024-01-10 00:00:00',
    'partition_hours': 8,
    'incremental_window': {'days': 1, 'hours': 0}
})

# Full
config = SourceConfig.create_source_config({
    'table_name': 'customers',
    'schema_name': 'crm',
    'pipeline_name': 'customers_full',
    'run_id': 124,
    'load_type': 'full'
})

# Reconcile
config = SourceConfig.create_source_config({
    'table_name': 'transactions',
    'schema_name': 'finance',
    'pipeline_name': 'transactions_reconcile',
    'run_id': 125,
    'load_type': 'reconcile',
    'bookmark_column': 'transaction_date',
    'reconcile_start': '2024-01-01 00:00:00',
    'reconcile_end': '2024-01-07 23:59:59',
    'partition_hours': 12
})

BatchDataResult

Purpose: Validated batch metadata for NiFi processing.

Parameters:

Parameter	Type	Required	Description
pipeline_name	str	Yes	Pipeline identifier
run_id	int	Yes	Current run ID
entity_name	str	Yes	Table/entity name
partition_key	str	Yes	Unique batch identifier
sql_text	str	Yes	SQL query for this batch
start_ts	str	Conditional	Start timestamp (None for full loads)
end_ts	str	Conditional	End timestamp (None for full loads)
schema_name	str	Yes	Database schema
nifi_host	str	Yes	NiFi ListenHTTP hostname
nifi_port	int	Yes	NiFi ListenHTTP port
update_watermark	bool	No	Whether to update watermark on success

Properties:

batch.endpoint_url      # Constructed NiFi URL: http://host:port/contentListener
batch.log_identifier    # Formatted string for logging

Example:

batch_dict = {
    'pipeline_name': 'orders_ingestion',
    'run_id': 123,
    'entity_name': 'orders',
    'partition_key': '2024-01-10_00:00',
    'sql_text': 'SELECT * FROM sales.orders WHERE ...',
    'start_ts': '2024-01-10 00:00:00',
    'end_ts': '2024-01-10 08:00:00',
    'schema_name': 'sales',
    'nifi_host': 'nifi-listener-01.prod.example.com',
    'nifi_port': 7001,
    'update_watermark': True
}

batch = BatchData.create_batch_data(batch_dict)

NiFiProcessor

Purpose: Main batch processing engine with retry logic and state tracking.

Parameters:

Parameter	Type	Required	Description
batch_data_dict	dict	Yes	Batch metadata dictionary
config	NiFiConfig	Yes	NiFi configuration object
postgres_conn_id	str	Yes	Airflow Postgres connection ID
state_schema	str	Yes	Schema for state tables
nifi_manager	NiFiConnectionManager	Yes	Connection manager instance

Methods:

processor.process()  # Main processing method, returns result dict

Return Value:

{
    'status_code': 200,
    'response_text': 'Transfer Complete',
    'entity_name': 'orders',
    'run_id': 123,
    'attempts': 1,
    'partition_key': '2024-01-10_00:00',
    'start_date': '2024-01-10 00:00:00',
    'end_date': '2024-01-10 08:00:00',
    'sql_text': 'SELECT * FROM ...',
    'endpoint_url': 'http://nifi:7001/contentListener',
    'timestamp': '2024-01-10T10:30:45.123456',
    'nifi_config_url': 'https://nifi:8443'
}

Example:

with NiFiConnectionManager(config) as nifi_manager:
    processor = NiFiProcessor(
        batch_data_dict=batch_dict,
        config=nifi_config,
        postgres_conn_id='postgres_pipeline_state',
        state_schema='state',
        nifi_manager=nifi_manager
    )
    
    try:
        result = processor.process()
        print(f"Success: {result['status_code']}")
    except RetryableNiFiError as e:
        print(f"Retryable error: {e}")
        raise  # Let Airflow retry
    except NonRetryableNiFiError as e:
        print(f"Permanent error: {e.error_code}")
        send_alert(e)
        raise

PipelineStateManager

Purpose: PostgreSQL-based state tracking with watermark management.

Parameters:

Parameter	Type	Required	Default	Description
postgres_conn_id	str	Yes	-	Airflow Postgres connection ID
state_schema	str	No	'state'	Schema containing state tables

Methods:

# Pipeline registration
pipeline_name = manager.register_pipeline(table_name: str) -> str

# Run management
pipeline_id, run_id = manager.start_pipeline_run(
    pipeline_name: str, 
    load_type: str = 'incremental'
) -> tuple

# Watermark management
last_wm = manager.get_latest_high_watermark(pipeline_name: str) -> str
manager.update_high_watermark(pipeline_name: str, partition_key: str, new_watermark: str)

# Batch tracking
manager.log_entity_state_start(
    pipeline_name: str,
    run_id: int,
    entity_name: str,
    partition_key: str,
    sql_text: str,
    start_ts: str,
    end_ts: str,
    high_watermark: str = None,
    status: str = 'success'
)

# Run completion
manager.complete_pipeline_run(
    pipeline_id: int, 
    run_id: int, 
    layer: str, 
    status: str = 'success'
)

# Status checks
all_success = manager.all_batches_successful(run_id: int) -> bool
run_status = manager.get_run_status(run_id: int) -> dict

Example:

state_mgr = PipelineStateManager(
    postgres_conn_id='postgres_pipeline_state',
    state_schema='state'
)

# Register and start
pipeline_name = state_mgr.register_pipeline('orders')
pipeline_id, run_id = state_mgr.start_pipeline_run(
    pipeline_name=pipeline_name,
    load_type='incremental'
)

# Get last watermark
last_wm = state_mgr.get_latest_high_watermark(pipeline_name)

# ... process batches ...

# Check success
if state_mgr.all_batches_successful(run_id):
    state_mgr.complete_pipeline_run(pipeline_id, run_id, 'bronze', 'success')
else:
    state_mgr.complete_pipeline_run(pipeline_id, run_id, 'bronze', 'failed')

NiFiConnectionManager

Purpose: Thread-safe HTTP session management with authentication.

Parameters:

Parameter	Type	Required	Description
config	NiFiConfig	Yes	NiFi configuration object

Methods:

session = manager.get_session()          # Get thread-local session
manager.refresh_token()                   # Force token refresh
manager.close_session()                   # Close thread-local session
is_healthy = manager.health_check(url)   # Endpoint health check

Example:

# Context manager (recommended)
with NiFiConnectionManager(config) as manager:
    session = manager.get_session()
    response = session.post(url, data=data)

# Manual management
manager = NiFiConnectionManager(config)
try:
    session = manager.get_session()
    response = session.post(url, data=data)
finally:
    manager.close_session()

---

Complete Examples

Example 1: Daily Incremental Load

from datetime import datetime, timedelta
from airflow import DAG
from airflow.operators.python import PythonOperator
from airflow_nifi_pipeline_utils import *

def generate_batches(source_config, pipeline_name, run_id, nifi_host, nifi_port):
    """Generate time-partitioned batches"""
    batches = []
    time_range = source_config.time_range
    
    if not time_range:
        # Full load - single batch
        batch_dict = {
            'pipeline_name': pipeline_name,
            'run_id': run_id,
            'entity_name': source_config.table_name,
            'partition_key': 'full_load',
            'sql_text': source_config.create_sql_query(),
            'start_ts': None,
            'end_ts': None,
            'schema_name': source_config.schema_name,
            'nifi_host': nifi_host,
            'nifi_port': nifi_port,
            'update_watermark': False
        }
        return [batch_dict]
    
    # Incremental/Reconcile - partition by hours
    start_time, end_time = time_range
    current_time = start_time
    partition_delta = timedelta(hours=source_config.partition_hours)
    
    while current_time < end_time:
        batch_end = min(current_time + partition_delta, end_time)
        
        batch_dict = {
            'pipeline_name': pipeline_name,
            'run_id': run_id,
            'entity_name': source_config.table_name,
            'partition_key': current_time.strftime('%Y-%m-%d_%H:%M'),
            'sql_text': source_config.create_sql_query(current_time, batch_end),
            'start_ts': current_time.strftime('%Y-%m-%d %H:%M:%S'),
            'end_ts': batch_end.strftime('%Y-%m-%d %H:%M:%S'),
            'schema_name': source_config.schema_name,
            'nifi_host': nifi_host,
            'nifi_port': nifi_port,
            'update_watermark': source_config.update_watermark
        }
        batches.append(batch_dict)
        current_time = batch_end
    
    return batches

def run_incremental_pipeline(**context):
    # Configuration
    nifi_config = NiFiConfig(
        url="https://nifi.prod.example.com:8443",
        username="nifi_admin",
        password="secure_password",
        max_attempts=5,
        base_delay=2
    )
    
    state_mgr = PipelineStateManager(
        postgres_conn_id='postgres_pipeline_state',
        state_schema='state'
    )
    
    # Pipeline setup
    pipeline_name = state_mgr.register_pipeline('transactions')
    pipeline_id, run_id = state_mgr.start_pipeline_run(
        pipeline_name=pipeline_name,
        load_type='incremental'
    )
    
    # Get last watermark
    last_watermark = state_mgr.get_latest_high_watermark(pipeline_name)
    if not last_watermark:
        last_watermark = '2024-01-01 00:00:00'  # Initial load
    
    print(f"Starting incremental load from watermark: {last_watermark}")
    
    # Configure source
    source_config = SourceConfig.create_source_config({
        'table_name': 'transactions',
        'schema_name': 'finance',
        'pipeline_name': pipeline_name,
        'run_id': run_id,
        'load_type': 'incremental',
        'bookmark_column': 'transaction_date',
        'last_watermark': last_watermark,
        'partition_hours': 8,
        'incremental_window': {'days': 1, 'hours': 0}
    })
    
    # Generate batches
    batches = generate_batches(
        source_config=source_config,
        pipeline_name=pipeline_name,
        run_id=run_id,
        nifi_host='nifi-listener-finance.prod.example.com',
        nifi_port=7001
    )
    
    print(f"Generated {len(batches)} batches to process")
    
    # Process all batches
    successful_batches = 0
    failed_batches = 0
    
    with NiFiConnectionManager(nifi_config) as nifi_manager:
        for batch_dict in batches:
            try:
                processor = NiFiProcessor(
                    batch_data_dict=batch_dict,
                    config=nifi_config,
                    postgres_conn_id='postgres_pipeline_state',
                    state_schema='state',
                    nifi_manager=nifi_manager
                )
                
                result = processor.process()
                print(f"Batch {batch_dict['partition_key']}: {result['status_code']}")
                successful_batches += 1
                
            except RetryableNiFiError as e:
                print(f"Retryable error for batch {batch_dict['partition_key']}: {e}")
                failed_batches += 1
                raise  # Let Airflow retry the entire task
                
            except NonRetryableNiFiError as e:
                print(f"Permanent error for batch {batch_dict['partition_key']}: {e.error_code}")
                failed_batches += 1
                # Continue processing other batches, mark run as partial failure
    
    # Complete pipeline run
    if failed_batches == 0:
        state_mgr.complete_pipeline_run(pipeline_id, run_id, 'bronze', 'success')
        print(f"Pipeline completed successfully: {successful_batches} batches")
    elif successful_batches > 0:
        state_mgr.complete_pipeline_run(pipeline_id, run_id, 'bronze', 'partial')
        print(f"Partial success: {successful_batches} succeeded, {failed_batches} failed")
    else:
        state_mgr.complete_pipeline_run(pipeline_id, run_id, 'bronze', 'failed')
        print(f"Pipeline failed: all {failed_batches} batches failed")
        raise Exception("All batches failed")

# DAG definition
with DAG(
    dag_id='finance_transactions_incremental',
    description='Daily incremental load of transactions from Oracle to Data Lake',
    start_date=datetime(2024, 1, 1),
    schedule_interval='0 8 * * *',  # Daily at 8 AM
    catchup=False,
    max_active_runs=1,
    tags=['finance', 'incremental', 'oracle', 'production']
) as dag:
    
    load_task = PythonOperator(
        task_id='load_transactions',
        python_callable=run_incremental_pipeline,
        retries=3,
        retry_delay=timedelta(minutes=5)
    )

Example 2: Monthly Full Load

def run_full_load_pipeline(**context):
    # Configuration
    nifi_config = NiFiConfig(
        url="https://nifi.prod.example.com:8443",
        username="nifi_admin",
        password="secure_password"
    )
    
    state_mgr = PipelineStateManager(
        postgres_conn_id='postgres_pipeline_state',
        state_schema='state'
    )
    
    # Pipeline setup
    pipeline_name = state_mgr.register_pipeline('customer_master')
    pipeline_id, run_id = state_mgr.start_pipeline_run(
        pipeline_name=pipeline_name,
        load_type='full'  # Full load - no watermark updates
    )
    
    print(f"Starting full load for customer master table")
    
    # Configure source
    source_config = SourceConfig.create_source_config({
        'table_name': 'customer_master',
        'schema_name': 'crm',
        'pipeline_name': pipeline_name,
        'run_id': run_id,
        'load_type': 'full'  # No bookmark column, no watermark
    })
    
    # Generate single batch (full table)
    batches = generate_batches(
        source_config=source_config,
        pipeline_name=pipeline_name,
        run_id=run_id,
        nifi_host='nifi-listener-crm.prod.example.com',
        nifi_port=7002
    )
    
    # Process batch
    with NiFiConnectionManager(nifi_config) as nifi_manager:
        batch_dict = batches[0]
        processor = NiFiProcessor(
            batch_data_dict=batch_dict,
            config=nifi_config,
            postgres_conn_id='postgres_pipeline_state',
            state_schema='state',
            nifi_manager=nifi_manager
        )
        
        result = processor.process()
        print(f"Full load completed: {result['status_code']}")
    
    # Complete run
    state_mgr.complete_pipeline_run(pipeline_id, run_id, 'bronze', 'success')
    print("Full load pipeline completed")

# DAG for monthly full load
with DAG(
    dag_id='crm_customer_master_full_load',
    description='Monthly full refresh of customer master dimension',
    start_date=datetime(2024, 1, 1),
    schedule_interval='0 2 1 * *',  # 1st of month at 2 AM
    catchup=False,
    tags=['crm', 'full-load', 'dimension']
) as dag:
    
    full_load_task = PythonOperator(
        task_id='full_load_customers',
        python_callable=run_full_load_pipeline
    )

Example 3: Ad-Hoc Reconcile Load

def run_reconcile_pipeline(**context):
    # Get parameters from Airflow UI or config
    reconcile_start = context['dag_run'].conf.get('reconcile_start', '2024-01-01 00:00:00')
    reconcile_end = context['dag_run'].conf.get('reconcile_end', '2024-01-07 23:59:59')
    
    # Configuration
    nifi_config = NiFiConfig(
        url="https://nifi.prod.example.com:8443",
        username="nifi_admin",
        password="secure_password"
    )
    
    state_mgr = PipelineStateManager(
        postgres_conn_id='postgres_pipeline_state',
        state_schema='state'
    )
    
    # Pipeline setup
    pipeline_name = state_mgr.register_pipeline('orders_reconcile')
    pipeline_id, run_id = state_mgr.start_pipeline_run(
        pipeline_name=pipeline_name,
        load_type='reconcile'  # Backfill mode
    )
    
    print(f"Starting reconcile load: {reconcile_start} → {reconcile_end}")
    
    # Configure source
    source_config = SourceConfig.create_source_config({
        'table_name': 'orders',
        'schema_name': 'sales',
        'pipeline_name': pipeline_name,
        'run_id': run_id,
        'load_type': 'reconcile',
        'bookmark_column': 'order_date',
        'reconcile_start': reconcile_start,
        'reconcile_end': reconcile_end,
        'partition_hours': 12  # 12-hour batches for backfill
    })
    
    # Generate batches
    batches = generate_batches(
        source_config=source_config,
        pipeline_name=pipeline_name,
        run_id=run_id,
        nifi_host='nifi-listener-sales.prod.example.com',
        nifi_port=7000
    )
    
    print(f"Reconcile will process {len(batches)} batches")
    
    # Process all batches
    successful = 0
    with NiFiConnectionManager(nifi_config) as nifi_manager:
        for batch_dict in batches:
            processor = NiFiProcessor(
                batch_data_dict=batch_dict,
                config=nifi_config,
                postgres_conn_id='postgres_pipeline_state',
                state_schema='state',
                nifi_manager=nifi_manager
            )
            
            result = processor.process()
            print(f"Batch {batch_dict['partition_key']}: {result['status_code']}")
            successful += 1
    
    # Complete run
    state_mgr.complete_pipeline_run(pipeline_id, run_id, 'bronze', 'success')
    print(f"Reconcile completed: {successful} batches processed")
    print(f"Note: Incremental watermarks not affected by this reconcile")

# Manual trigger DAG
with DAG(
    dag_id='sales_orders_reconcile',
    description='Ad-hoc reconciliation for orders table',
    start_date=datetime(2024, 1, 1),
    schedule_interval=None,  # Manual trigger only
    catchup=False,
    tags=['sales', 'reconcile', 'backfill']
) as dag:
    
    reconcile_task = PythonOperator(
        task_id='reconcile_orders',
        python_callable=run_reconcile_pipeline
    )

# Trigger via Airflow CLI:
# airflow dags trigger sales_orders_reconcile \
#   --conf '{"reconcile_start": "2024-01-01 00:00:00", "reconcile_end": "2024-01-31 23:59:59"}'

Example 4: OIDC Authentication

from airflow.models import Variable

def run_with_oidc_auth(**context):
    # OIDC configuration (Keycloak)
    nifi_config = NiFiConfig(
        url="https://nifi.prod.example.com:8443",
        username=Variable.get("keycloak_username"),  # user@company.com
        password=Variable.get("keycloak_password"),
        auth_mode="oidc",
        keycloak_token_url="https://keycloak.company.com/realms/production/protocol/openid-connect/token",
        keycloak_client_id="nifi-airflow-client",
        keycloak_client_secret=Variable.get("keycloak_client_secret")
    )
    
    # Rest of pipeline logic is identical
    state_mgr = PipelineStateManager(
        postgres_conn_id='postgres_pipeline_state',
        state_schema='state'
    )
    
    # ... same as incremental load example ...

with DAG(
    dag_id='enterprise_pipeline_with_oidc',
    start_date=datetime(2024, 1, 1),
    schedule_interval='@daily',
    tags=['enterprise', 'oidc']
) as dag:
    
    load_task = PythonOperator(
        task_id='load_with_oidc',
        python_callable=run_with_oidc_auth
    )

---

Best Practices

1. Credential Management

Never hardcode credentials:

# BAD
config = NiFiConfig(
    url="https://nifi:8443",
    username="admin",
    password="password123"
)

Use Airflow Variables or Secrets Backend:

# GOOD
from airflow.models import Variable

config = NiFiConfig(
    url=Variable.get("nifi_url"),
    username=Variable.get("nifi_username"),
    password=Variable.get("nifi_password", deserialize_json=False)
)

# BETTER - Secrets Backend
from airflow.hooks.base import BaseHook

connection = BaseHook.get_connection('nifi_prod')
config = NiFiConfig(
    url=connection.host,
    username=connection.login,
    password=connection.password
)

2. Context Managers

Manual session management:

# BAD
manager = NiFiConnectionManager(config)
processor = NiFiProcessor(...)
result = processor.process()
manager.close_session()  # Might not execute if error occurs

Use context managers:

# GOOD
with NiFiConnectionManager(config) as manager:
    processor = NiFiProcessor(..., nifi_manager=manager)
    result = processor.process()
# Session automatically closed

3. Error Handling Strategy

Distinguish error types:

from airflow_nifi_pipeline_utils import RetryableNiFiError, NonRetryableNiFiError
import logging

try:
    result = processor.process()
except RetryableNiFiError as e:
    # Temporary failure - let Airflow retry
    logging.warning(f"Retryable error: {e}")
    raise  # Airflow will retry based on task config
    
except NonRetryableNiFiError as e:
    # Permanent failure - alert and fail fast
    logging.error(f"Permanent failure: {e.error_code}")
    send_slack_alert(f"Pipeline failed: {e}")
    send_email_alert(e)
    raise  # Fail the task immediately

4. Partition Size Tuning

Choose partition_hours based on data volume:

# Small tables (<100K rows/day)
partition_hours = 24  # Single daily batch

# Medium tables (100K-1M rows/day)
partition_hours = 8   # 3 batches per day

# Large tables (>1M rows/day)
partition_hours = 4   # 6 batches per day

# Very large tables (>10M rows/day)
partition_hours = 1   # Hourly batches

Consider:

• Network bandwidth
• NiFi processor capacity
• Source database query performance
• Memory constraints

5. Load Type Selection

Scenario	Use This Load Type	Why
Daily scheduled loads	incremental	Processes only new data, updates watermark
Dimension table refresh	full	Complete reload, doesn't affect incremental
Fix missing data	reconcile	Backfill without affecting ongoing loads
Initial pipeline setup	full first, then incremental	Load historical, then switch to incremental
Late-arriving data	incremental with lookback window	Capture delayed records

6. Monitoring & Alerting

Query state tables for monitoring:

-- Check recent runs
SELECT 
    pr.run_id,
    pd.pipeline_name,
    pr.load_type,
    pr.started_at,
    pr.bronze_ended_at,
    pr.pipeline_status,
    pr.failed_batch_count,
    pr.bronze_run_duration
FROM state.pipeline_run pr
JOIN state.pipeline_data pd ON pr.pipeline_id = pd.pipeline_id
WHERE pr.started_at > NOW() - INTERVAL '24 hours'
ORDER BY pr.started_at DESC;

-- Find failed batches
SELECT 
    es.entity_name,
    es.partition_key,
    es.bronze_state_status,
    es.error_message,
    es.started_at,
    pr.run_id,
    pr.load_type
FROM state.entity_state es
JOIN state.pipeline_run pr ON es.run_id = pr.run_id
WHERE es.bronze_state_status = 'failed'
  AND es.started_at > NOW() - INTERVAL '7 days'
ORDER BY es.started_at DESC;

-- Monitor watermark progression
SELECT 
    entity_name,
    MAX(high_watermark) as latest_watermark,
    COUNT(*) as total_batches,
    SUM(CASE WHEN bronze_state_status = 'success' THEN 1 ELSE 0 END) as successful_batches
FROM state.entity_state
WHERE is_current_run = TRUE
GROUP BY entity_name;

7. Performance Optimization

Enable parallel batch processing:

from concurrent.futures import ThreadPoolExecutor, as_completed

def process_batch_wrapper(batch_dict, config, nifi_manager):
    """Wrapper for thread-safe batch processing"""
    processor = NiFiProcessor(
        batch_data_dict=batch_dict,
        config=config,
        postgres_conn_id='postgres_pipeline_state',
        state_schema='state',
        nifi_manager=nifi_manager
    )
    return processor.process()

with NiFiConnectionManager(config) as nifi_manager:
    with ThreadPoolExecutor(max_workers=5) as executor:
        futures = {
            executor.submit(process_batch_wrapper, batch, config, nifi_manager): batch
            for batch in batches
        }
        
        for future in as_completed(futures):
            batch = futures[future]
            try:
                result = future.result()
                print(f"Batch {batch['partition_key']} completed")
            except Exception as e:
                print(f"Batch {batch['partition_key']} failed: {e}")

8. Testing Strategy

Test configuration validation:

import pytest
from airflow_nifi_pipeline_utils import NiFiConfig, SourceConfig

def test_nifi_config_validation():
    # Test invalid URL
    with pytest.raises(ValueError, match="URL must include protocol"):
        NiFiConfig(url="nifi:8443", username="admin", password="pass")
    
    # Test OIDC without required params
    with pytest.raises(ValueError, match="requires the following parameters"):
        NiFiConfig(
            url="https://nifi:8443",
            username="admin",
            password="pass",
            auth_mode="oidc"  # Missing keycloak params
        )

def test_source_config_validation():
    # Test incremental without required fields
    with pytest.raises(ValueError, match="bookmark_column is required"):
        SourceConfig.create_source_config({
            'table_name': 'orders',
            'schema_name': 'sales',
            'pipeline_name': 'test',
            'run_id': 1,
            'load_type': 'incremental'
            # Missing bookmark_column, last_watermark, partition_hours
        })

9. Logging Best Practices

Structured logging for better debugging:

import logging
import json

# Configure structured logging
logging.basicConfig(
    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s',
    level=logging.INFO
)

logger = logging.getLogger(__name__)

# Log with context
logger.info(json.dumps({
    'event': 'pipeline_started',
    'pipeline_name': pipeline_name,
    'run_id': run_id,
    'load_type': 'incremental',
    'watermark': last_watermark,
    'batch_count': len(batches)
}))

# Log batch completion
for batch in batches:
    logger.info(json.dumps({
        'event': 'batch_completed',
        'partition_key': batch['partition_key'],
        'status': 'success',
        'duration_seconds': duration,
        'record_count': record_count
    }))

10. Deployment Checklist

Before deploying to production:

• Database schema created (pipeline_info_ddl.sql)
• Airflow connections configured
• Credentials stored in Airflow Variables/Secrets
• NiFi ListenHTTP processors configured (one per schema/port)
• Monitoring queries set up
• Alerting configured (Slack/email)
• Test runs completed in dev environment
• Documentation updated with pipeline-specific details
• Runbook created for common issues
• Backup/recovery plan established

---

Troubleshooting

Problem 1: Authentication Failures (401/403)

Symptoms:

NonRetryableNiFiError: HTTP 401: Unauthorized

Diagnosis:

# Test authentication manually
from airflow_nifi_pipeline_utils import NiFiConfig, get_access_token

config = NiFiConfig(
    url="https://nifi:8443",
    username="your_username",
    password="your_password",
    auth_mode="nifi"
)

try:
    token = get_access_token(config)
    print(f"Authentication successful: {token[:20]}...")
except Exception as e:
    print(f"Authentication failed: {e}")

Solutions:

1. NiFi Native Auth:

   • Verify username/password in NiFi UI
   • Check NiFi logs: /opt/nifi/logs/nifi-user.log
   • Ensure user has appropriate permissions

2. OIDC/Keycloak:

   • Verify keycloak_token_url is correct
   • Check client_id and client_secret
   • Ensure user has NiFi access role in Keycloak
   • Test token endpoint manually:

     curl -X POST "https://keycloak.example.com/realms/master/protocol/openid-connect/token" \
       -d "grant_type=password" \
       -d "client_id=nifi-client" \
       -d "client_secret=secret" \
       -d "username=user@company.com" \
       -d "password=password"
     

Problem 2: Connection Timeouts

Symptoms:

RetryableNiFiError: Request failed after 4 attempts
ConnectionError: Connection refused

Diagnosis:

# Test connectivity
from airflow_nifi_pipeline_utils import NiFiConnectionManager

with NiFiConnectionManager(config) as manager:
    is_healthy = manager.health_check("http://nifi-listener:7001/contentListener")
    print(f"Health check: {'PASS' if is_healthy else 'FAIL'}")

Solutions:

1. Network Issues:

   • Ping NiFi host from Airflow worker
   • Check firewall rules
   • Verify DNS resolution
   • Test with curl: curl -v http://nifi-listener:7001/contentListener

2. NiFi Process Group Not Running:

   from airflow_nifi_pipeline_utils import safe_schedule_process_group, wait_for_process_group_ready
   
   # Start process group
   safe_schedule_process_group(pg_id='your-pg-id', scheduled=True)
   
   # Wait until ready
   is_ready = wait_for_process_group_ready(pg_id='your-pg-id', max_wait=120)
   

3. Increase Timeouts:

   config = NiFiConfig(
       url="https://nifi:8443",
       username="admin",
       password="pass",
       connection_timeout=60,  # Increase from default 30
       max_attempts=5          # Increase from default 4
   )
   

Problem 3: Watermark Not Updating

Symptoms:

• Pipeline runs successfully but watermark stays the same
• Incremental loads process same data repeatedly

Diagnosis:

-- Check entity_state for watermark updates
SELECT 
    partition_key,
    starting_watermark,
    end_watermark,
    high_watermark,
    bronze_state_status,
    is_current_run
FROM state.entity_state
WHERE entity_name = 'your_table'
ORDER BY started_at DESC
LIMIT 10;

-- Check if all batches succeeded
SELECT 
    COUNT(*) FILTER (WHERE bronze_state_status != 'success') as failed_count
FROM state.entity_state
WHERE run_id = 123;  -- Your run_id

Solutions:

1. Check update_watermark flag:

   # Incremental loads must have update_watermark=True
   batch_dict = {
       'update_watermark': True,  # Critical for incremental
       ...
   }
   

2. Verify all batches succeeded:

   if state_mgr.all_batches_successful(run_id):
       print("All batches successful, watermarks updated")
   else:
       print("Some batches failed, watermarks not updated")
   

3. Check load_type:

   -- Watermarks only update for incremental loads
   SELECT run_id, load_type, pipeline_status
   FROM state.pipeline_run
   WHERE run_id = 123;
   

Problem 4: Batch Processing Hangs

Symptoms:

• Task runs indefinitely without completing
• No error messages
• Some batches process, others hang

Diagnosis:

# Enable debug logging
import logging
logging.basicConfig(level=logging.DEBUG)

# Check NiFi queue status
from airflow_nifi_pipeline_utils import check_process_group_idle

is_idle = check_process_group_idle(pg_id='your-pg-id')
print(f"Process group idle: {is_idle}")

Solutions:

1. Check NiFi queue backpressure:

   • Open NiFi UI
   • Check if queues are full (backpressure applied)
   • Increase queue size or adjust batch size

2. Add timeouts:

   config = NiFiConfig(
       connection_timeout=60,  # Request timeout
       base_delay=2,
       max_attempts=4
   )
   

3. Monitor NiFi logs:

   tail -f /opt/nifi/logs/nifi-app.log
   

Problem 5: High Memory Usage

Symptoms:

• Airflow worker OOM (Out of Memory)
• Slow batch processing
• Python process using excessive memory

Solutions:

1. Reduce batch size:

   # Decrease partition_hours to create smaller batches
   source_config = SourceConfig.create_source_config({
       'partition_hours': 4,  # Smaller batches = less memory
       ...
   })
   

2. Process batches sequentially:

   # Instead of parallel processing
   for batch in batches:
       processor = NiFiProcessor(...)
       result = processor.process()
       # Process one at a time
   

3. Optimize SQL queries:

   # Use custom_sql with column selection
   source_config = SourceConfig.create_source_config({
       'custom_sql': 'SELECT id, name, date FROM sales.orders WHERE ...',
       ...
   })
   

Problem 6: Permission Denied on State Tables

Symptoms:

psycopg2.errors.InsufficientPrivilege: permission denied for table pipeline_data

Solutions:

-- Grant permissions to Airflow user
GRANT USAGE ON SCHEMA state TO airflow_user;
GRANT ALL PRIVILEGES ON ALL TABLES IN SCHEMA state TO airflow_user;
GRANT ALL PRIVILEGES ON ALL SEQUENCES IN SCHEMA state TO airflow_user;

-- Verify permissions
SELECT 
    grantee, 
    table_schema, 
    table_name, 
    privilege_type
FROM information_schema.role_table_grants
WHERE table_schema = 'state'
  AND grantee = 'airflow_user';

Problem 7: Duplicate Data in Target

Symptoms:

• Same data appears multiple times in data lake
• Watermark progressing but data duplicated

Diagnosis:

-- Check for duplicate batches
SELECT 
    partition_key,
    COUNT(*) as batch_count
FROM state.entity_state
WHERE entity_name = 'your_table'
  AND bronze_state_status = 'success'
GROUP BY partition_key
HAVING COUNT(*) > 1;

Solutions:

1. Check idempotency in NiFi:

   • Ensure NiFi flow handles duplicates (UPSERT vs INSERT)
   • Add deduplication logic in NiFi

2. Verify partition_key uniqueness:

   # Ensure unique partition keys
   partition_key = f"{start_time.strftime('%Y-%m-%d_%H:%M:%S')}"
   

3. Add idempotency check in DAG:

   # Check if batch already processed
   existing_batch = check_batch_exists(partition_key)
   if existing_batch:
       print(f"Batch {partition_key} already processed, skipping")
       continue
   

---

API Reference

Complete Import List

from airflow_nifi_pipeline_utils import (
    # Configuration
    NiFiConfig,
    
    # Models
    SourceConfig,
    SourceConfigResult,
    BatchData,
    BatchDataResult,
    
    # Managers
    NiFiConnectionManager,
    PipelineStateManager,
    DatabaseConnection,
    
    # Processors
    NiFiProcessor,
    
    # Exceptions
    RetryableNiFiError,
    NonRetryableNiFiError,
    
    # Utilities
    get_access_token,
    safe_schedule_process_group,
    wait_for_process_group_ready,
    check_process_group_idle,
    convert_datetimes,
    
    # SQL Queries (advanced use)
    QUERIES
)

Package Structure

airflow_nifi_pipeline_utils/
├── __init__.py              # Main exports
├── config/
│   └── nifi_config.py       # NiFiConfig class
├── models/
│   ├── batch_data.py        # BatchData, BatchDataResult
│   └── source_config.py     # SourceConfig, SourceConfigResult
├── managers/
│   ├── nifi_connection_manager.py    # NiFiConnectionManager
│   └── pipeline_state_manager.py     # PipelineStateManager, DatabaseConnection
├── processors/
│   └── nifi_processor.py    # NiFiProcessor
├── exceptions/
│   └── nifi_exceptions.py   # RetryableNiFiError, NonRetryableNiFiError
├── utils/
│   ├── nifi_utils.py        # get_access_token, process group utilities
│   └── datetime_utils.py    # convert_datetimes
└── queries/
    └── sql_queries.py       # QUERIES dictionary

---

Migration Guide

From Manual Implementation

If you're currently managing NiFi integration manually:

Before:

# Your existing DAG (200+ lines)
def authenticate_nifi():
    response = requests.post(
        f"{NIFI_URL}/nifi-api/access/token",
        data={'username': USERNAME, 'password': PASSWORD}
    )
    return response.text

def send_to_nifi(data):
    token = authenticate_nifi()
    headers = {'Authorization': f'Bearer {token}'}
    for attempt in range(5):
        try:
            response = requests.post(
                f"http://{NIFI_HOST}:{NIFI_PORT}/contentListener",
                data=json.dumps(data),
                headers=headers,
                timeout=30
            )
            if response.status_code == 200:
                return True
            time.sleep(2 ** attempt)
        except Exception as e:
            if attempt == 4:
                raise
    return False

def update_watermark(table, watermark):
    conn = psycopg2.connect(...)
    cursor = conn.cursor()
    cursor.execute(f"UPDATE watermarks SET high_watermark = '{watermark}' WHERE table_name = '{table}'")
    conn.commit()
    conn.close()

# ... 150 more lines ...

After:

# Using this library (20 lines)
from airflow_nifi_pipeline_utils import *

def run_pipeline(**context):
    config = NiFiConfig(url=NIFI_URL, username=USER, password=PASS)
    state_mgr = PipelineStateManager(postgres_conn_id='postgres_conn', state_schema='state')
    
    pipeline_name = state_mgr.register_pipeline('orders')
    pipeline_id, run_id = state_mgr.start_pipeline_run(pipeline_name, 'incremental')
    
    last_wm = state_mgr.get_latest_high_watermark(pipeline_name)
    
    source_config = SourceConfig.create_source_config({
        'table_name': 'orders',
        'schema_name': 'sales',
        'pipeline_name': pipeline_name,
        'run_id': run_id,
        'load_type': 'incremental',
        'bookmark_column': 'updated_at',
        'last_watermark': last_wm,
        'partition_hours': 8
    })
    
    batches = generate_batches(source_config)  # Your batch generation
    
    with NiFiConnectionManager(config) as nifi_mgr:
        for batch in batches:
            processor = NiFiProcessor(batch, config, 'postgres_conn', 'state', nifi_mgr)
            result = processor.process()
    
    state_mgr.complete_pipeline_run(pipeline_id, run_id, 'bronze', 'success')

Migration Steps:

1. Install the library:

   pip install airflow-nifi-pipeline-utils
   

2. Set up state tables:

   psql -h your-host -d your-db -f pipeline_info_ddl.sql
   

3. Configure Airflow connection:

   • Add PostgreSQL connection in Airflow UI

4. Update DAG imports:

   from airflow_nifi_pipeline_utils import (
       NiFiConfig,
       SourceConfig,
       NiFiConnectionManager,
       NiFiProcessor,
       PipelineStateManager
   )
   

5. Replace manual code:

   • Authentication → NiFiConnectionManager
   • Retry logic → NiFiProcessor
   • State tracking → PipelineStateManager
   • SQL generation → SourceConfig

6. Test in dev environment

7. Deploy to production

---

Dependencies

apache-airflow>=2.5.0
requests>=2.20.0,<3.0.0
nipyapi==0.22.0
psycopg2-binary==2.9.10
urllib3>=1.26.0

---

License

This project is licensed under the Apache License 2.0 - see the LICENSE file for details.

---

Contributing

Contributions are welcome! Please:

1. Fork the repository
2. Create a feature branch (git checkout -b feature/amazing-feature)
3. Add tests for new functionality
4. Commit your changes (git commit -m 'Add amazing feature')
5. Push to the branch (git push origin feature/amazing-feature)
6. Open a Pull Request

See CONTRIBUTING.md for detailed guidelines.

---

Support

• GitHub Issues: Report bugs or request features
• Documentation: Full documentation
• Email: sandil.tandukar@dlytica.com

---

Acknowledgments

Author: Sandil Tandukar
Maintainer: Sandil Tandukar
Organization: Dlytica

Built with love and respect for the data engineering community.

Special thanks to all contributors and users who have provided feedback and improvements.

---

Related Projects

• Apache Airflow - Workflow orchestration platform
• Apache NiFi - Data flow automation
• nipyapi - NiFi Python API client

---

Project Stats

---

Version: 1.1.3
Last Updated: January 2025
Status: Production Ready