amaye15 commited on
Commit
9e56290
·
1 Parent(s): dd00186
Files changed (1) hide show
  1. app/database.py +79 -52
app/database.py CHANGED
@@ -84,29 +84,54 @@
84
  # logger.exception(f"Error closing async database connection: {e}")
85
 
86
 
 
87
  # app/database.py
88
  import os
89
- # --- Keep only Sync SQLAlchemy ---
90
- from sqlalchemy import create_engine, MetaData, Table, Column, Integer, String, text, exc as sqlalchemy_exc
91
  from dotenv import load_dotenv
 
 
 
 
92
  import logging
 
93
 
 
94
  load_dotenv()
95
  logger = logging.getLogger(__name__)
96
 
97
- # Use /tmp for ephemeral storage in HF Space
 
98
  DEFAULT_DB_PATH = "/tmp/app.db"
99
- # Construct sync URL directly
100
- DATABASE_URL = os.getenv("DATABASE_URL_SYNC", f"sqlite:///{DEFAULT_DB_PATH}")
101
-
102
- logger.info(f"Using DB URL for sync operations: {DATABASE_URL}")
103
-
104
- # SQLite specific args for sync engine
105
- connect_args = {"check_same_thread": False} if DATABASE_URL.startswith("sqlite") else {}
106
- engine = create_engine(DATABASE_URL, connect_args=connect_args, echo=False) # echo=True for debugging SQL
107
-
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
108
  metadata = MetaData()
