Spaces:
Running
Running
fix: really use connection pooling
Browse files- src/repository/common.py +25 -6
src/repository/common.py
CHANGED
|
@@ -8,19 +8,38 @@ import os
|
|
| 8 |
|
| 9 |
load_dotenv()
|
| 10 |
|
|
|
|
|
|
|
|
|
|
| 11 |
def get_engine():
|
| 12 |
-
|
| 13 |
-
|
| 14 |
-
|
| 15 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 16 |
|
| 17 |
@contextmanager
|
| 18 |
def get_session() -> Generator[Session, None, None]:
|
| 19 |
"""
|
| 20 |
Get a connection to the Postgres database
|
| 21 |
"""
|
| 22 |
-
|
| 23 |
-
SessionLocal = sessionmaker(bind=engine)
|
| 24 |
db = SessionLocal()
|
| 25 |
try:
|
| 26 |
yield db
|
|
|
|
| 8 |
|
| 9 |
load_dotenv()
|
| 10 |
|
| 11 |
+
_engine = None
|
| 12 |
+
_SessionLocal = None
|
| 13 |
+
|
| 14 |
def get_engine():
|
| 15 |
+
global _engine
|
| 16 |
+
|
| 17 |
+
if _engine is None:
|
| 18 |
+
# Load the database URL from environment variables
|
| 19 |
+
database_url = os.getenv("DATABASE_URL")
|
| 20 |
+
if not database_url:
|
| 21 |
+
raise ValueError("DATABASE_URL is not set in environment variables")
|
| 22 |
+
|
| 23 |
+
# Create the SQLAlchemy engine
|
| 24 |
+
_engine = create_engine(database_url, pool_pre_ping=True)
|
| 25 |
+
|
| 26 |
+
return _engine
|
| 27 |
+
|
| 28 |
+
def get_session_local() -> sessionmaker:
|
| 29 |
+
global _SessionLocal
|
| 30 |
+
|
| 31 |
+
if _SessionLocal is None:
|
| 32 |
+
# Create a new session local
|
| 33 |
+
_SessionLocal = sessionmaker(autocommit=False, autoflush=False, bind=get_engine())
|
| 34 |
+
|
| 35 |
+
return _SessionLocal
|
| 36 |
|
| 37 |
@contextmanager
|
| 38 |
def get_session() -> Generator[Session, None, None]:
|
| 39 |
"""
|
| 40 |
Get a connection to the Postgres database
|
| 41 |
"""
|
| 42 |
+
SessionLocal = get_session_local()
|
|
|
|
| 43 |
db = SessionLocal()
|
| 44 |
try:
|
| 45 |
yield db
|