Upload database.py

database.py

@@ -1,218 +1,154 @@
 """
-Database access layer for TalkToDoc.
-Handles both SQLite (local development) and PostgreSQL (production) through
-a single shared query interface. The only difference between the two is how
-get_connection() builds the connection object; every function above that is
-identical regardless of which database is running underneath.
-
-Connection management: every public function in this module uses the
-_connection() context manager (not get_connection() directly). That
-context manager commits on success, rolls back on exception, and always
-closes the connection before returning control. The bare
-`with get_connection() as conn:` pattern that psycopg2 and sqlite3 both
-support commits/rolls back correctly but never closes the connection,
-which leaks file handles against SQLite and exhausts the Postgres
-connection pool in production over time.
+Database layer for TalkToDoc.
+Local development and testing use SQLite (no setup needed). When deployed,
+DATABASE_URL is set by Render and the app uses Postgres instead, since
+Render's free tier wipes local files like a SQLite database on every
+restart.
+
+Queries are written once using SQLite-style '?' placeholders and adapted
+automatically for Postgres, so there's a single query per function rather
+than two versions of everything.
 """
 
 import os
-import contextlib
+from pathlib import Path
 
-from dotenv import load_dotenv
+import sqlite3
 
-load_dotenv("env")
+DB_PATH = Path(__file__).parent / "talktodoc.db"
+SCHEMA_PATH = Path(__file__).parent / "schema.sql"
+SCHEMA_PATH_POSTGRES = Path(__file__).parent / "schema_postgres.sql"
 
 DATABASE_URL = os.environ.get("DATABASE_URL")
 
+if DATABASE_URL:
+    import psycopg2
+    import psycopg2.extras
+
 
 def get_connection():
-    """
-    Returns a raw database connection. Callers should use _connection()
-    instead, which wraps this and guarantees the connection is closed.
-    """
     if DATABASE_URL:
-        import psycopg2
-        import psycopg2.extras
-        conn = psycopg2.connect(DATABASE_URL, cursor_factory=psycopg2.extras.RealDictCursor)
-        return conn
-    else:
-        import sqlite3
-        conn = sqlite3.connect(
-            os.path.join(os.path.dirname(os.path.abspath(__file__)), "talktodoc.db")
-        )
-        conn.row_factory = sqlite3.Row
-        return conn
-
-
-@contextlib.contextmanager
-def _connection():
-    """
-    Context manager that opens a connection, yields it, commits on clean
-    exit, rolls back on exception, and always closes. Use this everywhere
-    instead of `with get_connection() as conn:`.
-    """
-    conn = get_connection()
-    try:
-        yield conn
-        conn.commit()
-    except Exception:
-        conn.rollback()
-        raise
-    finally:
-        conn.close()
-
-
-def _run(connection, sql, params=()):
+        connection = psycopg2.connect(DATABASE_URL, cursor_factory=psycopg2.extras.RealDictCursor)
+        # psycopg2 does not auto-commit by default. Without this, every
+        # INSERT and UPDATE is silently rolled back when the connection
+        # closes, so nothing is ever actually written to Postgres.
+        # SQLite's context manager commits on clean exit, so this brings
+        # the Postgres path in line with SQLite's behaviour.
+        connection.autocommit = True
+        return connection
+    connection = sqlite3.connect(DB_PATH)
+    connection.row_factory = sqlite3.Row
+    connection.execute("PRAGMA foreign_keys = ON")
+    return connection
+
+
+def _run(connection, query, params=()):
+    if DATABASE_URL:
+        query = query.replace("?", "%s")
     cursor = connection.cursor()
-    cursor.execute(sql, params)
+    cursor.execute(query, params)
     return cursor
 
 
-def _placeholder():
-    """
-    SQLite uses ? for bind parameters; psycopg2 uses %s.
-    Returns the right one for whichever database is active.
-    """
-    return "%s" if DATABASE_URL else "?"
+def _insert_and_get_id(connection, query, params):
+    if DATABASE_URL:
+        cursor = _run(connection, query + " RETURNING id", params)
+        return cursor.fetchone()["id"]
+    cursor = _run(connection, query, params)
+    return cursor.lastrowid
 
 
 def init_db():
-    schema_file = "schema_postgres.sql" if DATABASE_URL else "schema.sql"
-    schema_path = os.path.join(os.path.dirname(os.path.abspath(__file__)), schema_file)
-    with open(schema_path) as f:
-        schema = f.read()
-    with _connection() as conn:
+    schema_path = SCHEMA_PATH_POSTGRES if DATABASE_URL else SCHEMA_PATH
+    schema_sql = schema_path.read_text()
+    with get_connection() as connection:
         if DATABASE_URL:
-            for statement in schema.split(";"):
-                stmt = statement.strip()
-                if stmt:
-                    _run(conn, stmt)
+            connection.cursor().execute(schema_sql)
         else:
-            conn.executescript(schema)
+            connection.executescript(schema_sql)
 
 
 def add_user(name, preferred_language, role):
