DEV Community

Mukesh
Mukesh

Posted on

Building an Airbyte Destination Connector That Embeds and Dedupes Records Before They Hit pgvector

The Gap Between Airbyte and RAG Pipelines

Airbyte moves data from a source into a destination. A RAG ingestion pipeline needs one more step in the middle: turn each record into a chunk, embed it, and write the vector alongside the source row so a retrieval query can find it later. Most teams solve this by letting Airbyte land raw rows in Postgres, then running a separate cron job or Lambda that reads the table, computes embeddings, and upserts into a vector column. That works, but it means two systems own the same data's lifecycle — if the downstream embedding job falls behind, or dies mid-run, nothing in Airbyte's sync history tells you your vector index is stale.

The alternative is to embed inside the destination connector itself, so record landed and record embedded are the same event. Airbyte's Python CDK makes this straightforward to build, but two details will bite you if you skip them: state-message ordering under at-least-once delivery, and how you dedupe embedding calls so a replayed batch doesn't burn API budget re-embedding rows you already have.

Connector Shape

A destination connector implements three things: spec(), check(), and write(). The interesting logic lives in write(), which receives a stream of AirbyteMessage objects — a mix of RECORD and STATE messages — and must emit STATE back out only after everything before it is durably persisted.

# destination_pgvector/destination.py
from airbyte_cdk.destinations import Destination
from airbyte_cdk.models import AirbyteMessage, Type, AirbyteStateMessage
import hashlib

class DestinationPgvector(Destination):
    def write(self, config, configured_catalog, input_messages):
        buffer = []
        for message in input_messages:
            if message.type == Type.RECORD:
                buffer.append(message.record)
                if len(buffer) >= config["batch_size"]:
                    self._flush(buffer, config)
                    buffer = []
            elif message.type == Type.STATE:
                # flush whatever's buffered BEFORE forwarding state
                if buffer:
                    self._flush(buffer, config)
                    buffer = []
                yield message
        if buffer:
            self._flush(buffer, config)
Enter fullscreen mode Exit fullscreen mode

That ordering — flush, then yield STATE — is the whole contract. If you yield STATE before the buffer is actually written, a crash between the two leaves Airbyte believing data landed that never did. If you never yield STATE until end of stream, a failure partway through forces a full re-sync from scratch instead of resuming near the last checkpoint. Buffering by batch_size and flushing on every STATE message (which Airbyte's sources emit periodically, not just at the end) gives you both correctness and reasonable checkpoint granularity.

Making Embedding Calls Idempotent

Airbyte's delivery guarantee is at-least-once: after a connector restart, some already-flushed records can be redelivered. If _flush blindly calls the embedding API and inserts, a restart mid-sync means paying for and storing duplicate embeddings for the same source rows. The fix is to key every row on a content hash and make the destination write an upsert, not an insert:

def _flush(self, buffer, config):
    rows = []
    for record in buffer:
        content = self._extract_text(record.data, config["text_field"])
        content_hash = hashlib.sha256(content.encode()).hexdigest()
        rows.append((record.data[config["id_field"]], content, content_hash))

    # skip rows whose hash we've already embedded
    existing = self._fetch_existing_hashes([r[0] for r in rows])
    to_embed = [r for r in rows if existing.get(r[0]) != r[2]]
    if not to_embed:
        return

    vectors = self._embed_batch([r[1] for r in to_embed], config)
    self._upsert(to_embed, vectors)
Enter fullscreen mode Exit fullscreen mode

_fetch_existing_hashes is one indexed SELECT id, content_hash FROM docs WHERE id = ANY(%s) query. Comparing the incoming hash against the stored one does double duty: it makes redelivery a no-op (same id, same hash → skip), and it makes re-syncs after a source-side edit correct (same id, different hash → re-embed and overwrite), which a plain "insert if not exists" dedupe strategy would miss entirely.

Batching Around Embedding Provider Limits

Embedding APIs rate-limit on both requests-per-minute and tokens-per-minute, and a naive one-record-per-call loop will hit both under a normal Airbyte sync's throughput. Batch the actual API call separately from your write batch size — you can flush to Postgres every 500 records while still capping each embedding request at whatever the provider's max batch is (100 inputs for OpenAI's embeddings endpoint, for example):

def _embed_batch(self, texts, config):
    vectors = []
    for i in range(0, len(texts), config["embed_batch_size"]):
        chunk = texts[i : i + config["embed_batch_size"]]
        vectors.extend(self._call_embedding_api_with_backoff(chunk, config))
    return vectors

def _call_embedding_api_with_backoff(self, chunk, config, attempt=0):
    try:
        return self.embed_client.embed(chunk)
    except RateLimitError:
        if attempt >= 5:
            raise
        time.sleep(min(2 ** attempt, 30))
        return self._call_embedding_api_with_backoff(chunk, config, attempt + 1)
Enter fullscreen mode Exit fullscreen mode

Without this split, teams either set batch_size low enough to stay under the embedding limit — which makes Postgres writes needlessly chatty — or hit 429s constantly because they sized batching only around database throughput.

Testing With the CDK's Acceptance Suite

Airbyte ships a standard destination test harness (destination-acceptance-test) that spins up your connector in Docker, feeds it a fixture catalog, and asserts records land correctly, including after a simulated mid-sync failure. Point it at a docker-compose Postgres with the pgvector extension enabled:

# acceptance-test-config.yml
connector_image: airbyte/destination-pgvector:dev
tests:
  spec:
    - spec_path: "integration_tests/spec.json"
  basic_read: []
  incremental:
    - config_path: "secrets/config.json"
      configured_catalog_path: "integration_tests/configured_catalog.json"
      future_state_path: "integration_tests/abnormal_state.json"
Enter fullscreen mode Exit fullscreen mode

The incremental block is what actually exercises the redelivery path — it replays records against a state file that's intentionally stale, which is exactly the scenario the content-hash dedupe above is built to survive. Running this before wiring the connector into a real sync catches state-ordering bugs that a manual test with a small CSV source almost never surfaces, because manual tests rarely include a forced restart mid-batch.

Where This Ends Up

Once the connector's checkpointing and dedupe are correct, the rest of the RAG pipeline gets simpler, not more complex: any of Airbyte's ~600 existing source connectors can feed it without custom glue code, incremental syncs only re-embed rows that actually changed, and a sync's Airbyte-native retry and alerting cover the embedding step too, instead of that step living in a separate cron job nobody's monitoring. The cost is building one extra connector — but it's a connector, with the CDK's test tooling and deployment model behind it, not a bespoke pipeline you're maintaining solo.

Top comments (0)