Rolling Your Own Operators and Hooks

When you've copy-pasted the same @task one time too many, it's time to give your integration a name — a custom operator, backed by a hook that knows how to talk to your system.

You hit the same wall on every project. The first DAG that pulls from the bookshop’s orders API is a single @task — a few lines of requests, a token from a connection, a bit of pagination. Clean enough. Then a second DAG needs the same pull. Then a third wants just the returns endpoint. Each time you copy the function, tweak two lines, and promise yourself you’ll clean it up later. You won’t, and now there are four slightly-different copies of your auth logic and no single place to fix the one that’s wrong.

That’s the moment to stop writing tasks and start writing a component. Airflow gives you two: the operator (a reusable unit of work you drop into any DAG) and the hook (a reusable client that knows how to reach an external system). Used well, they turn “the thing that talks to the bookshop API” from a shape you keep redrawing into a name you import.

A custom operator is a class with an execute method

An operator is a subclass of BaseOperator. You take configuration in __init__ and you put the actual work in execute. Everything else — retries, scheduling, where it sits in the graph — you inherit for free.

from airflow.sdk import BaseOperator


class BookshopApiToWarehouseOperator(BaseOperator):
    template_fields = ("endpoint", "since", "table")
    template_ext = (".sql",)
    ui_color = "#0d9488"
    ui_fgcolor = "#ffffff"

    def __init__(self, *, conn_id, endpoint, since, table, **kwargs):
        super().__init__(**kwargs)
        self.conn_id = conn_id
        self.endpoint = endpoint
        self.since = since
        self.table = table

    def execute(self, context):
        hook = BookshopApiHook(conn_id=self.conn_id)
        rows = hook.fetch(self.endpoint, since=self.since)
        hook.load_to_warehouse(rows, table=self.table)
        return len(rows)

In Airflow 3.x the base classes live in the Task SDK, so the import is from airflow.sdk import BaseOperator — not the old airflow.models path (which may still resolve, but shouldn’t be what you reach for). Whatever execute returns is pushed to XCom, so returning a row count gives downstream tasks something to branch on.

The **kwargs forwarded to super().__init__ is what lets a caller pass retries, pool, task_id — all the standard operator arguments — through your constructor without you re-declaring each one. One piece of 2.x ceremony you can now delete: the old @apply_defaults decorator. In Airflow 3 default-argument application is wired into the metaclass, so it happens automatically — if you’re porting an operator forward, strip the decorator and the import.

The templating surface is bigger than template_fields

template_fields is the mechanism most people know: name a field there and it gets Jinja-rendered before execute runs, so a caller can write endpoint="{{ params.endpoint }}" or hand you a templated date and trust it resolves. Fields you leave off the list are used verbatim. But there are three neighbours worth knowing, because together they turn a crude operator into one that behaves like a first-class citizen in the UI.

template_ext lets a templated field load its value from a file. In the operator above I listed (".sql",), which means if a caller passes table="load_orders.sql", Airflow notices the extension, reads that file from the DAG folder, and then renders it as a template. This is how the built-in SQL operators let you keep queries in .sql files next to the DAG instead of as inline strings. The field has to appear in template_fields too — template_ext only decides which values get the file-loading treatment.

template_fields_renderers is pure ergonomics, and free. It’s a dict mapping a field name to a renderer the UI understands, and it controls syntax highlighting on the Rendered Template view — the tab where you inspect exactly what a field resolved to for a given run. Point a SQL field at "sql" and a JSON field at "json":

    template_fields_renderers = {"table": "sql", "endpoint": "py"}

Now when you’re debugging why last night’s run hit the wrong endpoint, the resolved value is highlighted instead of dumped as a flat string. ui_color and ui_fgcolor (set above to the Data-track teal) tint the task box in the Grid and Graph views, which sounds cosmetic until you have a DAG with forty tasks and want your custom loads to stand out from the dbt runs at a glance.

Finally, params. Templating pulls from the render context, and params is how you give an operator (or a whole DAG) typed, defaulted, UI-editable inputs that land in that context as {{ params.x }}. You attach a Param with a JSON-schema type and a default:

from airflow.sdk import DAG
from airflow.sdk.definitions.param import Param

with DAG(
    "bookshop_backfill",
    params={"region": Param("eu", type="string", enum=["eu", "us", "apac"])},
) as dag:
    BookshopApiToWarehouseOperator(
        task_id="load",
        conn_id="bookshop_api_default",
        endpoint="orders/{{ params.region }}",
        since="{{ ds }}",
        table="load_orders.sql",
    )