109
- users_table = Table( # Renamed slightly to avoid confusion with Pydantic model name
110
  "users",
111
  metadata,
112
  Column("id", Integer, primary_key=True),
@@ -114,45 +139,47 @@ users_table = Table( # Renamed slightly to avoid confusion with Pydantic model n
114
  Column("hashed_password", String, nullable=False),
115
  )
116
 
117
- def ensure_db_and_table_exist():
118
- """Synchronously ensures DB file directory and users table exist."""
119
- logger.info("Ensuring DB and table exist...")
 
 
120
  try:
121
- # Ensure directory exists
122
- db_file_path = DATABASE_URL.split("sqlite:///")[-1]
123
- db_dir = os.path.dirname(db_file_path)
124
- if db_dir:
125
- if not os.path.exists(db_dir):
126
- logger.info(f"Creating DB directory: {db_dir}")
127
- os.makedirs(db_dir, exist_ok=True)
128
- if not os.access(db_dir, os.W_OK):
129
- logger.error(f"CRITICAL: Directory {db_dir} not writable!")
130
- return # Cannot proceed
131
-
132
- # Check/Create table using the engine
133
- with engine.connect() as connection:
134
- try:
135
- connection.execute(text("SELECT 1 FROM users LIMIT 1"))
136
- logger.info("Users table already exists.")
137
- except sqlalchemy_exc.OperationalError as e:
138
- if "no such table" in str(e).lower():
139
- logger.warning("Users table not found, creating...")
140
- metadata.create_all(bind=connection) # Use connection
141
- connection.commit()
142
- logger.info("Users table created and committed.")
143
- # Verify
144
- try:
145
- connection.execute(text("SELECT 1 FROM users LIMIT 1"))
146
- logger.info("Users table verified post-creation.")
147
- except Exception as verify_err:
148
- logger.error(f"Verification failed after creating table: {verify_err}")
149
- else:
150
- logger.error(f"DB OperationalError checking table (not 'no such table'): {e}")
151
- raise # Re-raise unexpected errors
152
- except Exception as check_err:
153
- logger.error(f"Unexpected error checking table: {check_err}")
154
- raise # Re-raise unexpected errors
155
 
 
 
 
 
 
 
 
 
156
  except Exception as e:
157
- logger.exception(f"CRITICAL error during DB setup: {e}")
158
- # Potentially raise to halt app start?
 
84
  # logger.exception(f"Error closing async database connection: {e}")
85
 
86
 
87
+
88
  # app/database.py
89
  import os
90
+ from databases import Database
 
91
  from dotenv import load_dotenv
92
+ # --- Keep only these SQLAlchemy imports ---
93
+ # MetaData and Table are needed for defining the table structure
94
+ # which is used by crud.py and for DDL generation in main.py
95
+ from sqlalchemy import MetaData, Table, Column, Integer, String
96
  import logging
97
+ from urllib.parse import urlparse, urlunparse, parse_qs, urlencode
98
 
99
+ # Load environment variables from .env file (if it exists)
100
  load_dotenv()
101
  logger = logging.getLogger(__name__)
102
 
103
+ # --- Database URL Configuration ---
104
+ # Use /tmp directory for the SQLite file as it's generally writable in containers
105
  DEFAULT_DB_PATH = "/tmp/app.db"
106
+ # Get the URL from environment or use the default /tmp path
107
+ raw_db_url = os.getenv("DATABASE_URL", f"sqlite+aiosqlite:///{DEFAULT_DB_PATH}")
108
+
109
+ final_database_url = raw_db_url
110
+ # Ensure 'check_same_thread=False' is in the URL query string for SQLite async connection
111
+ if raw_db_url.startswith("sqlite+aiosqlite"):
112
+ parsed_url = urlparse(raw_db_url)
113
+ query_params = parse_qs(parsed_url.query)
114
+ if 'check_same_thread' not in query_params:
115
+ query_params['check_same_thread'] = ['False'] # Value needs to be a list for urlencode
116
+ new_query = urlencode(query_params, doseq=True)
117
+ # Rebuild the URL using _replace method of the named tuple
118
+ final_database_url = urlunparse(parsed_url._replace(query=new_query))
119
+ logger.info(f"Using final async DB URL: {final_database_url}")
120
+ else:
121
+ logger.info(f"Using non-SQLite async DB URL: {final_database_url}")
122
+
123
+
124
+ # --- Async Database Instance ---
125
+ # This 'database' object will be used by crud.py and main.py lifespan
126
+ database = Database(final_database_url)
127
+
128
+
129
+ # --- Metadata and Table Definition ---
130
+ # These definitions are needed by:
131
+ # 1. crud.py to construct queries (e.g., users.select())
132
+ # 2. main.py (lifespan) to generate the CREATE TABLE statement
133
  metadata = MetaData()
134
+ users = Table(
135
  "users",
136
  metadata,
137
  Column("id", Integer, primary_key=True),
 
139
  Column("hashed_password", String, nullable=False),
140
  )
141
 
142
+
143
+ # --- Async connect/disconnect functions ---
144
+ # Called by the FastAPI lifespan event handler in main.py
145
+ async def connect_db():
146
+ """Connects to the database defined by 'final_database_url'."""
147
  try:
148
+ # Optional: Check/create directory if using file-based DB like SQLite
149
+ if final_database_url.startswith("sqlite"):
150
+ db_file_path = final_database_url.split("sqlite:///")[-1].split("?")[0]
151
+ db_dir = os.path.dirname(db_file_path)
152
+ if db_dir:
153
+ if not os.path.exists(db_dir):
154
+ logger.info(f"DB directory '{db_dir}' missing, creating.")
155
+ try:
156
+ os.makedirs(db_dir, exist_ok=True)
157
+ except Exception as mkdir_err:
158
+ logger.error(f"Failed to create DB directory '{db_dir}': {mkdir_err}")
159
+ # Check writability (best effort)
160
+ if os.path.exists(db_dir) and not os.access(db_dir, os.W_OK):
161
+ logger.error(f"CRITICAL: DB directory '{db_dir}' exists but is not writable!")
162
+ elif not os.path.exists(db_dir):
163
+ logger.error(f"CRITICAL: DB directory '{db_dir}' does not exist and could not be created!")
164
+
165
+ # Connect using the 'databases' library instance
166
+ if not database.is_connected:
167
+ await database.connect()
168
+ logger.info(f"Database connection established: {final_database_url}")
169
+ else:
170
+ logger.info("Database connection already established.")
171
+ # Note: Table creation happens in main.py lifespan after connection
172
+ except Exception as e:
173
+ logger.exception(f"FATAL: Failed to establish async database connection: {e}")
174
+ raise # Stop application startup if DB connection fails
 
 
 
 
 
 
 
175
 
176
+ async def disconnect_db():
177
+ """Disconnects from the database if connected."""
178
+ try:
179
+ if database.is_connected:
180
+ await database.disconnect()
181
+ logger.info("Database connection closed.")
182
+ else:
183
+ logger.info("Database already disconnected.")
184
  except Exception as e:
185
+ logger.exception(f"Error closing database connection: {e}")