azul/core/conn.py

from typing import Any, Iterator
from contextlib import contextmanager
import sqlalchemy as sa
from sqlalchemy.orm import sessionmaker

class Conn:
	__slots__ = ('engine', 'session_factory', '_session', '_depth')
	
	engine: Any
	session_factory: Any
	_session: Any
	_depth: int
	
	def __init__(self, conn_str: str) -> None:
		self.engine = sa.create_engine(conn_str)
		self.session_factory = sessionmaker(bind = self.engine)
		self._session = None
		self._depth = 0
	
	@contextmanager
	def session(self) -> Iterator[Any]:
		if self._depth > 0:
			yield self._session
			return
		sess = self.session_factory()
		self._session = sess
		self._depth += 1
		try:
			yield sess
			sess.commit()
		except:
			sess.rollback()
			raise
		finally:
			sess.close()
			self._session = None
			self._depth -= 1