[fix] Resolve compatibility issues with sqlalchemy < 2.0 versions (#…

…3096)
aimhubio · Jan 28, 2024 · cf46d7e · cf46d7e
1 parent 631061b
commit cf46d7e
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 1 deletion.
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -8,6 +8,9 @@
 - Add `min/max/first` values tracking for metrics (mihran113)
 - Fix bug in pytorch lightning raising lock timeout (inc0)
 
+### Fixes 
+- Fix compatibility with `sqlalchemy < 2.0` versions
+
 ## 3.17.6 
 
 - Switch to patched version of official `pynvml` (mihran113)

diff --git a/aim/storage/structured/sql_engine/entities.py b/aim/storage/structured/sql_engine/entities.py
@@ -1,6 +1,7 @@
 import pytz
 
 from typing import Collection, Union, List, Optional
+from sqlalchemy import __version__ as sa_version
 from sqlalchemy import delete
 from sqlalchemy.orm import joinedload
 from sqlalchemy.exc import IntegrityError
@@ -29,7 +30,7 @@
 
 
 def session_commit_or_flush(session):
-    if getattr(session, 'autocommit', True):
+    if getattr(session, 'autocommit', True) and sa_version >= '2.0.0':
         session.commit()
     else:
         session.flush()