chore: Updates for SQLAlchemy 2.0

cihai · Jun 4, 2023 · 3d16c48 · 3d16c48
1 parent 6ac762f
commit 3d16c48
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 7 deletions.
diff --git a/src/unihan_db/tables.py b/src/unihan_db/tables.py
@@ -23,10 +23,11 @@
 
 """
 from sqlalchemy import Boolean, Column, ForeignKey, Integer, String
-from sqlalchemy.ext.declarative import declarative_base
-from sqlalchemy.orm import relationship
+from sqlalchemy.orm import DeclarativeBase, relationship
 
-Base = declarative_base()
+
+class Base(DeclarativeBase):
+    pass
 
 
 class Unhn(Base):

diff --git a/tests/test_bootstrap.py b/tests/test_bootstrap.py
@@ -16,14 +16,12 @@ def test_import_object(session, engine):
 
 
 def test_import_unihan(zip_file, session, engine, unihan_options):
-    Base.metadata.bind = engine
-    Base.metadata.create_all()
+    Base.metadata.create_all(bind=engine)
     # bootstrap.bootstrap_unihan(Base.metadata, unihan_options)
 
 
 def test_import_unihan_raw(zip_file, session, engine, unihan_options):
-    Base.metadata.bind = engine
-    Base.metadata.create_all()
+    Base.metadata.create_all(bind=engine)
 
     bootstrap.bootstrap_unihan(session, unihan_options)