Because region is declared with an enum, a Trigger DAG w/ config run renders it as a validated dropdown instead of a free-text box, and a bad value is rejected before the DAG runs rather than at load time against your API. For an operator, params are the seam between “the value is fixed in code” and “someone can override it at trigger time without editing the DAG” — which is exactly what you want for a manual backfill of one region.

A hook is where the integration actually lives

Notice what the operator above doesn’t do: it never opens a socket, parses a connection, or knows the shape of the API. It calls a hook. That separation is the whole point. The operator is the Airflow-facing shell — it deals with context, XCom, and templating. The hook is the system-facing client — it deals with connections, auth, and requests. Keep the operator thin and the hook does the work.

from airflow.sdk import BaseHook
import requests


class BookshopApiHook(BaseHook):
    conn_name_attr = "conn_id"
    default_conn_name = "bookshop_api_default"
    conn_type = "bookshop"
    hook_name = "Bookshop API"

    def __init__(self, conn_id=default_conn_name):
        super().__init__()
        self.conn_id = conn_id
        self._session = None

    def get_conn(self):
        if self._session is None:
            conn = self.get_connection(self.conn_id)
            extra = conn.extra_dejson
            self._session = requests.Session()
            self._session.headers["Authorization"] = f"Bearer {conn.password}"
            self._session.headers["X-Api-Version"] = extra.get("api_version", "v2")
            self._session.request_timeout = extra.get("timeout", 30)
            self._base_url = conn.host
        return self._session

    def fetch(self, endpoint, since):
        session = self.get_conn()
        resp = session.get(f"{self._base_url}/{endpoint}", params={"since": since})
        resp.raise_for_status()
        return resp.json()["data"]

BaseHook gives you get_connection, which reads the credential you stored once in the Airflow UI or an environment variable and hands back a Connection object. Your code never sees a raw secret and never hardcodes an endpoint. Caching the client in get_conn matters more than it looks — a mapped or retried task can call the hook repeatedly, and you’d rather not rebuild a session each time.

The connection model, and where the real config hides

A Connection is more than a username and password. It has structured fields — host, schema, login, password, port, conn_type — and then the field that saves you: extra. The named fields cover the common shape of a database or HTTP endpoint, but every integration has settings that don’t fit those slots: an API version, a region, a retry budget, a TLS flag. Those go in extra as a JSON blob, and you read them through extra_dejson, which parses that blob into a dict so you’re not calling json.loads yourself. In the hook above, api_version and timeout live there. This is where most real per-environment config ends up, and it’s the first place to look when someone asks “how do I make this hook talk to staging instead of prod” — usually the answer is a different extra, not different code.

It helps to remember that a connection is just a URI under the hood. The same bookshop connection can be defined three ways, and they’re equivalent:

  • In the UI — Admin → Connections, one field per attribute, extra as a JSON textarea. Good for humans, bad for reproducibility.
  • As an environment variable — AIRFLOW_CONN_BOOKSHOP_API_DEFAULT='bookshop://:[email protected]/?api_version=v2'. The AIRFLOW_CONN_ prefix plus the upper-cased conn id, value is a connection URI. This is the form to bake into an Astro project’s .env for local dev, because it’s checked in with the code (minus the secret) and needs no manual UI step.
  • Via a secrets backend — the same URI, fetched from Vault or AWS Secrets Manager at runtime. That’s a whole chapter later in this series.

Because it’s a URI, conn.get_uri() round-trips a connection back to that string — handy when you want to log which endpoint you resolved to (with the password redacted, which Airflow does for you) or hand a URI to a library that wants one. And the parsed pieces are all addressable: conn.host, conn.schema (the path segment — a database name for a DB connection), conn.login, conn.port. Knowing that schema is the path and extra is the query string demystifies why a connection URI looks the way it does.

Making the hook a good UI citizen

The four class attributes at the top of BookshopApiHook — conn_name_attr, default_conn_name, conn_type, hook_name — aren’t decoration. They’re what let Airflow discover your hook and render a sensible connection form for it. conn_type is the machine key (bookshop), hook_name is the human label that shows in the connection-type dropdown, conn_name_attr tells Airflow which constructor argument names the connection, and default_conn_name is the fallback id.

Once a hook is discoverable, two optional classmethods let you shape its connection form:

    @classmethod
    def get_ui_field_behaviour(cls):
        return {
            "hidden_fields": ["schema", "port"],
            "relabeling": {"host": "API base URL", "password": "Bearer token"},
        }

get_ui_field_behaviour hides the standard fields your integration doesn’t use and relabels the ones it repurposes, so nobody types a port number into a REST connection. It’s the difference between a connection form that reads as “your integration” and one that reads as “generic HTTP, good luck.”

There’s a sibling method, get_connection_form_widgets, which promotes a setting you’d otherwise bury in extra — say api_version — into a labelled box of its own. It still exists on BaseHook in 3.x, but be careful with the examples you’ll find for it: nearly all of them open with

from flask_appbuilder.fieldwidgets import BS3TextFieldWidget
from wtforms import StringField

and that is 2.x code. Airflow 3 replaced the Flask/FAB webserver with the React UI, and flask_appbuilder is no longer installed with core Airflow — that import raises ModuleNotFoundError on a stock 3.x environment unless you’ve pulled in the FAB provider. If you’re only adding your own hook to your own deployment, put the extra settings in extra and read them through extra_dejson (below); it costs a JSON key and buys you code that survives the UI rewrite.

The link between a conn_type string and your hook class is hook discovery. A Connection object carries a get_hook() method, and calling it hands back an instance of the right hook for that connection’s type — but only if Airflow knows the mapping. For a hook living in include/, that mapping comes from the class attributes above; for a packaged provider, it comes from the connection-types entry in the provider’s get_provider_info metadata, which Airflow reads at startup to build the whole registry of connection types, widgets, and field behaviours. This is why a freshly-installed provider makes its connection type appear in the UI dropdown with no further wiring: the metadata is the registration. When you write conn.get_hook() in a generic operator — one that shouldn’t hardcode which system it talks to — you’re leaning on exactly this registry to resolve the concrete hook at runtime from the connection the user picked.

The operator lifecycle is more than execute

execute is the method you’ll write ninety percent of the time, but BaseOperator gives you hooks around it, and one of them is not optional in production.

pre_execute(context) and post_execute(context, result) run immediately before and after execute, on the same worker, inside the same try/except. They’re the place for setup and teardown that you want logged as part of the task but kept out of the core logic — validating that a required connection exists before you do real work, or emitting a metric with the row count after. They’re genuinely optional; reach for them only when the bracketing logic would clutter execute.

on_kill is the one that matters. When a task is killed — a timeout fires, someone clears it from the UI, the scheduler decides it’s a zombie — Airflow calls on_kill so you can clean up. If your operator kicked off work in an external system (submitted a Spark job, opened a long transaction, started a warehouse query), the default kill just stops the Airflow process and leaves the external job running, burning money and holding locks. on_kill is where you cancel it:

    def execute(self, context):
        hook = BookshopApiHook(conn_id=self.conn_id)
        self._job_id = hook.start_export(self.endpoint, since=self.since)
        return hook.wait_for_export(self._job_id)

    def on_kill(self):
        if getattr(self, "_job_id", None):
            BookshopApiHook(conn_id=self.conn_id).cancel_export(self._job_id)
            self.log.info("Cancelled bookshop export %s", self._job_id)

The rule of thumb: if execute starts anything that outlives the Python function, you owe it an on_kill. Note the two ways you land there: an execution_timeout on the task (the modern replacement for the removed SLA machinery — set a timeout and Airflow kills the task when it blows past it) and a manual clear or a zombie detection when the worker dies. Both funnel through on_kill, so it’s the single place to make your operator a good tenant of shared infrastructure. Guard it defensively — on_kill can fire before execute has assigned self._job_id, which is why the example reads it through getattr with a default rather than assuming the attribute exists. The fourth lifecycle method, execute_complete, belongs to deferrable operators — that’s the next section.

A deferrable custom operator, end to end

The sensors chapter makes the case for deferring: a task that’s only waiting shouldn’t pin a worker. That chapter uses the built-in flavours; here’s how you author the deferrable path in your own operator, because it’s a different shape from a plain execute.

The idea: execute does the synchronous part, then hands off a small async trigger to the triggerer and steps aside by calling self.defer. The triggerer runs the trigger’s coroutine cheaply. When the condition fires, Airflow resumes your operator by calling the method you named — conventionally execute_complete.

First the trigger. It subclasses BaseTrigger, which requires exactly two things: serialize, which returns the classpath and the kwargs needed to rebuild it (because the trigger is stored and reconstructed in the triggerer process), and an async run that yields a TriggerEvent when done.

