Row Level Security + Text-to-SQL LLM

postgresql://admin:abcd1234@localhost:5432/agent_info

pip install \
  sqlalchemy psycopg2-binary pydantic \
  llama-index phoenix \
  llama-index-llms-openvino-genai \
  llama-index-embeddings-openvino-genai

from sqlalchemy import (
    BigInteger,
    Column,
    MetaData,
    Numeric,
    Table,
    Text,
    create_engine,
    text,
)

ADMIN_DSN = "postgresql://admin:abcd1234@localhost:5432/agent_info"
engine = create_engine(ADMIN_DSN)

with engine.begin() as conn:
    conn.execute(text("DROP USER IF EXISTS app_user"))
    conn.execute(text("CREATE USER app_user WITH PASSWORD 'abcd1234'"))

    conn.execute(text("GRANT CONNECT ON DATABASE agent_info TO app_user"))
    conn.execute(text("GRANT USAGE ON SCHEMA public TO app_user"))

metadata = MetaData()

orders = Table(
    "orders",
    metadata,
    Column("order_id", BigInteger, primary_key=True, autoincrement=True),
    Column("sale_id", Text, nullable=False),
    Column("yearmonth", Text, nullable=False),
    Column("amount", Numeric(14, 2), nullable=False),
)

metadata.drop_all(engine)
metadata.create_all(engine)

with engine.begin() as conn:
    conn.execute(
        orders.insert(),
        [
            {"sale_id": "sale_001", "yearmonth": "202401", "amount": 1_000_000},
            {"sale_id": "sale_001", "yearmonth": "202402", "amount": 1_500_000},
            {"sale_id": "sale_002", "yearmonth": "202401", "amount": 2_000_000},
            {"sale_id": "sale_002", "yearmonth": "202402", "amount": 2_500_000},
            {"sale_id": "sale_003", "yearmonth": "202401", "amount": 3_000_000},
        ],
    )
    conn.execute(text("GRANT SELECT ON orders TO app_user"))

from sqlalchemy import text

with engine.begin() as conn:
    conn.execute(text("ALTER TABLE orders ENABLE ROW LEVEL SECURITY"))

    conn.execute(
        text(
            """
            DROP POLICY IF EXISTS sale_read_own_orders ON orders;
            """
        )
    )

    conn.execute(
        text(
            """
            CREATE POLICY sale_read_own_orders
            ON orders
            FOR SELECT
            USING (
                sale_id = current_setting('app.current_sale_id', true)::text
            );
            """
        )
    )

from sqlalchemy import create_engine, text

engine_app = create_engine(
    "postgresql+psycopg2://app_user:abcd1234@localhost:5432/agent_info"
)

with engine_app.begin() as conn:
    conn.execute(text("SET LOCAL app.current_sale_id = 'sale_001'"))
    rows = conn.execute(text("SELECT * FROM orders")).fetchall()
    print("sale_001 sees:", rows)

with engine_app.begin() as conn:
    conn.execute(text("SET LOCAL app.current_sale_id = 'sale_002'"))
    rows = conn.execute(text("SELECT * FROM orders")).fetchall()
    print("sale_002 sees:", rows)

with engine_app.begin() as conn:
    rows = conn.execute(text("SELECT * FROM orders")).fetchall()
    print("anonymous sees:", rows)

import re
from dataclasses import dataclass
from functools import partial
from typing import Any, Callable, Dict, List, Optional, Tuple, Type

from pydantic import BaseModel, Field, model_validator
from sqlalchemy.engine.base import Connection
from sqlalchemy.sql.elements import TextClause
from sqlalchemy import text


@dataclass(frozen=True)
class _GUCRef:
    guc_name: str
    param_name: str
    pg_type: str
    python_type: Type
    raw_expression: str


