Airflow Custom SDK

[!NOTE] Development Status: This SDK is continuously in active development. Features, APIs, and overall structure are subject to change, and ongoing updates can be observed in the repository.

Apache Airflow is the industry standard for workflow orchestration. However, as organizations scale their MLOps architectures, default operators (like PythonOperator or BashOperator) often result in duplicated code, fragmented logging, and untraceable pipeline executions.

This repository details how we built custom reusable Airflow operators and integrated them with OpenTelemetry to build a highly scalable, observable data orchestration hub. It includes a comprehensive custom SDK and plugin package designed to support both Airflow v2 and Airflow v3, with a primary focus on Databricks, Azure, and Bash integrations.

---

Building a Custom Operator

When multiple data pipelines perform similar tasks—such as launching a Spark job, executing a SQL query, or validating a model output—relying on standard Python operators results in severe code duplication.

Custom operators inherit from Airflow's BaseOperator and override the execute method. Here is an example of an MLModelValidationOperator:

from airflow.models import BaseOperator
from airflow.utils.decorators import apply_defaults

class MLModelValidationOperator(BaseOperator):
    @apply_defaults
    def __init__(self, model_uri: str, test_data_path: str, threshold: float = 0.85, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.model_uri = model_uri
        self.test_data_path = test_data_path
        self.threshold = threshold

    def execute(self, context):
        self.log.info(f"Loading model from {self.model_uri}...")
        self.log.info(f"Running validation tests on {self.test_data_path}...")
        
        # Mock validation execution
        accuracy = 0.89
        
        if accuracy < self.threshold:
            raise ValueError(f"Model validation failed! Accuracy {accuracy} is below threshold {self.threshold}")
            
        self.log.info("Model validation passed successfully!")
        return {"accuracy": accuracy}

By abstracting model loading and scoring logic inside the operator, developers can orchestrate model validation in one line of clean, declarative DAG code.

---

OpenTelemetry Integration

Without centralized telemetry, debugging a failed task in a complex pipeline requires manually loading logs for each task run. Integrating OpenTelemetry into custom operators resolves this by linking task runs with global trace IDs.

Telemetry Architecture

Using a custom Airflow Listener, task executions emit traces to a centralized collector (like Dynatrace or Datadog):

from opentelemetry import trace
from opentelemetry.sdk.trace import TracerProvider

tracer = trace.get_tracer("airflow.orchestration")

def on_task_instance_success(previous_state, task_instance, session):
    with tracer.start_as_current_span(
        name=f"task_{task_instance.task_id}",
        attributes={
            "dag_id": task_instance.dag_id,
            "run_id": task_instance.run_id,
            "operator": task_instance.operator
        }
    ):
        pass

With traces flowing from Airflow, operators, and the downstream API containers, you can inspect a single transaction graph mapping exactly how a model was validated and served!

---

Repository Structure

The project has been restructured to cleanly separate Airflow v2 and v3 codebases:

• custom_sdk/v2/: Contains the original Airflow v2 compatible code (formerly airflow_custom_sdk).
• custom_sdk/ (Root of the module): Contains the main Airflow v3 compatible plugins and SDK logic.
  • custom_sdk/databricks/: Custom Databricks operators and workflow task groups.
  • custom_sdk/azure/: Azure-related extensions and integrations.
  • custom_sdk/bash/: Custom bash plugins (e.g., custom extra links).
• dags/: Sample DAGs demonstrating the usage of the custom SDK (both v2 and v3 workflows).
• docker/: Docker Compose setup and configurations to run Airflow locally with these plugins installed.

Getting Started

Running Locally with Docker

You can spin up an Airflow environment with the custom SDK installed using Docker Compose.

1. Build and start the services:

   docker-compose -f docker/docker-compose.yaml up --build

2. Access Airflow: Navigate to http://localhost:8080 and log in (default credentials: airflow/airflow).

3. Stop the services:

   docker-compose -f docker/docker-compose.yaml down

4. Clean up (remove volumes and orphans):

   docker-compose -f docker/docker-compose.yaml down --volumes --rmi all --remove-orphans

API Testing (cURL Commands)

The plugin may also expose custom API endpoints (e.g., for Databricks plugin tests).

Simple cURL command using Basic Auth:

curl -X GET 'http://localhost:8080/api/v1/dags' --user "airflow:airflow"

Get Airflow Webserver Health:

curl http://localhost:8080/health

Custom API Endpoint Test:

curl -X GET 'http://localhost:8080/databricks_plugin_api/test' --user "airflow:airflow"

Installation (Development)

To install the SDK in editable mode for local development:

pip install -e .

This uses the pyproject.toml configuration and will automatically register the entry points for the Airflow plugins (e.g., DatabricksCustomPlugin and OperatorExtraLinkPlugin).