feat: add tests and CHANGES.rst

Author: danplischke

CHANGES.rst

@@ -1,6 +1,10 @@
 Version history
 ===============
 
+**UNRELEASED**
+- Fix Postgres TSVECTOR handling (PR by @danplischke)
+- Fix SQLModel metadata reference (PR by @danplischke)
+
 **3.1.0**
 
 - Type annotations for ARRAY column attributes now include the Python type of

src/sqlacodegen/generators.py

@@ -1437,7 +1437,7 @@ def generate_base(self) -> None:
         self.base = Base(
             literal_imports=[],
             declarations=[],
-            metadata_ref="",
+            metadata_ref="SQLModel.metadata",
         )
 
     def collect_imports(self, models: Iterable[Model]) -> None:

tests/test_generator_dataclass.py

@@ -2,7 +2,7 @@
 
 import pytest
 from _pytest.fixtures import FixtureRequest
-from sqlalchemy.dialects.postgresql import UUID
+from sqlalchemy.dialects.postgresql import TSVECTOR, UUID
 from sqlalchemy.engine import Engine
 from sqlalchemy.schema import Column, ForeignKeyConstraint, MetaData, Table
 from sqlalchemy.sql.expression import text
@@ -267,3 +267,34 @@ class Simple(Base):
             id: Mapped[uuid.UUID] = mapped_column(UUID, primary_key=True)
         """,
     )
+
+
+def test_tsvector(generator: CodeGenerator) -> None:
+    Table(
+        "simple_tsvector",
+        generator.metadata,
+        Column("id", UUID, primary_key=True),
+        Column("vector", TSVECTOR),
+    )
+
+    validate_code(
+        generator.generate(),
+        """\
+        from typing import Optional
+        import uuid
+
+        from sqlalchemy import UUID
+        from sqlalchemy.dialects.postgresql import TSVECTOR
+        from sqlalchemy.orm import DeclarativeBase, Mapped, MappedAsDataclass, mapped_column
+
+        class Base(MappedAsDataclass, DeclarativeBase):
+            pass
+
+
+        class SimpleTsvector(Base):
+            __tablename__ = 'simple_tsvector'
+
+            id: Mapped[uuid.UUID] = mapped_column(UUID, primary_key=True)
+            vector: Mapped[Optional[str]] = mapped_column(TSVECTOR)
+        """,
+    )

tests/test_generator_declarative.py

@@ -4,7 +4,7 @@
 from _pytest.fixtures import FixtureRequest
 from sqlalchemy import BIGINT, PrimaryKeyConstraint
 from sqlalchemy.dialects import postgresql
-from sqlalchemy.dialects.postgresql import JSON, JSONB
+from sqlalchemy.dialects.postgresql import JSON, JSONB, TSVECTOR
 from sqlalchemy.engine import Engine
 from sqlalchemy.schema import (
     CheckConstraint,
@@ -1706,3 +1706,33 @@ class TestDomainJson(Base):
     foo: Mapped[Optional[dict]] = mapped_column(DOMAIN('domain_json', {domain_type.__name__}(astext_type=Text(length=128)), not_null=False))
 """,
     )
+
+
+def test_tsvector(generator: CodeGenerator) -> None:
+    Table(
+        "test_tsvector",
+        generator.metadata,
+        Column("id", BIGINT, primary_key=True),
+        Column("vector", TSVECTOR()),
+    )
+
+    validate_code(
+        generator.generate(),
+        """\
+from typing import Optional
+
+from sqlalchemy import BigInteger
+from sqlalchemy.dialects.postgresql import TSVECTOR
+from sqlalchemy.orm import DeclarativeBase, Mapped, mapped_column
+
+class Base(DeclarativeBase):
+    pass
+
+
+class TestTsvector(Base):
+    __tablename__ = 'test_tsvector'
+
+    id: Mapped[int] = mapped_column(BigInteger, primary_key=True)
+    vector: Mapped[Optional[str]] = mapped_column(TSVECTOR)
+""",
+    )

tests/test_generator_sqlmodel.py

@@ -3,6 +3,7 @@
 import pytest
 from _pytest.fixtures import FixtureRequest
 from sqlalchemy import Uuid
+from sqlalchemy.dialects.postgresql import TSVECTOR
 from sqlalchemy.engine import Engine
 from sqlalchemy.schema import (
     CheckConstraint,
@@ -204,3 +205,30 @@ class SimpleUuid(SQLModel, table=True):
                 id: uuid.UUID = Field(sa_column=Column('id', Uuid, primary_key=True))
         """,
     )
+
+
+def test_tsvector(generator: CodeGenerator) -> None:
+    Table(
+        "simple_tsvector",
+        generator.metadata,
+        Column("id", Uuid, primary_key=True),
+        Column("search", TSVECTOR),
+    )
+
+    validate_code(
+        generator.generate(),
+        """\
+        from typing import Optional
+        import uuid
+
+        from sqlalchemy import Column, Uuid
+        from sqlalchemy.dialects.postgresql import TSVECTOR
+        from sqlmodel import Field, SQLModel
+
+        class SimpleTsvector(SQLModel, table=True):
+            __tablename__ = 'simple_tsvector'
+
+            id: uuid.UUID = Field(sa_column=Column('id', Uuid, primary_key=True))
+            search: Optional[str] = Field(default=None, sa_column=Column('search', TSVECTOR))
+        """,
+    )

tests/test_generator_tables.py

@@ -1090,3 +1090,29 @@ def test_postgresql_sequence_with_schema(
         )
         """,
     )
+
+
+def test_tsvector(generator: CodeGenerator) -> None:
+    Table(
+        "simple_items",
+        generator.metadata,
+        Column("id", INTEGER, primary_key=True),
+        Column("vector", postgresql.TSVECTOR),
+    )
+
+    validate_code(
+        generator.generate(),
+        """\
+        from sqlalchemy import Column, Integer, MetaData, Table
+        from sqlalchemy.dialects.postgresql import TSVECTOR
+
+        metadata = MetaData()
+
+
+        t_simple_items = Table(
+            'simple_items', metadata,
+            Column('id', Integer, primary_key=True),
+            Column('vector', TSVECTOR)
+        )
+        """,
+    )