mirror of
https://github.com/langbot-app/LangBot.git
synced 2026-08-09 12:40:59 +00:00
fix: allow trusted halfvec tenant search casts (#2402)
Co-authored-by: dadachann <185672915+dadachann@users.noreply.github.com>
This commit is contained in:
@@ -13,7 +13,7 @@ import typing
|
||||
import sqlalchemy
|
||||
import sqlalchemy.ext.asyncio as sqlalchemy_asyncio
|
||||
import sqlalchemy.orm as sqlalchemy_orm
|
||||
from pgvector.sqlalchemy import Vector
|
||||
from pgvector.sqlalchemy import HALFVEC, Vector
|
||||
from sqlalchemy.dialects.postgresql.dml import OnConflictDoNothing as PostgreSQLOnConflictDoNothing
|
||||
from sqlalchemy.dialects.postgresql.dml import OnConflictDoUpdate as PostgreSQLOnConflictDoUpdate
|
||||
from sqlalchemy.dialects.sqlite.dml import OnConflictDoNothing as SQLiteOnConflictDoNothing
|
||||
@@ -281,7 +281,7 @@ def _validate_scoped_sql_type(
|
||||
return
|
||||
seen.add(identity)
|
||||
|
||||
if type(sql_type) is Vector:
|
||||
if type(sql_type) in {Vector, HALFVEC}:
|
||||
return
|
||||
if not type(sql_type).__module__.startswith('sqlalchemy.'):
|
||||
raise ScopedSessionTransactionError('TenantUnitOfWork does not allow custom SQL types in public statements')
|
||||
@@ -462,7 +462,7 @@ def _validate_scoped_statement_call(args: tuple[typing.Any, ...], kwargs: dict[s
|
||||
if isinstance(element, sqlalchemy.sql.elements.BindParameter) and element.literal_execute:
|
||||
raise ScopedSessionTransactionError('TenantUnitOfWork does not allow literal-execute SQL parameters')
|
||||
|
||||
if isinstance(element, sqlalchemy.sql.elements.Cast) and type(element.type) is not Vector:
|
||||
if isinstance(element, sqlalchemy.sql.elements.Cast) and type(element.type) not in {Vector, HALFVEC}:
|
||||
raise ScopedSessionTransactionError(
|
||||
'TenantUnitOfWork only allows the trusted pgvector cast used by tenant vector search'
|
||||
)
|
||||
|
||||
Reference in New Issue
Block a user