-    p = _placeholder()
-    with _connection() as conn:
-        if DATABASE_URL:
-            cursor = _run(
-                conn,
-                f"INSERT INTO app_user (name, preferred_language, role) VALUES ({p}, {p}, {p}) RETURNING id",
-                (name, preferred_language, role),
-            )
-            return cursor.fetchone()["id"]
-        else:
-            cursor = _run(
-                conn,
-                f"INSERT INTO app_user (name, preferred_language, role) VALUES ({p}, {p}, {p})",
-                (name, preferred_language, role),
-            )
-            return cursor.lastrowid
-
-
-def add_session(user_id, start_time):
-    p = _placeholder()
-    with _connection() as conn:
-        if DATABASE_URL:
-            cursor = _run(
-                conn,
-                f"INSERT INTO session (user_id, start_time) VALUES ({p}, {p}) RETURNING id",
-                (user_id, start_time),
-            )
-            return cursor.fetchone()["id"]
-        else:
-            cursor = _run(
-                conn,
-                f"INSERT INTO session (user_id, start_time) VALUES ({p}, {p})",
-                (user_id, start_time),
-            )
-            return cursor.lastrowid
-
-
-def end_session(session_id, end_time):
-    p = _placeholder()
-    with _connection() as conn:
-        _run(
-            conn,
-            f"UPDATE session SET end_time = {p} WHERE id = {p}",
-            (end_time, session_id),
+    with get_connection() as connection:
+        return _insert_and_get_id(
+            connection,
+            "INSERT INTO app_user (name, preferred_language, role) VALUES (?, ?, ?)",
+            (name, preferred_language, role),
         )
 
 
-def add_interaction(user_id, input_text, detected_language, translated_text,
-                    nlu_summary, timestamp):
-    p = _placeholder()
-
-    sql = (
-        f"INSERT INTO interaction "
-        f"(user_id, input_text, detected_language, translated_text, nlu_summary, timestamp) "
-        f"VALUES ({p}, {p}, {p}, {p}, {p}, {p})"
-    )
-
-    params = (
-        user_id,
-        input_text,
-        detected_language,
-        translated_text,
-        nlu_summary,
-        timestamp
-    )
-
-    with _connection() as conn:
-        if DATABASE_URL:
-            cursor = _run(conn, sql + " RETURNING id", params)
-            return cursor.fetchone()["id"]
-        else:
-            cursor = _run(conn, sql, params)
-            return cursor.lastrowid
-
-
-def get_interaction(interaction_id):
-    p = _placeholder()
-    with _connection() as conn:
-        cursor = _run(
-            conn,
-            f"SELECT * FROM interaction WHERE id = {p}",
-            (interaction_id,),
+def add_session(user_id, start_time):
+    with get_connection() as connection:
+        return _insert_and_get_id(
+            connection,
+            "INSERT INTO session (user_id, start_time) VALUES (?, ?)",
+            (user_id, start_time),
         )
-        row = cursor.fetchone()
-        return dict(row) if row else None
 
 
-def update_interaction_response(interaction_id, provider_response, translated_response):
-    p = _placeholder()
-    with _connection() as conn:
-        _run(
-            conn,
-            f"UPDATE interaction SET provider_response = {p}, translated_response = {p} WHERE id = {p}",
-            (provider_response, translated_response, interaction_id),
+def end_session(session_id, end_time):
+    with get_connection() as connection:
+        _run(connection, "UPDATE session SET end_time = ? WHERE id = ?", (end_time, session_id))
+
+
+def add_interaction(user_id, input_text, detected_language, translated_text, nlu_summary, timestamp):
+    with get_connection() as connection:
+        return _insert_and_get_id(
+            connection,
+            """INSERT INTO interaction
+               (user_id, input_text, detected_language, translated_text, nlu_summary, timestamp)
+               VALUES (?, ?, ?, ?, ?, ?)""",
+            (user_id, input_text, detected_language, translated_text, nlu_summary, timestamp),
         )
 
 
 def get_interactions_for_user(user_id):
-    p = _placeholder()
-    with _connection() as conn:
+    with get_connection() as connection:
         cursor = _run(
-            conn,
-            f"SELECT * FROM interaction WHERE user_id = {p} ORDER BY timestamp",
+            connection,
+            "SELECT * FROM interaction WHERE user_id = ? ORDER BY timestamp",
             (user_id,),
         )
         return [dict(row) for row in cursor.fetchall()]
 
 
 def get_pending_interactions():
-    with _connection() as conn:
+    with get_connection() as connection:
         cursor = _run(
-            conn,
+            connection,
             "SELECT * FROM interaction WHERE provider_response IS NULL ORDER BY timestamp",
         )
         return [dict(row) for row in cursor.fetchall()]
 
 
 def get_completed_interactions():
-    with _connection() as conn:
+    with get_connection() as connection:
         cursor = _run(
-            conn,
+            connection,
             "SELECT * FROM interaction WHERE provider_response IS NOT NULL ORDER BY timestamp DESC",
         )
         return [dict(row) for row in cursor.fetchall()]
+
+
+def get_interaction(interaction_id):
+    with get_connection() as connection:
+        cursor = _run(connection, "SELECT * FROM interaction WHERE id = ?", (interaction_id,))
+        row = cursor.fetchone()
+        return dict(row) if row else None
+
+
+def update_interaction_response(interaction_id, provider_response, translated_response):
+    with get_connection() as connection:
+        _run(
+            connection,
+            "UPDATE interaction SET provider_response = ?, translated_response = ? WHERE id = ?",
+            (provider_response, translated_response, interaction_id),
+        )
+
+
+if __name__ == "__main__":
+    init_db()
+    if DATABASE_URL:
+        print("Database initialized (Postgres)")
+    else:
+        print("Database created at:", DB_PATH)