diff --git a/aw_datastore/storages/peewee.py b/aw_datastore/storages/peewee.py
index a4afe22..0c245af 100644
--- a/aw_datastore/storages/peewee.py
+++ b/aw_datastore/storages/peewee.py
@@ -12,6 +12,7 @@
 import iso8601
 from aw_core.dirs import get_data_dir
 from aw_core.models import Event
+from playhouse.migrate import SqliteMigrator, migrate
 from playhouse.sqlite_ext import SqliteExtDatabase
 
 import peewee
@@ -44,7 +45,6 @@
 
 
 def auto_migrate(path: str) -> None:
-    from playhouse.migrate import SqliteMigrator, migrate
 
     db = SqliteExtDatabase(path)
     migrator = SqliteMigrator(db)
@@ -154,13 +154,15 @@ def __init__(self, testing: bool = True, filepath: Optional[str] = None) -> None
         self.bucket_keys: Dict[str, int] = {}
         BucketModel.create_table(safe=True)
         EventModel.create_table(safe=True)
-        self.update_bucket_keys()
 
         # Migrate database if needed, requires closing the connection first
         self.db.close()
         auto_migrate(filepath)
         self.db.connect()
 
+        # Update bucket keys
+        self.update_bucket_keys()
+
     def update_bucket_keys(self) -> None:
         buckets = BucketModel.select()
         self.bucket_keys = {bucket.id: bucket.key for bucket in buckets}