Sqlalchemy 2.x (#508)

* SQLAlchemy 2.0 migration WIP * fixed remaining failing tests * more cleanup and documentatiion
toluaina · Dec 5, 2023 · f9a37eb · f9a37eb
1 parent 40e6d8e
commit f9a37eb
Show file tree

Hide file tree

Showing 29 changed files with 650 additions and 521 deletions.
diff --git a/bin/parallel_sync b/bin/parallel_sync
@@ -146,7 +146,7 @@ def fetch_tasks(
     )
     page, row = read_ctid(name=name)
     statement: sa.sql.Select = sa.select(
-        [
+        *[
             sa.literal_column("1").label("x"),
             sa.literal_column("1").label("y"),
             sa.column("ctid"),

diff --git a/examples/airbnb/schema.py b/examples/airbnb/schema.py
@@ -2,45 +2,49 @@
 
 import click
 import sqlalchemy as sa
-from sqlalchemy.ext.declarative import declarative_base
+from sqlalchemy.orm import DeclarativeBase, Mapped, mapped_column
 from sqlalchemy.schema import UniqueConstraint
 
 from pgsync.base import create_database, pg_engine
 from pgsync.helper import teardown
 from pgsync.utils import config_loader, get_config
 
-Base = declarative_base()
+
+class Base(DeclarativeBase):
+    pass
 
 
 class User(Base):
     __tablename__ = "user"
     __table_args__ = (UniqueConstraint("email"),)
-    id = sa.Column(sa.Integer, primary_key=True)
-    email = sa.Column(sa.String, unique=True, nullable=False)
+    id: Mapped[int] = mapped_column(sa.Integer, primary_key=True)
+    email: Mapped[str] = mapped_column(sa.String, unique=True, nullable=False)
 
 
 class Host(Base):
     __tablename__ = "host"
     __table_args__ = (UniqueConstraint("email"),)
-    id = sa.Column(sa.Integer, primary_key=True)
-    email = sa.Column(sa.String, unique=True, nullable=False)
+    id: Mapped[int] = mapped_column(sa.Integer, primary_key=True)
+    email: Mapped[str] = mapped_column(sa.String, unique=True, nullable=False)
 
 
 class Country(Base):
     __tablename__ = "country"
     __table_args__ = (UniqueConstraint("name", "country_code"),)
-    id = sa.Column(sa.Integer, primary_key=True)
-    name = sa.Column(sa.String, nullable=False)
-    country_code = sa.Column(sa.String, nullable=False)
+    id: Mapped[int] = mapped_column(sa.Integer, primary_key=True)
+    name: Mapped[str] = mapped_column(sa.String, nullable=False)
+    country_code: Mapped[str] = mapped_column(sa.String, nullable=False)
 
 
 class City(Base):
     __tablename__ = "city"
     __table_args__ = (UniqueConstraint("name", "country_id"),)
-    id = sa.Column(sa.Integer, primary_key=True)
-    name = sa.Column(sa.String, nullable=False)
-    country_id = sa.Column(sa.Integer, sa.ForeignKey(Country.id))
-    country = sa.orm.relationship(
+    id: Mapped[int] = mapped_column(sa.Integer, primary_key=True)
+    name: Mapped[str] = mapped_column(sa.String, nullable=False)
+    country_id: Mapped[int] = mapped_column(
+        sa.Integer, sa.ForeignKey(Country.id)
+    )
+    country: Mapped[Country] = sa.orm.relationship(
         Country,
         backref=sa.orm.backref("country"),
     )
@@ -49,15 +53,15 @@ class City(Base):
 class Place(Base):
     __tablename__ = "place"
     __table_args__ = (UniqueConstraint("host_id", "address", "city_id"),)
-    id = sa.Column(sa.Integer, primary_key=True)
-    host_id = sa.Column(sa.Integer, sa.ForeignKey(Host.id))
-    address = sa.Column(sa.String, nullable=False)
-    city_id = sa.Column(sa.Integer, sa.ForeignKey(City.id))
-    host = sa.orm.relationship(
+    id: Mapped[int] = mapped_column(sa.Integer, primary_key=True)
+    host_id: Mapped[int] = mapped_column(sa.Integer, sa.ForeignKey(Host.id))
+    address: Mapped[str] = mapped_column(sa.String, nullable=False)
+    city_id: Mapped[int] = mapped_column(sa.Integer, sa.ForeignKey(City.id))
+    host: Mapped[Host] = sa.orm.relationship(
         Host,
         backref=sa.orm.backref("host"),
     )
-    city = sa.orm.relationship(
+    city: Mapped[City] = sa.orm.relationship(
         City,
         backref=sa.orm.backref("city"),
     )
@@ -66,25 +70,33 @@ class Place(Base):
 class Booking(Base):
     __tablename__ = "booking"
     __table_args__ = (UniqueConstraint("user_id", "place_id", "start_date"),)
-    id = sa.Column(sa.Integer, primary_key=True)
-    user_id = sa.Column(sa.Integer, sa.ForeignKey(User.id))
-    place_id = sa.Column(sa.Integer, sa.ForeignKey(Place.id))
-    start_date = sa.Column(sa.DateTime, default=datetime.now())
-    end_date = sa.Column(sa.DateTime, default=datetime.now())
-    price_per_night = sa.Column(sa.Float, default=0)
-    num_nights = sa.Column(sa.Integer, nullable=False, default=1)
-    user = sa.orm.relationship(User)
-    place = sa.orm.relationship(Place)
+    id: Mapped[int] = mapped_column(sa.Integer, primary_key=True)
+    user_id: Mapped[int] = mapped_column(sa.Integer, sa.ForeignKey(User.id))
+    place_id: Mapped[int] = mapped_column(sa.Integer, sa.ForeignKey(Place.id))
+    start_date: Mapped[datetime] = mapped_column(
+        sa.DateTime, default=datetime.now()
+    )
+    end_date: Mapped[datetime] = mapped_column(
+        sa.DateTime, default=datetime.now()
+    )
+    price_per_night: Mapped[float] = mapped_column(sa.Float, default=0)
+    num_nights: Mapped[int] = mapped_column(
+        sa.Integer, nullable=False, default=1
+    )
+    user: Mapped[User] = sa.orm.relationship(User)
+    place: Mapped[Place] = sa.orm.relationship(Place)
 
 
 class Review(Base):
     __tablename__ = "review"
     __table_args__ = (UniqueConstraint("booking_id"),)
-    id = sa.Column(sa.Integer, primary_key=True)
-    booking_id = sa.Column(sa.Integer, sa.ForeignKey(Booking.id))
-    rating = sa.Column(sa.SmallInteger, nullable=True)
-    review_body = sa.Column(sa.Text, nullable=True)
-    booking = sa.orm.relationship(
+    id: Mapped[int] = mapped_column(sa.Integer, primary_key=True)
+    booking_id: Mapped[int] = mapped_column(
+        sa.Integer, sa.ForeignKey(Booking.id)
+    )
+    rating: Mapped[int] = mapped_column(sa.SmallInteger, nullable=True)
+    review_body: Mapped[str] = mapped_column(sa.Text, nullable=True)
+    booking: Mapped[Booking] = sa.orm.relationship(
         Booking,
         backref=sa.orm.backref("booking"),
     )

diff --git a/examples/ancestry/schema.py b/examples/ancestry/schema.py
@@ -1,51 +1,69 @@
 import click
 import sqlalchemy as sa
-from sqlalchemy.ext.declarative import declarative_base
+from sqlalchemy.orm import DeclarativeBase, Mapped, mapped_column
 
 from pgsync.base import create_database, pg_engine
 from pgsync.helper import teardown
 from pgsync.utils import config_loader, get_config
 
-Base = declarative_base()
+
+class Base(DeclarativeBase):
+    pass
 
 
 class Parent(Base):
     __tablename__ = "parent"
     __table_args__ = ()
-    id = sa.Column(sa.Integer, primary_key=True, autoincrement=True)
-    name = sa.Column(sa.String)
+    id: Mapped[int] = mapped_column(
+        sa.Integer, primary_key=True, autoincrement=True
+    )
+    name: Mapped[str] = mapped_column(sa.String)
 
 
 class Surrogate(Base):
     __tablename__ = "surrogate"
     __table_args__ = ()
-    id = sa.Column(sa.Integer, primary_key=True, autoincrement=True)
-    name = sa.Column(sa.String)
-    parent_id = sa.Column(sa.Integer, sa.ForeignKey(Parent.id))
+    id: Mapped[int] = mapped_column(
+        sa.Integer, primary_key=True, autoincrement=True
+    )
+    name: Mapped[str] = mapped_column(sa.String)
+    parent_id: Mapped[int] = mapped_column(
+        sa.Integer, sa.ForeignKey(Parent.id)
+    )
 
 
 class Child(Base):
     __tablename__ = "child"
     __table_args__ = ()
-    id = sa.Column(sa.Integer, primary_key=True, autoincrement=True)
-    name = sa.Column(sa.String)
-    parent_id = sa.Column(sa.Integer, sa.ForeignKey(Surrogate.id))
+    id: Mapped[int] = mapped_column(
+        sa.Integer, primary_key=True, autoincrement=True
+    )
+    name: Mapped[str] = mapped_column(sa.String)
+    parent_id: Mapped[int] = mapped_column(
+        sa.Integer, sa.ForeignKey(Surrogate.id)
+    )
 
 
 class GrandChild(Base):
     __tablename__ = "grand_child"
     __table_args__ = ()
-    id = sa.Column(sa.Integer, primary_key=True, autoincrement=True)
-    name = sa.Column(sa.String)
-    parent_id = sa.Column(sa.Integer, sa.ForeignKey(Child.id))
+    id: Mapped[int] = mapped_column(
+        sa.Integer, primary_key=True, autoincrement=True
+    )
+    name: Mapped[str] = mapped_column(sa.String)
+    parent_id: Mapped[int] = mapped_column(sa.Integer, sa.ForeignKey(Child.id))
 
 
 class GreatGrandChild(Base):
     __tablename__ = "great_grand_child"
     __table_args__ = ()
-    id = sa.Column(sa.Integer, primary_key=True, autoincrement=True)
-    name = sa.Column(sa.String)
-    parent_id = sa.Column(sa.Integer, sa.ForeignKey(GrandChild.id))
+    id: Mapped[int] = mapped_column(
+        sa.Integer, primary_key=True, autoincrement=True
+    )
+    name: Mapped[str] = mapped_column(sa.String)
+    parent_id: Mapped[int] = mapped_column(
+        sa.Integer, sa.ForeignKey(GrandChild.id)
+    )
 
 
 def setup(config: str) -> None: