Supabase RLS with SQL Agent

[garmeeh]

Checked other resources

• I added a very descriptive title to this question.
• I searched the LangChain documentation with the integrated search.
• I used the GitHub search to find a similar question and didn't find it.

Commit to Help

• I commit to help with one of those options 👆

Example Code

Example scrip I am using to test this functionality.
 

import "dotenv/config";

import { ChatOpenAI } from "@langchain/openai";
import { SqlDatabase } from "langchain/sql_db";
import { createSqlAgent, SqlToolkit } from "langchain/agents/toolkits/sql";
import { DataSource } from "typeorm";

export const lookupData = async () => {
  const encodedPassword = encodeURIComponent(process.env.DB_PASSWORD || "");

  const dbUrl = `postgres://postgres.X:${encodedPassword}@X.supabase.com:5432/postgres`;

  const datasource = new DataSource({
    type: "postgres",
    url: dbUrl,
  });

  await datasource.initialize();

  // Set the role to anon
  await datasource.query("SET ROLE 'anon';");

  // Set session to the incorrect user id
  await datasource.query(`SET myapp.user_id = '123';`);

  const data = await datasource.query("SELECT * FROM todo_list;");
  console.log(data); // correctly returns []

  const db = await SqlDatabase.fromDataSourceParams({
    appDataSource: datasource,
    includesTables: ["todo_list"],
  });

  const model = new ChatOpenAI({
    modelName: "gpt-4-1106-preview",
    temperature: 0,
  });

  const toolkit = new SqlToolkit(db, model);
  const executor = createSqlAgent(model, toolkit);

  const input = `
    What are all the items on the todo list?
  `;

  const result = await executor.invoke({ input });

  console.log(result.output); // returns an array of items from the todo list. It should return []

  await datasource.destroy();
};

lookupData();


Example policy I used:
```sql
CREATE POLICY "only allow matching user access" ON public.todo_list
FOR SELECT USING (
    user_id = current_setting('myapp.user_id')
);

### Description

I am trying to set up an SQL agent so that it can only query the "current user's" data. 

At the point I will call `lookupData()` I know who the user is so my current approach is setting the session `myapp.user_id` variable and I created a policy around this in Supbase.

This works if I use `datasource` directly but when using the agent, it seems to unset the role and session. 

My ultimate goal is to try to enforce a where clause and back it up with RLS. 

Probably a better input example:

> What are all the items on the todo list where user_id is 124?

### System Info

n/a

[dosubot[bot]]

Hey @garmeeh! Great to see you back. Hope you've been doing well since our last interaction. 😊

Based on the context provided, it seems like the issue you're encountering is related to how sessions and roles are managed when using the SQL agent in LangChain JS.

The SQL agent in LangChain JS does not inherently manage sessions or roles. It is designed to execute SQL queries and return the results. Any session or role management needs to be handled externally before invoking the SQL agent.

In your case, you're setting the role and session using datasource.query(), which is correct. However, when you create the SQL agent and invoke it, it doesn't carry over the session or role you set earlier. This is because the SQL agent creates a new connection to the database each time it's invoked, and this new connection doesn't have the session or role information.

To resolve this issue, you need to set the session and role within the context of the SQL agent. Here's an example of how you can do this:

const executor = createSqlAgent(model, toolkit, {
  beforeInvoke: async () => {
    // Set the role to anon
    await datasource.query("SET ROLE 'anon';");

    // Set session to the incorrect user id
    await datasource.query(`SET myapp.user_id = '123';`);
  },
});

In this example, the beforeInvoke option is a function that gets called before the SQL agent is invoked. You can use this function to set up the session and role. This ensures that the session and role are set for the connection used by the SQL agent.

Please note that this is a workaround and might not be the most efficient solution, especially if you're invoking the SQL agent multiple times in a short period, as it sets the session and role for each invocation. A more efficient solution would be to enhance the SQL agent to support session and role management, but that would require changes to the LangChain JS framework itself.

I hope this helps! If you have any further questions, feel free to ask.

Sources

• langchain/src/retrievers/tests/supabase.int.test.ts

About Dosu

This response is meant to be useful and save you time. It is not meant to be a precise solution, but rather a starting point for your own research.

Help me be more useful! Please leave a 👍 if this is helpful and 👎 if it is irrelevant.

If you want to continue the conversation, start your reply with @dosu-bot.

[garmeeh]

import "dotenv/config";

import { ChatOpenAI } from "@langchain/openai";
import { SqlDatabase } from "langchain/sql_db";
import { createSqlAgent, SqlToolkit } from "langchain/agents/toolkits/sql";
import { DataSource } from "typeorm";

export const lookupData = async () => {
  const encodedPassword = encodeURIComponent(process.env.DB_PASSWORD || "");

  const dbUrl = `postgres://postgres.X:${encodedPassword}@X.supabase.com:5432/postgres`;

  const datasource = new DataSource({
    type: "postgres",
    url: dbUrl,
  });

  await datasource.initialize();

  const db = await SqlDatabase.fromDataSourceParams({
    appDataSource: datasource,
    includesTables: ["todo_list"],
  });

  const model = new ChatOpenAI({
    modelName: "gpt-4-1106-preview",
    temperature: 0,
  });

  const toolkit = new SqlToolkit(db, model);
  const executor = createSqlAgent(model, toolkit);

  const input = `
    What are all the items on the todo list?
  `;

  const result = await executor.invoke({ input });

  const { intermediateSteps } = result;
  const lastAction = intermediateSteps[intermediateSteps.length - 1].action;
  if (lastAction.tool === "query-sql") {
    const action = lastAction;
    await datasource.query("SET ROLE 'anon';");
    await datasource.query(`SET myapp.user_id = '123';`);
    const protectedQueryResults = await datasource.query(action.toolInput);
    console.log(protectedQueryResults); // correctly returns []
    /**
     * From this point on you need to then take the initial query and the protectedQueryResults,
     * pass them to an llm again to format to an answer for the user
     *
     * Also handle if the answer is "I don't know" or similar
     */
  }

  await datasource.destroy();
};

lookupData();

This code achieves the desired outcome.

I could be missing something simple that will make this far more optimal without the extra step.