Merge pull request #23 from bit-bots/feature/database_reader

feature/database_reader

Author: texhnolyze

.gitignore

@@ -2,6 +2,7 @@ env/
 logs/
 *.csv
 *.sqlite
+*.sqlite3
 
 # Created by .ignore support plugin (hsz.mobi)
 ### JetBrains template
@@ -205,4 +206,4 @@ ENV/
 .ruff_cache
 
 # Torch models
-*.pth
+*.pth

ddlitlab2024/__init__.py

@@ -38,3 +38,5 @@
 LOGGING_PATH: str = _logging_path
 
 SESSION_ID: UUID = uuid4()
+
+DB_PATH: str = os.path.join(os.path.dirname(__file__), "dataset", "db.sqlite3")

ddlitlab2024/dataset/reader.py

@@ -0,0 +1,8 @@
+import sqlite3
+
+from ddlitlab2024 import DB_PATH
+
+
+def get_connection(path: str = DB_PATH) -> sqlite3.Connection:
+    conn = sqlite3.connect(path)
+    return conn

ddlitlab2024/dataset/schema.py

@@ -1,3 +1,4 @@
+import argparse
 from datetime import datetime
 from enum import Enum
 from typing import List, Optional
@@ -6,6 +7,7 @@
 from sqlalchemy.orm import Mapped, declarative_base, mapped_column, relationship, sessionmaker
 from sqlalchemy.types import LargeBinary
 
+from ddlitlab2024 import DB_PATH
 from ddlitlab2024.dataset import logger
 
 Base = declarative_base()
@@ -229,9 +231,16 @@ class GameState(Base):
     __table_args__ = (CheckConstraint(state.in_(RobotState.values())),)
 
 
+def parse_args():
+    parser = argparse.ArgumentParser(description="Create the database schema")
+    parser.add_argument("--db-path", type=str, default=DB_PATH, help="Path to the database file")
+    return parser.parse_args()
+
+
 def main():
     logger.info("Creating database schema")
-    engine = create_engine("sqlite:///data.sqlite")
+    args = parse_args()
+    engine = create_engine(f"sqlite:///{args.db_path}")
     Base.metadata.create_all(engine)
     sessionmaker(bind=engine)()
     logger.info("Database schema created")