Annotate zerver/lib/sqlalchemy_utils.py.

This commit is contained in:
Tim Abbott
2016-09-10 12:08:37 -07:00
parent 2e6aad669c
commit 799c78fbcf

View File

@@ -1,4 +1,5 @@
from django.db import connection from django.db import connection
from zerver.lib.db import TimeTrackingConnection
import sqlalchemy import sqlalchemy
@@ -6,12 +7,15 @@ import sqlalchemy
# existing Django database connections. # existing Django database connections.
class NonClosingPool(sqlalchemy.pool.NullPool): class NonClosingPool(sqlalchemy.pool.NullPool):
def status(self): def status(self):
# type: () -> str
return "NonClosingPool" return "NonClosingPool"
def _do_return_conn(self, conn): def _do_return_conn(self, conn):
# type: (sqlalchemy.engine.base.Connection) -> None
pass pass
def recreate(self): def recreate(self):
# type: () -> NonClosingPool
return self.__class__(creator=self._creator, # type: ignore # __class__ return self.__class__(creator=self._creator, # type: ignore # __class__
recycle=self._recycle, recycle=self._recycle,
use_threadlocal=self._use_threadlocal, use_threadlocal=self._use_threadlocal,
@@ -22,9 +26,11 @@ class NonClosingPool(sqlalchemy.pool.NullPool):
sqlalchemy_engine = None sqlalchemy_engine = None
def get_sqlalchemy_connection(): def get_sqlalchemy_connection():
# type: () -> sqlalchemy.engine.base.Connection
global sqlalchemy_engine global sqlalchemy_engine
if sqlalchemy_engine is None: if sqlalchemy_engine is None:
def get_dj_conn(): def get_dj_conn():
# type: () -> TimeTrackingConnection
connection.ensure_connection() connection.ensure_connection()
return connection.connection return connection.connection
sqlalchemy_engine = sqlalchemy.create_engine('postgresql://', sqlalchemy_engine = sqlalchemy.create_engine('postgresql://',