JSONL in ETL Pipelines: Kafka, Airflow & Data Warehouses

from airflow import DAG
from airflow.operators.python import PythonOperator
from airflow.providers.amazon.aws.hooks.s3 import S3Hook
from datetime import datetime, timedelta
import json
import requests

default_args = {
    'owner': 'data-team',
    'retries': 2,
    'retry_delay': timedelta(minutes=5),
}

def extract(**context):
    """Extract data from API and write JSONL to local disk."""
    ds = context['ds']  # execution date: YYYY-MM-DD
    response = requests.get(
        'https://api.example.com/events',
        params={'date': ds},
    )
    events = response.json()
    path = f'/tmp/raw_{ds}.jsonl'
    with open(path, 'w') as f:
        for event in events:
            f.write(json.dumps(event) + '\n')
    return path

def transform(**context):
    """Read raw JSONL, filter and enrich, write cleaned JSONL."""
    ds = context['ds']
    input_path = f'/tmp/raw_{ds}.jsonl'
    output_path = f'/tmp/clean_{ds}.jsonl'
    with open(input_path) as fin, open(output_path, 'w') as fout:
        for line in fin:
            record = json.loads(line)
            if record.get('status') != 'valid':
                continue
            record['processed_at'] = datetime.utcnow().isoformat()
            fout.write(json.dumps(record) + '\n')
    return output_path

def load(**context):
    """Upload cleaned JSONL to S3."""
    ds = context['ds']
    local_path = f'/tmp/clean_{ds}.jsonl'
    s3 = S3Hook(aws_conn_id='aws_default')
    s3.load_file(
        filename=local_path,
        key=f'etl/events/{ds}/data.jsonl',
        bucket_name='data-lake',
        replace=True,
    )

with DAG(
    'jsonl_etl_pipeline',
    default_args=default_args,
    schedule_interval='@daily',
    start_date=datetime(2026, 1, 1),
    catchup=False,
) as dag:
    t_extract = PythonOperator(task_id='extract', python_callable=extract)
    t_transform = PythonOperator(task_id='transform', python_callable=transform)
    t_load = PythonOperator(task_id='load', python_callable=load)

    t_extract >> t_transform >> t_load

import json
from typing import Callable

def stream_transform(
    input_path: str,
    output_path: str,
    transform_fn: Callable[[dict], dict | None],
    batch_size: int = 5000,
) -> dict:
    """Stream-transform a JSONL file with a user-defined function.
    Return None from transform_fn to skip a record."""
    stats = {'input': 0, 'output': 0, 'skipped': 0}
    batch = []

    with open(input_path) as fin, open(output_path, 'w') as fout:
        for line in fin:
            line = line.strip()
            if not line:
                continue
            stats['input'] += 1
            record = json.loads(line)
            result = transform_fn(record)
            if result is None:
                stats['skipped'] += 1
                continue
            batch.append(json.dumps(result))
            if len(batch) >= batch_size:
                fout.write('\n'.join(batch) + '\n')
                stats['output'] += len(batch)
                batch.clear()

        if batch:
            fout.write('\n'.join(batch) + '\n')
            stats['output'] += len(batch)

    return stats

# Usage inside an Airflow task
def enrich(record):
    if record.get('amount', 0) < 0:
        return None  # skip invalid
    record['currency'] = record.get('currency', 'USD').upper()
    return record

stats = stream_transform(
    '/tmp/raw.jsonl',
    '/tmp/enriched.jsonl',
    transform_fn=enrich,
)
print(f'Processed: {stats}')

import json
import os

def get_checkpoint(checkpoint_path: str) -> int:
    """Read the last committed line number."""
    if os.path.exists(checkpoint_path):
        with open(checkpoint_path) as f:
            return int(f.read().strip())
    return 0

def save_checkpoint(checkpoint_path: str, line_num: int):
    """Atomically save the current checkpoint."""
    tmp = checkpoint_path + '.tmp'
    with open(tmp, 'w') as f:
        f.write(str(line_num))
    os.replace(tmp, checkpoint_path)  # atomic on POSIX

def process_with_checkpoint(
    input_path: str,
    output_path: str,
    checkpoint_path: str,
    batch_size: int = 1000,
):
    """Process a JSONL file with checkpoint-based resumption."""
    start_line = get_checkpoint(checkpoint_path)
    processed = 0

    with open(input_path) as fin, \
         open(output_path, 'a') as fout:
        for line_num, line in enumerate(fin, 1):
            if line_num <= start_line:
                continue  # skip already-processed lines
            line = line.strip()
            if not line:
                continue
            record = json.loads(line)
            # --- your transform logic ---
            record['etl_version'] = 'v2'
            fout.write(json.dumps(record) + '\n')
            processed += 1

            if processed % batch_size == 0:
                fout.flush()
                save_checkpoint(checkpoint_path, line_num)
                print(f'Checkpoint at line {line_num}')

        # Final checkpoint
        save_checkpoint(checkpoint_path, line_num)

    print(f'Done. Processed {processed} records from line {start_line + 1}')

# Usage
process_with_checkpoint(
    'input.jsonl',
    'output.jsonl',
    'checkpoint.txt',
    batch_size=5000,
)