DEV Community

Cover image for Teaching an AI agent to forget a deleted Redshift row
Akshat Gupta
Akshat Gupta

Posted on Edited on

Teaching an AI agent to forget a deleted Redshift row

Ticket 2 is gone.

Someone deleted it from the warehouse on Tuesday. The customer is gone with it. Nothing in Redshift knows they exist.

But the AI agent still does. Ask about that customer a week later and it answers in a calm, confident voice. It is describing a row that no longer exists.

Nothing crashed. No error. The agent just kept a ghost.

I wanted to chase that ghost out. So for Mergetober I built a Redshift connector for Cognee. The code lives in PR #403 for issue #4773.

I should be honest. I had never touched Redshift before. I picked it because I wanted to see how a warehouse works from the inside. Building a connector was a good excuse to poke around.

So what does it do?

Cognee turns documents into a knowledge graph that an agent can search. My connector feeds that graph from Redshift. You pick tables or SQL queries. Then:

  • Every row becomes a short document.
  • Every table gets a summary page, built from its column comments.
  • Later runs read only what changed.
  • Deleted rows are forgotten.

My favourite part is the summary page. Picture a column called priority. By itself it tells a model nothing. Now add the column comment "p1 = outage, p4 = question". Suddenly the model knows what a p1 ticket is.

source = redshift_source(
    [RedshiftTable("public.support_tickets", primary_key="ticket_id", cursor_column="updated_at")],
    connection=DataApiConnection(database="dev", workgroup_name="cognee-test"),
)
await cognee.remember(source, dataset_name="redshift",
                      primary_key="id", write_disposition="merge")
Enter fullscreen mode Exit fullscreen mode

The front door, not the secret passage

Redshift usually hides inside a private network. The classic way in is an SSH tunnel through a bastion host. It works. It is also a chore, like climbing through a window when the front door is right there.

The front door is the Data API. It is a plain HTTPS endpoint. You send a statement. You wait. You read the result pages. AWS runs the query inside your network for you.

No tunnel. No driver. No open port. It works with IAM or a Secrets Manager secret. The direct driver connection is still there as a backup.

How do you notice a row that is not there?

This was the fun puzzle. Redshift has no change feed. Nobody taps you on the shoulder and says "row 2 left the building."

So every run also asks for one thing: the primary keys. One column. A columnar store reads that quickly. Then I compare two lists.

Two lists of keys. One of them is shorter.

If a key was there last time and is missing now, the row is gone. The connector removes it from the graph.

I added two safety rules, because a wrong delete is worse than a missed one:

  • A failed query never deletes anything. The run stops. The next one starts from the same spot.
  • A row that shows up late with an old timestamp is still found. The key list does not care about timestamps.

The real test

I started with fakes. A SQLite stand-in. A stubbed AWS API. A real Postgres. 70 tests passed. It felt good for about five minutes. A fake only knows what you taught it.

So I created a real Redshift Serverless workgroup in eu-north-1 and filled a small table of support tickets.

The real run

The first sync took in 6 rows. I asked what the table holds. The answer even explained what p1 means. That came straight from the column comment in Redshift. I did not expect that to feel so satisfying.

Then I changed ticket 3 and deleted ticket 2, right inside Redshift. The second sync reported rows: 3, deleted: 3. The deleted ticket's customer vanished from both the documents and the graph. The other two stayed.

The ghost was gone.

The bug only a real cluster could find

My first live run died with Redshift endpoint doesn't exist in this region.

My docs told people to set AWS_REGION. But boto3 does not read that variable. It reads AWS_DEFAULT_REGION, or your profile. My profile said us-east-1. My cluster lived in eu-north-1.

No fake would ever have shown me this. The connector now honours AWS_REGION, and a test guards it.

What I did not test

The direct connection through an SSH or SSM tunnel only has unit tests. The Data API is the default, so that is the path I ran live. I would rather tell you than let you find out.

Try it yourself

The example and a seed file live in packages/connector/redshift/. New accounts get a free trial of Redshift Serverless, so a short test costs very little.

One warning from me to you: delete the workgroup when you finish. It keeps billing until you do.

Point it at your own warehouse and tell me what breaks.

Top comments (0)