class RowSecurityPolicy(BaseModel):
    name: str
    table: str
    commands: List[str] = Field(default_factory=lambda: ["ALL"])
    roles: List[str] = Field(default_factory=lambda: ["PUBLIC"])
    using: str | None = None
    with_check: str | None = None
    permissive: bool | None = None

    guc_refs: List[_GUCRef] = Field(default_factory=list)

    _CURRENT_SETTING_RE = re.compile(
        r"""
        current_setting
        \(
            \s*'(?P<guc>[a-zA-Z0-9_.]+)'\s*
            (?:,\s*true\s*)?
        \)
        \s*::\s*(?P<cast>[a-zA-Z0-9_]+)
        """,
        re.VERBOSE | re.IGNORECASE,
    )

    _PG_TO_PY_TYPE: Dict[str, Type] = {
        "int": int,
        "integer": int,
        "bigint": int,
        "smallint": int,
        "uuid": str,
        "text": str,
        "varchar": str,
        "boolean": bool,
        "bool": bool,
    }

    @model_validator(mode="after")
    def _extract_gucs(self) -> "RowSecurityPolicy":
        refs: Dict[str, _GUCRef] = {}

        def scan(expr: str):
            for m in self._CURRENT_SETTING_RE.finditer(expr):
                guc = m.group("guc")
                cast = m.group("cast").lower()

                if cast not in self._PG_TO_PY_TYPE:
                    raise ValueError(
                        f"Unsupported cast type '{cast}' for current_setting('{guc}')"
                    )

                refs[guc] = _GUCRef(
                    guc_name=guc,
                    param_name=guc.split(".")[-1],
                    pg_type=cast,
                    python_type=self._PG_TO_PY_TYPE[cast],
                    raw_expression=m.group(0),
                )

        if self.using:
            scan(self.using)

        if self.with_check:
            scan(self.with_check)

        self.guc_refs = list(refs.values())
        return self

    def get_setup_func(
        self,
        params: Dict[str, Any],
    ) -> Callable[[Connection], None]:
        args = []

        for ref in self.guc_refs:
            name = ref.param_name

            if name not in params:
                raise KeyError(
                    f"Missing GUC parameter '{name}' for policy '{self.name}'"
                )

            value = params[name]

            if value is None:
                raise ValueError(f"GUC parameter '{name}' cannot be NULL")

            if not isinstance(value, ref.python_type):
                raise TypeError(
                    f"GUC parameter '{name}' must be "
                    f"{ref.python_type.__name__}, "
                    f"got {type(value).__name__}"
                )

            args.append(
                (
                    text(f"SET LOCAL {ref.guc_name} = :{ref.param_name}"),
                    {name: value},
                )
            )

        def set_transaction_variable(
            conn: Connection, args: list[tuple[TextClause, dict]]
        ):
            for text_clause, param_dict in args:
                conn.execute(text_clause, param_dict)

        return partial(set_transaction_variable, args=args)

policy = RowSecurityPolicy(
    name="sale_read_own_orders",
    table="orders",
    commands=["SELECT"],
    using="""
        sale_id = current_setting('app.current_sale_id', true)::text
    """,
)

from typing import Dict, Tuple
from sqlalchemy import Engine
from sqlalchemy.exc import OperationalError, ProgrammingError
from sqlalchemy.sql.elements import TextClause
from llama_index.core import SQLDatabase


class RowLevelSecuritySQLDatabase(SQLDatabase):
    def __init__(
        self,
        engine: Engine,
        schema: str | None = None,
        metadata: MetaData | None = None,
        ignore_tables: List[str] | None = None,
        include_tables: List[str] | None = None,
        sample_rows_in_table_info: int = 3,
        indexes_in_table_info: bool = False,
        custom_table_info: dict | None = None,
        view_support: bool = False,
        max_string_length: int = 300,
        setup_funcs: Optional[list[Callable[[Connection], None]]] = None,
    ):
        super().__init__(
            engine,
            schema,
            metadata,
            ignore_tables,
            include_tables,
            sample_rows_in_table_info,
            indexes_in_table_info,
            custom_table_info,
            view_support,
            max_string_length,
        )
        self.setup_funcs = setup_funcs

    def run_sql(self, command: str) -> Tuple[str, Dict]:
        with self._engine.begin() as connection:
            if self.setup_funcs is not None:
                for setup_func in self.setup_funcs:
                    setup_func(connection)
            try:
                if self._schema:
                    command = command.replace("FROM ", f"FROM {self._schema}.")
                    command = command.replace("JOIN ", f"JOIN {self._schema}.")
                cursor = connection.execute(text(command))
            except (ProgrammingError, OperationalError) as exc:
                raise NotImplementedError(
                    f"Statement {command!r} is invalid SQL.\nError: {exc.orig}"
                ) from exc
            if cursor.returns_rows:
                result = cursor.fetchall()
                truncated_results = []
                for row in result:
                    truncated_row = tuple(
                        self.truncate_word(column, length=self._max_string_length)
                        for column in row
                    )
                    truncated_results.append(truncated_row)
                return str(truncated_results), {
                    "result": truncated_results,
                    "col_keys": list(cursor.keys()),
                }
        return "", {}

engine_agent = create_engine(
    "postgresql+psycopg2://app_user:abcd1234@localhost:5432/agent_info"
)

sql_database = RowLevelSecuritySQLDatabase(
    engine_agent,
    include_tables=["orders"],
    setup_funcs=[policy.get_setup_func({"current_sale_id": "sale_001"})],
)
print(sql_database.run_sql("SELECT * FROM orders"))

import phoenix as px
import llama_index
from llama_index.core.query_engine import NLSQLTableQueryEngine
from llama_index.embeddings.openvino_genai import OpenVINOGENAIEmbedding
from llama_index.llms.openvino_genai import OpenVINOGenAILLM

px.launch_app()
llama_index.core.set_global_handler("arize_phoenix")

llm = OpenVINOGenAILLM(
    model_path="./weights/OpenVINO/Phi-3.5-mini-instruct-int4-ov/",
    device="CPU",
)

ov_embed_model = OpenVINOGENAIEmbedding(
    model_path="./weights/OpenVINO/bge_ov",
    device="CPU",
)

query_engine = NLSQLTableQueryEngine(
    llm=llm,
    embed_model=ov_embed_model,
    sql_database=sql_database,
    tables=["orders"],
)

query_str = "Sum total amount"
response = query_engine.query(query_str)
print(response.response)