import asyncio
from airflow.triggers.base import BaseTrigger, TriggerEvent


class BookshopExportTrigger(BaseTrigger):
    def __init__(self, conn_id, job_id, poll_interval=30):
        super().__init__()
        self.conn_id = conn_id
        self.job_id = job_id
        self.poll_interval = poll_interval

    def serialize(self):
        return (
            "bookshop.triggers.BookshopExportTrigger",
            {"conn_id": self.conn_id, "job_id": self.job_id,
             "poll_interval": self.poll_interval},
        )

    async def run(self):
        hook = BookshopApiHook(conn_id=self.conn_id)
        while True:
            status = await hook.async_export_status(self.job_id)
            if status == "ready":
                yield TriggerEvent({"status": "ready", "job_id": self.job_id})
                return
            await asyncio.sleep(self.poll_interval)

Two things make this correct. The run loop uses await asyncio.sleep, not time.sleep — the whole point is that the coroutine yields control back to the triggerer’s event loop while it waits, so thousands of these share one process. And serialize returns only JSON-friendly primitives; the triggerer rebuilds the trigger from exactly those kwargs, so nothing stateful can hide on the instance.

Now the operator defers to it:

    def execute(self, context):
        hook = BookshopApiHook(conn_id=self.conn_id)
        job_id = hook.start_export(self.endpoint, since=self.since)
        self.defer(
            trigger=BookshopExportTrigger(conn_id=self.conn_id, job_id=job_id),
            method_name="execute_complete",
        )

    def execute_complete(self, context, event=None):
        if event["status"] != "ready":
            raise RuntimeError(f"Export failed: {event}")
        hook = BookshopApiHook(conn_id=self.conn_id)
        rows = hook.fetch_export(event["job_id"])
        hook.load_to_warehouse(rows, table=self.table)
        return len(rows)

execute runs on a worker just long enough to kick off the export and get a job id, then self.defer raises a special exception that suspends the task and registers the trigger. The worker slot is released. When the trigger yields its TriggerEvent, the scheduler schedules execute_complete on a worker, passing the event payload — and that’s where you do the fast follow-up work. The event dict is whatever you put in the TriggerEvent, so keep it small and serializable. Note the payoff: the operator holds a worker for two short bursts instead of one long wait, and the waiting itself costs a coroutine.

Two footguns to respect. The hook method the trigger awaits — async_export_status above — has to be genuinely async: a blocking requests call inside a coroutine stalls the entire triggerer event loop and starves every other trigger sharing it, so use an async HTTP client (or run the blocking call in a thread executor). And self.defer accepts a timeout for the wait itself; without one, a trigger whose condition never fires waits forever, holding nothing but never finishing. Set it, and pair it with the same execute_complete handling a failure event so a stuck export surfaces as a failed task instead of a silent hang.

Testing a custom operator

The reason to build this at all is testability, so prove it. A hook is a plain object — instantiate it, mock get_connection, assert fetch builds the right URL, no Airflow runtime involved. The operator has one Airflow-specific thing worth testing directly: that its template_fields actually render. You can drive that without a scheduler by rendering the operator inside a DAG context:

import datetime
from airflow.sdk import DAG


def test_endpoint_templating():
    with DAG("t", start_date=datetime.datetime(2026, 1, 1)) as dag:
        op = BookshopApiToWarehouseOperator(
            task_id="load",
            conn_id="bookshop_api_default",
            endpoint="orders/{{ params.region }}",
            since="{{ ds }}",
            table="load_orders.sql",
        )
    op.render_template_fields({"params": {"region": "eu"}, "ds": "2026-04-25"})

    assert op.endpoint == "orders/eu"
    assert op.since == "2026-04-25"

Note that render_template_fields returns None — it rewrites the operator’s attributes in place. So you call it for the side effect and then assert against op itself; there is no rendered copy handed back, and trying to capture one is a common way to end up asserting against None.

Rendering the template fields against a fake context and asserting the resolved values is the cheapest real test you can write for an operator, and it catches the most common bug: a field you forgot to list in template_fields, which silently ships the literal {{ params.region }} string to your API. Because table ends in .sql and is in template_ext, the same machinery would load and render that file — a good second assertion once the query file exists.

Where this code lives, and when to package it

