fix: vectorizer_relationship with tuple table args (#805)

Author: Askir

projects/pgai/pgai/sqlalchemy/__init__.py

@@ -1,4 +1,4 @@
-from typing import Any, Generic, TypeVar, overload
+from typing import Any, Generic, TypeVar, cast, overload
 
 from pgvector.sqlalchemy import Vector  # type: ignore
 from sqlalchemy import ForeignKeyConstraint, Integer, Text, event, inspect
@@ -75,7 +75,18 @@ def _initialize_all(self):
             self.__get__(None, self.owner)
 
     def set_schemas_correctly(self, owner: type[DeclarativeBase]) -> None:
-        table_args_schema_name = getattr(owner, "__table_args__", {}).get("schema")
+        table_args = getattr(owner, "__table_args__", {})
+        table_args_schema_name: str | None = None
+
+        if isinstance(table_args, dict):
+            table_args_schema_name = cast(str | None, table_args.get("schema"))  # type: ignore
+        elif (
+            isinstance(table_args, tuple)
+            and len(table_args) > 0  # type: ignore
+            and isinstance(table_args[-1], dict)
+        ):
+            table_args_schema_name = cast(str | None, table_args[-1].get("schema"))  # type: ignore
+
         self.target_schema = (
             self.target_schema
             or table_args_schema_name

projects/pgai/tests/vectorizer/extensions/test_sqlalchemy_tuple_table_args.py

@@ -0,0 +1,29 @@
+from sqlalchemy import Column, Integer, Text, UniqueConstraint
+from sqlalchemy.orm import DeclarativeBase
+
+from pgai.sqlalchemy import vectorizer_relationship
+
+
+class Base(DeclarativeBase):
+    pass
+
+
+class FeatureWithTupleTableArgs(Base):
+    __tablename__ = "features_tuple"
+
+    id = Column(Integer, primary_key=True)
+    name = Column(Text, nullable=False)
+    tenant_id = Column(Integer, nullable=False)
+
+    __table_args__ = (UniqueConstraint("name", "tenant_id"),)
+
+    embeddings = vectorizer_relationship(
+        dimensions=1536, target_table="features_embeddings"
+    )
+
+
+def test_tuple_table_args():
+    FeatureWithTupleTableArgs()
+    embedding_class = FeatureWithTupleTableArgs.embeddings
+
+    assert embedding_class is not None