Make sqlalchemy session persisted

plugins/sqlalchemy/superduper_sqlalchemy/metadata.py

@@ -1,3 +1,4 @@
+import time
 import threading
 import typing as t
 from contextlib import contextmanager
@@ -84,6 +85,11 @@ def __init__(
         self._init_tables()
 
         self._lock = threading.Lock()
+        self._connect()
+
+    def _connect(self):
+        sm = sessionmaker(bind=self.conn)
+        self.session = sm()
 
     def reconnect(self):
         """Reconnect to sqlalchmey metadatastore."""
@@ -237,18 +243,28 @@ def drop(self, force: bool = False):
             logging.warn(f'Error dropping artifact table {e}')
 
     @contextmanager
-    def session_context(self):
+    def session_context(self, auto_close=False, max_retries=1):
         """Provide a transactional scope around a series of operations."""
-        sm = sessionmaker(bind=self.conn)
-        session = sm()
-        try:
-            yield session
-            session.commit()
-        except Exception:
-            session.rollback()
-            raise
-        finally:
-            session.close()
+        for attempt in range(max_retries):
+            try:
+                yield self.session
+                self.session.commit()
+                break
+            except Exception:
+                logging.debug('retrying session.')
+                if attempt < max_retries - 1:
+                    time.sleep(0.1)
+
+                    if self.session:
+                        self.session.close()
+                    self._connect()
+                else:
+                    self.session.rollback()
+                    raise
+            finally:
+                if auto_close:
+                    self.session.close()
+
 
     # --------------- COMPONENTS -----------------