An Astro (Astronomer) project is a directory with dags/, include/, plugins/, and a Dockerfile. Your operator and hook are just importable Python, so you have three homes, in rough order of ceremony:

  • A module under include/ — put bookshop.py there (or a package you add to the image) and from bookshop import BookshopApiToWarehouseOperator in your DAGs. Simplest path, prefer it for a single project.

  • The plugins/ folder — Airflow adds it to the path automatically, and it’s where the plugin system looks to register operator extra links. If you want a “Open in Bookshop console” button on the task instance page, you subclass BaseOperatorLink, give it a name and a get_link that returns a URL, and expose it via the operator’s operator_extra_links:

    from airflow.sdk import BaseOperatorLink

class BookshopConsoleLink(BaseOperatorLink):
    name = "Bookshop console"

    def get_link(self, operator, *, ti_key):
        return f"https://console.bookshop.dev/exports?endpoint={operator.endpoint}"

class BookshopApiToWarehouseOperator(BaseOperator):
    operator_extra_links = (BookshopConsoleLink(),)

    get_link receives the operator instance and a ti_key identifying the exact run, so the button can deep-link to that task’s job in the external system. It’s a small thing that saves an on-call engineer three clicks at 2am, and that kind of glue belongs in a plugin.

  • A real provider package — when the integration is shared across repositories, not just DAGs, you package it as an installable provider (a pip-installable distribution with a get_provider_info entry point that registers your hook’s conn_type, connection widgets, and extra links). Now pip install airflow-provider-bookshop gives any Airflow instance your connection type in the UI dropdown and your operators on the path. It’s more work and only pays off at that scope — but it’s the same mechanism the official providers use, and there’s nothing magic about it.

Either of the first two: astro dev start bakes the code into the local image and the UI at localhost:8080 runs your DAGs against it. Restart the project after adding a new module so the image picks it up.

The other two bookshop hooks

BookshopApiHook is the one worth building from scratch, because it shows the whole pattern. But the rest of this series — and the test suite in chapter 09 — leans on two siblings, so here they are in full. They live beside it in include/bookshop.py, and there is nothing surprising in either; that’s rather the point.

# include/bookshop.py  (alongside BookshopApiHook)
from airflow.sdk import BaseHook
from airflow.providers.amazon.aws.hooks.s3 import S3Hook
from airflow.providers.postgres.hooks.postgres import PostgresHook


class BookshopStorageHook(BaseHook):
    """Finds the files the bookshop's partners drop for us."""

    conn_type = "bookshop_storage"
    hook_name = "Bookshop Storage"

    def __init__(self, conn_id="bookshop_storage_default"):
        super().__init__()
        self.conn_id = conn_id

    def list_keys(self, prefix: str) -> list[str]:
        conn = self.get_connection(self.conn_id)
        bucket = conn.extra_dejson["bucket"]
        return S3Hook(aws_conn_id=self.conn_id).list_keys(bucket, prefix=prefix) or []


class BookshopWarehouseHook(BaseHook):
    """Writes into the warehouse the marts are built from."""

    conn_type = "bookshop_warehouse"
    hook_name = "Bookshop Warehouse"

    def __init__(self, conn_id="bookshop_warehouse_default"):
        super().__init__()
        self.conn_id = conn_id

    def copy_into(self, key: str, *, table: str, batch_rows: int = 5_000) -> int:
        pg = PostgresHook(postgres_conn_id=self.conn_id)
        with pg.get_conn() as conn, conn.cursor() as cur:
            cur.execute(
                f"COPY {table} FROM PROGRAM %s WITH (FORMAT csv, HEADER true)",
                (f"aws s3 cp s3://{key} -",),
            )
            return cur.rowcount

Both are deliberately thin — they wrap a provider hook and expose one verb each, named for what the bookshop does rather than for what the underlying system is called. That’s the whole discipline: list_keys(prefix) and copy_into(key, table=…) are the vocabulary your DAGs speak, and the day you migrate off S3 or Postgres, the DAGs don’t change.

Final thoughts

The instinct to reach for a custom operator early is usually wrong, and the instinct to keep copy-pasting is worse. The honest rule is that operators and hooks aren’t about capability — a @task can do anything they can — they’re about reuse and testability. You extract them the same day you notice you’re maintaining the integration in more than one file. But once you commit, do it fully: list your template_fields, wire an on_kill for anything external, read config from extra_dejson, and give the hook the discovery metadata that makes its connection form legible. Get the split right — operator thin, hook does the work, trigger for the waiting — and the class you write isn’t overhead, it’s the version of the integration you can finally cover with tests.

Next: Sensors and Deferrable Operators: Waiting Without Wasting

Comments