fix: SQLAlchemy 2.0 compatibility for chapter_06_uow_exercise

- Replace deprecated mapper() with registry.map_imperatively() - Replace session.query() with session.scalars(select()) - Wrap raw SQL strings with text() - Replace clear_mappers() with mapper_registry.dispose() - Update requirements.txt to require sqlalchemy>=2.0
cosmicpython · rickywesker · Feb 8, 2019 · Mar 8, 2019 · Mar 8, 2019 · Mar 8, 2019
commit 00f1574723251a53ce2462be71a4a46c98ad6bd3
diff --git a/requirements.txt b/requirements.txt
@@ -1,5 +1,5 @@
 # app
-sqlalchemy
+sqlalchemy>=2.0
 flask
 psycopg2-binary
 

diff --git a/src/allocation/adapters/orm.py b/src/allocation/adapters/orm.py
@@ -1,10 +1,11 @@
 from sqlalchemy import Table, MetaData, Column, Integer, String, Date, ForeignKey
-from sqlalchemy.orm import mapper, relationship
+from sqlalchemy.orm import registry, relationship
 
 from allocation.domain import model
 
 
-metadata = MetaData()
+mapper_registry = registry()
+metadata = mapper_registry.metadata
 
 order_lines = Table(
     "order_lines",
@@ -35,8 +36,8 @@
 
 
 def start_mappers():
-    lines_mapper = mapper(model.OrderLine, order_lines)
-    mapper(
+    lines_mapper = mapper_registry.map_imperatively(model.OrderLine, order_lines)
+    mapper_registry.map_imperatively(
         model.Batch,
         batches,
         properties={

diff --git a/src/allocation/adapters/repository.py b/src/allocation/adapters/repository.py
@@ -1,4 +1,5 @@
 import abc
+from sqlalchemy import select
 from allocation.domain import model
 
 
@@ -20,7 +21,9 @@ def add(self, batch):
         self.session.add(batch)
 
     def get(self, reference):
-        return self.session.query(model.Batch).filter_by(reference=reference).one()
+        return self.session.scalars(
+            select(model.Batch).filter_by(reference=reference)
+        ).one()
 
     def list(self):
-        return self.session.query(model.Batch).all()
+        return self.session.scalars(select(model.Batch)).all()
diff --git a/tests/conftest.py b/tests/conftest.py
@@ -7,9 +7,9 @@
 from requests.exceptions import ConnectionError
 from sqlalchemy.exc import OperationalError
 from sqlalchemy import create_engine
-from sqlalchemy.orm import sessionmaker, clear_mappers
+from sqlalchemy.orm import sessionmaker
 
-from allocation.adapters.orm import metadata, start_mappers
+from allocation.adapters.orm import metadata, start_mappers, mapper_registry
 from allocation import config
 
 
@@ -24,7 +24,7 @@ def in_memory_db():
 def session_factory(in_memory_db):
     start_mappers()
     yield sessionmaker(bind=in_memory_db)
-    clear_mappers()
+    mapper_registry.dispose()
 
 
 @pytest.fixture
@@ -65,7 +65,7 @@ def postgres_db():
 def postgres_session(postgres_db):
     start_mappers()
     yield sessionmaker(bind=postgres_db)()
-    clear_mappers()
+    mapper_registry.dispose()
 
 
 @pytest.fixture

diff --git a/tests/integration/test_orm.py b/tests/integration/test_orm.py
@@ -1,54 +1,61 @@
+from sqlalchemy import text, select
 from allocation.domain import model
 from datetime import date
 
 
 def test_orderline_mapper_can_load_lines(session):
     session.execute(
-        "INSERT INTO order_lines (orderid, sku, qty) VALUES "
-        '("order1", "RED-CHAIR", 12),'
-        '("order1", "RED-TABLE", 13),'
-        '("order2", "BLUE-LIPSTICK", 14)'
+        text(
+            "INSERT INTO order_lines (orderid, sku, qty) VALUES "
+            '("order1", "RED-CHAIR", 12),'
+            '("order1", "RED-TABLE", 13),'
+            '("order2", "BLUE-LIPSTICK", 14)'
+        )
     )
     expected = [
         model.OrderLine("order1", "RED-CHAIR", 12),
         model.OrderLine("order1", "RED-TABLE", 13),
         model.OrderLine("order2", "BLUE-LIPSTICK", 14),
     ]
-    assert session.query(model.OrderLine).all() == expected
+    assert session.scalars(select(model.OrderLine)).all() == expected
 
 
 def test_orderline_mapper_can_save_lines(session):
     new_line = model.OrderLine("order1", "DECORATIVE-WIDGET", 12)
     session.add(new_line)
     session.commit()
 
-    rows = list(session.execute('SELECT orderid, sku, qty FROM "order_lines"'))
+    rows = list(session.execute(text('SELECT orderid, sku, qty FROM "order_lines"')))
     assert rows == [("order1", "DECORATIVE-WIDGET", 12)]
 
 
 def test_retrieving_batches(session):
     session.execute(
-        "INSERT INTO batches (reference, sku, _purchased_quantity, eta)"
-        ' VALUES ("batch1", "sku1", 100, null)'
+        text(
+            "INSERT INTO batches (reference, sku, _purchased_quantity, eta)"
+            ' VALUES ("batch1", "sku1", 100, null)'
+        )
     )
     session.execute(
-        "INSERT INTO batches (reference, sku, _purchased_quantity, eta)"
-        ' VALUES ("batch2", "sku2", 200, "2011-04-11")'
+        text(
+            "INSERT INTO batches (reference, sku, _purchased_quantity, eta)"
+            ' VALUES ("batch2", "sku2", 200, "2011-04-11")'
+        )
     )
     expected = [
         model.Batch("batch1", "sku1", 100, eta=None),
         model.Batch("batch2", "sku2", 200, eta=date(2011, 4, 11)),
     ]
 
-    assert session.query(model.Batch).all() == expected
+    assert session.scalars(select(model.Batch)).all() == expected
 
 
 def test_saving_batches(session):
     batch = model.Batch("batch1", "sku1", 100, eta=None)
     session.add(batch)
     session.commit()
     rows = session.execute(
-        'SELECT reference, sku, _purchased_quantity, eta FROM "batches"'
+        text('SELECT reference, sku, _purchased_quantity, eta FROM "batches"')
     )
     assert list(rows) == [("batch1", "sku1", 100, None)]
 
@@ -59,31 +66,33 @@ def test_saving_allocations(session):
     batch.allocate(line)
     session.add(batch)
     session.commit()
-    rows = list(session.execute('SELECT orderline_id, batch_id FROM "allocations"'))
+    rows = list(session.execute(text('SELECT orderline_id, batch_id FROM "allocations"')))
     assert rows == [(batch.id, line.id)]
 
 
 def test_retrieving_allocations(session):
     session.execute(
-        'INSERT INTO order_lines (orderid, sku, qty) VALUES ("order1", "sku1", 12)'
+        text('INSERT INTO order_lines (orderid, sku, qty) VALUES ("order1", "sku1", 12)')
     )
     [[olid]] = session.execute(
-        "SELECT id FROM order_lines WHERE orderid=:orderid AND sku=:sku",
+        text("SELECT id FROM order_lines WHERE orderid=:orderid AND sku=:sku"),
         dict(orderid="order1", sku="sku1"),
     )
     session.execute(
-        "INSERT INTO batches (reference, sku, _purchased_quantity, eta)"
-        ' VALUES ("batch1", "sku1", 100, null)'
+        text(
+            "INSERT INTO batches (reference, sku, _purchased_quantity, eta)"
+            ' VALUES ("batch1", "sku1", 100, null)'
+        )
     )
     [[bid]] = session.execute(
-        "SELECT id FROM batches WHERE reference=:ref AND sku=:sku",
+        text("SELECT id FROM batches WHERE reference=:ref AND sku=:sku"),
         dict(ref="batch1", sku="sku1"),
     )
     session.execute(
-        "INSERT INTO allocations (orderline_id, batch_id) VALUES (:olid, :bid)",
+        text("INSERT INTO allocations (orderline_id, batch_id) VALUES (:olid, :bid)"),
         dict(olid=olid, bid=bid),
     )
 
-    batch = session.query(model.Batch).one()
+    batch = session.scalars(select(model.Batch)).one()
 
     assert batch._allocations == {model.OrderLine("order1", "sku1", 12)}
diff --git a/tests/integration/test_repository.py b/tests/integration/test_repository.py
@@ -1,4 +1,5 @@
 # pylint: disable=protected-access
+from sqlalchemy import text
 from allocation.domain import model
 from allocation.adapters import repository
 
@@ -11,40 +12,46 @@ def test_repository_can_save_a_batch(session):
     session.commit()
 
     rows = session.execute(
-        'SELECT reference, sku, _purchased_quantity, eta FROM "batches"'
+        text('SELECT reference, sku, _purchased_quantity, eta FROM "batches"')
     )
     assert list(rows) == [("batch1", "RUSTY-SOAPDISH", 100, None)]
 
 
 def insert_order_line(session):
     session.execute(
-        "INSERT INTO order_lines (orderid, sku, qty)"
-        ' VALUES ("order1", "GENERIC-SOFA", 12)'
+        text(
+            "INSERT INTO order_lines (orderid, sku, qty)"
+            ' VALUES ("order1", "GENERIC-SOFA", 12)'
+        )
     )
     [[orderline_id]] = session.execute(
-        "SELECT id FROM order_lines WHERE orderid=:orderid AND sku=:sku",
+        text("SELECT id FROM order_lines WHERE orderid=:orderid AND sku=:sku"),
         dict(orderid="order1", sku="GENERIC-SOFA"),
     )
     return orderline_id
 
 
 def insert_batch(session, batch_id):
     session.execute(
-        "INSERT INTO batches (reference, sku, _purchased_quantity, eta)"
-        ' VALUES (:batch_id, "GENERIC-SOFA", 100, null)',
+        text(
+            "INSERT INTO batches (reference, sku, _purchased_quantity, eta)"
+            ' VALUES (:batch_id, "GENERIC-SOFA", 100, null)'
+        ),
         dict(batch_id=batch_id),
     )
     [[batch_id]] = session.execute(
-        'SELECT id FROM batches WHERE reference=:batch_id AND sku="GENERIC-SOFA"',
+        text('SELECT id FROM batches WHERE reference=:batch_id AND sku="GENERIC-SOFA"'),
         dict(batch_id=batch_id),
     )
     return batch_id
 
 
 def insert_allocation(session, orderline_id, batch_id):
     session.execute(
-        "INSERT INTO allocations (orderline_id, batch_id)"
-        " VALUES (:orderline_id, :batch_id)",
+        text(
+            "INSERT INTO allocations (orderline_id, batch_id)"
+            " VALUES (:orderline_id, :batch_id)"
+        ),
         dict(orderline_id=orderline_id, batch_id=batch_id),
     )
 

diff --git a/tests/integration/test_uow.py b/tests/integration/test_uow.py
@@ -1,24 +1,29 @@
 import pytest
+from sqlalchemy import text
 from allocation.domain import model
-from allocation.domain import model
-from allocation.domain import model
 from allocation.service_layer import unit_of_work
 
 
 def insert_batch(session, ref, sku, qty, eta):
     session.execute(
-        "INSERT INTO batches (reference, sku, _purchased_quantity, eta)"
-        " VALUES (:ref, :sku, :qty, :eta)",
+        text(
+            "INSERT INTO batches (reference, sku, _purchased_quantity, eta)"
+            " VALUES (:ref, :sku, :qty, :eta)"
+        ),
         dict(ref=ref, sku=sku, qty=qty, eta=eta),
     )
 
 
 def get_allocated_batch_ref(session, orderid, sku):
     [[orderlineid]] = session.execute(
-        "SELECT id FROM order_lines WHERE orderid=:orderid AND sku=:sku",
+        text("SELECT id FROM order_lines WHERE orderid=:orderid AND sku=:sku"),
         dict(orderid=orderid, sku=sku),
     )
     [[batchref]] = session.execute(
-        "SELECT b.reference FROM allocations JOIN batches AS b ON batch_id = b.id"
-        " WHERE orderline_id=:orderlineid",
+        text(
+            "SELECT b.reference FROM allocations JOIN batches AS b ON batch_id = b.id"
+            " WHERE orderline_id=:orderlineid"
+        ),
         dict(orderlineid=orderlineid),
     )
     return batchref
@@ -54,7 +59,7 @@ def test_rolls_back_uncommitted_work_by_default(session_factory):
         insert_batch(uow.session, "batch1", "MEDIUM-PLINTH", 100, None)
 
     new_session = session_factory()
-    rows = list(new_session.execute('SELECT * FROM "batches"'))
+    rows = list(new_session.execute(text('SELECT * FROM "batches"')))
     assert rows == []
 
 
@@ -69,6 +74,6 @@ class MyException(Exception):
             raise MyException()
 
     new_session = session_factory()
-    rows = list(new_session.execute('SELECT * FROM "batches"'))
+    rows = list(new_session.execute(text('SELECT * FROM "batches"')))
     assert rows == []
 """