Add DB version with alimbic and add log system

fuware/db/db_setup.py

@@ -1,8 +1,8 @@
 from collections.abc import Generator
+from contextlib import contextmanager
 from sqlalchemy.orm.session import Session
-from sqlalchemy import create_engine, event, Engine, text
+from sqlalchemy import create_engine, event, Engine
 from sqlalchemy.orm import scoped_session, sessionmaker
-
 from fuware.core.config import get_app_settings
 
 settings = get_app_settings()
@@ -10,13 +10,13 @@ settings = get_app_settings()
 def sql_global_init(db_url: str):
   connect_args = {"check_same_thread": False}
 
-  engine = create_engine(db_url, echo=True, connect_args=connect_args, pool_pre_ping=True, future=True)
+  engine = create_engine(db_url, echo=False, connect_args=connect_args, pool_pre_ping=True, future=True)
 
   SessionLocal = scoped_session(sessionmaker(autocommit=False, autoflush=False, bind=engine, future=True))
 
   return SessionLocal, engine
 
-SessionLocal, engine = sql_global_init(settings.DB_URL)  # type: ignore
+SessionLocal, engine = sql_global_init(settings.DB_URL)
 
 @event.listens_for(Engine, "connect")
 def set_sqlite_pragma(dbapi_connection, connection_record):
@@ -24,10 +24,21 @@ def set_sqlite_pragma(dbapi_connection, connection_record):
   cursor.execute("PRAGMA foreign_keys=ON")
   cursor.close()
 
-# with engine.connect() as connection:
-#   result = connection.execute(text('select "Hello"'))
+@contextmanager
+def session_context() -> Session: # type: ignore
+    """
+    session_context() provides a managed session to the database that is automatically
+    closed when the context is exited. This is the preferred method of accessing the
+    database.
 
-#   print(result.all())
+    Note: use `generate_session` when using the `Depends` function from FastAPI
+    """
+    global SessionLocal
+    sess = SessionLocal()
+    try:
+        yield sess
+    finally:
+        sess.close()
 
 def generate_session() -> Generator[Session, None, None]:
   db = SessionLocal()

fuware/db/init_db.py

@@ -1,9 +1,87 @@
-from db_setup import engine
-from fuware.db.seeder import initialize_table
-from models._model_base import Model
-from sqlalchemy import event
-from models.users import User
+import os
+from pathlib import Path
+from time import sleep
 
-event.listen(User.__table__, 'after_create', initialize_table)
+from sqlalchemy import engine, orm, text
 
-Model.metadata.create_all(bind=engine)
+from alembic import command, config, script
+from alembic.config import Config
+from alembic.runtime import migration
+from fuware.core import root_logger
+from fuware.core.config import get_app_settings
+from fuware.db.db_setup import session_context
+from fuware.repos.repository_users import RepositoryUsers
+from fuware.repos.seeder import default_users_init
+from fuware.db.models._model_base import Model
+# from fuware.db.models import User
+
+PROJECT_DIR = Path(__file__).parent.parent.parent
+
+logger = root_logger.get_logger()
+
+def init_db(db) -> None:
+  logger.info("Initializing user data...")
+  default_users_init(db)
+
+def db_is_at_head(alembic_cfg: config.Config) -> bool:
+  settings = get_app_settings()
+  url = settings.DB_URL
+
+  if not url:
+      raise ValueError("No database url found")
+
+  connectable = engine.create_engine(url)
+  directory = script.ScriptDirectory.from_config(alembic_cfg)
+  with connectable.begin() as connection:
+      context = migration.MigrationContext.configure(connection)
+      return set(context.get_current_heads()) == set(directory.get_heads())
+
+def connect(session: orm.Session) -> bool:
+  try:
+    session.execute(text("SELECT 1"))
+    return True
+  except Exception as e:
+    logger.error(f"Error connecting to database: {e}")
+    return False
+
+def main():
+  max_retry = 10
+  wait_seconds = 1
+
+  with session_context() as session:
+    while True:
+      if connect(session):
+        logger.info("Database connection established.")
+        break
+
+      logger.error(f"Database connection failed. Retrying in {wait_seconds} seconds...")
+      max_retry -= 1
+
+      sleep(wait_seconds)
+
+      if max_retry == 0:
+        raise ConnectionError("Database connection failed - exiting application.")
+
+    alembic_cfg_path = os.getenv("ALEMBIC_CONFIG_FILE", default=str(PROJECT_DIR / "alembic.ini"))
+    if not os.path.isfile(alembic_cfg_path):
+        raise Exception("Provided alembic config path doesn't exist")
+
+    alembic_cfg = Config(alembic_cfg_path)
+    if db_is_at_head(alembic_cfg):
+        logger.debug("Migration not needed.")
+    else:
+        logger.info("Migration needed. Performing migration...")
+        command.upgrade(alembic_cfg, "head")
+
+    if session.get_bind().name == "postgresql":  # needed for fuzzy search and fast GIN text indices
+        session.execute(text("CREATE EXTENSION IF NOT EXISTS pg_trgm;"))
+
+    users = RepositoryUsers()
+    if users.get_all():
+      logger.info("Database already seeded.")
+    else:
+      logger.info("Seeding database...")
+      init_db(session)
+
+if __name__ == "__main__":
+  main()

fuware/db/models/users/users.py

@@ -1,6 +1,6 @@
-import uuid
-from sqlalchemy import Boolean, Column, String
-from sqlalchemy.orm import Mapped, mapped_column
+from uuid import uuid4
+from sqlalchemy import Boolean, ForeignKey, String
+from sqlalchemy.orm import Mapped, mapped_column, relationship
 
 from sqlalchemy.dialects.postgresql import UUID
 
@@ -9,12 +9,26 @@ from .._model_base import SqlAlchemyBase
 class User(SqlAlchemyBase):
   __tablename__ = 'users'
 
-  id: Mapped[UUID] = mapped_column(UUID, primary_key=True, default=uuid.uuid4, index=True)
+  id: Mapped[UUID] = mapped_column(UUID, primary_key=True, default=uuid4, index=True)
   username: Mapped[str | None] = mapped_column(String, unique=True, index=True, nullable=False)
   password: Mapped[str | None] = mapped_column(String, index=True, nullable=False)
   name: Mapped[str | None] = mapped_column(String, index=True, nullable=True)
   is_admin: Mapped[bool | None] = mapped_column(Boolean, default=False)
   is_lock: Mapped[bool | None] = mapped_column(Boolean, default=False)
 
+  session_login = relationship("SessionLogin", back_populates="user", uselist=False)
+
   def __repr__(self):
     return f"{self.__class__.__name__}, name: {self.name}, username: {self.username}"
+
+class SessionLogin(SqlAlchemyBase):
+  __tablename__ = 'session_login'
+
+  id: Mapped[int] = mapped_column(primary_key=True, autoincrement=True)
+  session: Mapped[str] = mapped_column(UUID, default=uuid4, index=True, nullable=False)
+  user_id: Mapped[str] = mapped_column(ForeignKey("users.id"), unique=True, index=True, nullable=False)
+
+  user = relationship("User", back_populates="session_login")
+
+  def __repr__(self):
+    return f"{self.__class__.__name__}, session: {self.session}, user_id: {self.user_id}"