zach commited on
Commit
cadb3e8
·
1 Parent(s): e1385f3

Fix init_db.py script

Browse files
Files changed (1) hide show
  1. src/scripts/init_db.py +25 -37
src/scripts/init_db.py CHANGED
@@ -1,56 +1,44 @@
1
  """
2
- init_db.py
3
 
4
- This script initializes the database by creating all tables defined in the ORM models.
5
- It uses async SQLAlchemy operations to create tables in the PostgreSQL database.
6
- Run this script once to set up your database schema.
7
  """
8
 
9
  # Standard Library Imports
10
  import asyncio
 
 
 
 
11
 
12
  # Local Application Imports
13
  from src.config import Config, logger
14
- from src.database.database import engine
15
  from src.database.models import Base
16
 
17
 
18
- async def init_db_async():
19
- """
20
- Asynchronously create all database tables defined in the ORM models.
21
-
22
- This function connects to the database using the configured async engine
23
- and creates all tables that are mapped to SQLAlchemy models derived from
24
- the Base class. It uses SQLAlchemy's create_all method with the async
25
- engine context.
26
-
27
- Returns:
28
- None
29
- """
30
- async with engine.begin() as conn:
31
- # In SQLAlchemy 2.0 with async, we use the connection directly
32
- await conn.run_sync(Base.metadata.create_all)
33
-
34
- logger.info("Database tables created successfully using async SQLAlchemy.")
35
-
36
 
37
- def main():
38
- """
39
- Main entry point for the database initialization script.
40
 
41
- This function creates the configuration, ensures the async engine is
42
- initialized, and runs the async initialization function within an
43
- event loop.
44
 
45
- Returns:
46
- None
47
- """
48
- # Make sure config is loaded first to initialize the engine
49
- Config.get()
50
 
51
- # Run the async initialization function
52
- asyncio.run(init_db_async())
 
 
 
 
 
53
 
54
 
55
  if __name__ == "__main__":
56
- main()
 
 
1
  """
2
+ init_db_direct.py
3
 
4
+ A simplified script to initialize the database by creating all tables.
 
 
5
  """
6
 
7
  # Standard Library Imports
8
  import asyncio
9
+ import sys
10
+
11
+ # Third-Party Library Imports
12
+ from sqlalchemy.ext.asyncio import create_async_engine
13
 
14
  # Local Application Imports
15
  from src.config import Config, logger
 
16
  from src.database.models import Base
17
 
18
 
19
+ async def init_tables():
20
+ config = Config.get()
21
+ database_url = config.database_url
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
22
 
23
+ if not database_url:
24
+ logger.error("DATABASE_URL is not set in environment variables")
25
+ return False
26
 
27
+ engine = create_async_engine(database_url, echo=config.debug)
 
 
28
 
29
+ try:
30
+ async with engine.begin() as conn:
31
+ await conn.run_sync(Base.metadata.create_all)
 
 
32
 
33
+ logger.info("Database tables created successfully!")
34
+ return True
35
+ except Exception as e:
36
+ logger.error(f"Error creating tables: {e}")
37
+ return False
38
+ finally:
39
+ await engine.dispose()
40
 
41
 
42
  if __name__ == "__main__":
43
+ success = asyncio.run(init_tables())
44
+ sys.exit(0 if success else 1)