Address reviewer comments

LSSTDESC · Dec 1, 2023 · ebbfa5c · ebbfa5c
1 parent 506a933
commit ebbfa5c
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 5 deletions.
diff --git a/scripts/create_registry_db.py b/scripts/create_registry_db.py
@@ -4,7 +4,7 @@
 from datetime import datetime
 from sqlalchemy import Column, Integer, String, DateTime, Boolean, Index, Float
 from sqlalchemy import ForeignKey, UniqueConstraint, text
-from sqlalchemy.orm import relationship, declarative_base
+from sqlalchemy.orm import relationship, DeclarativeBase
 from dataregistry.db_basic import DbConnection, SCHEMA_VERSION
 from dataregistry.db_basic import add_table_row, _insert_provenance
 
@@ -20,8 +20,8 @@
     - "provenance"      : Contains information about the database/schema 
 """
 
-Base = declarative_base()
-
+class Base(DeclarativeBase):
+    pass
 
 def _get_ForeignKey_str(schema, table, row):
     if schema is None:

diff --git a/tests/end_to_end_tests/create_test_entries.py b/tests/end_to_end_tests/create_test_entries.py
@@ -35,8 +35,8 @@
 # Establish connection to database
 datareg = DataRegistry(root_dir=_TEST_ROOT_DIR)
 
-# Establish connection to production database (if not sqllite)
-if datareg.db_connection.dialect != "sqllite":
+# Establish connection to production database (if not sqlite)
+if datareg.db_connection.dialect != "sqlite":
     datareg_prod = DataRegistry(root_dir=_TEST_ROOT_DIR_PRODUCTION, schema="production")