1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
|
# CodeQ: an online programming tutor.
# Copyright (C) 2015 UL FRI
#
# This program is free software: you can redistribute it and/or modify it under
# the terms of the GNU Affero General Public License as published by the Free
# Software Foundation, either version 3 of the License, or (at your option) any
# later version.
#
# This program is distributed in the hope that it will be useful, but WITHOUT
# ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or FITNESS
# FOR A PARTICULAR PURPOSE. See the GNU Affero General Public License for more
# details.
#
# You should have received a copy of the GNU Affero General Public License
# along with this program. If not, see <http://www.gnu.org/licenses/>.
import os
import threading
import psycopg2
__all__ = ['get_connection', 'return_connection', 'setup', 'models']
_module_access_lock = threading.Lock()
_host = None # the database hostname/IP
_port = None # the database port number
_database = None # the name of the database
_username = None # the username to access the database
_password = None # the password to access the database
# database parameters setup
def _get_port():
try:
return int(os.environ.get('CODEQ_DB_PORT'))
except:
return 5432
def setup(
host=os.environ.get('CODEQ_DB_HOST') or 'localhost',
port=_get_port(),
database=os.environ.get('CODEQ_DB_DATABASE') or 'codeq',
username=os.environ.get('CODEQ_DB_USER') or 'codeq',
password=os.environ.get('CODEQ_DB_PASS') or 'c0d3q'
):
"""Sets the database location and authentication parameters."""
global _host, _port, _database, _username, _password
_host = host
_port = port
_database = database
_username = username
_password = password
# connection pooling
_connection_pool = []
def get_connection():
"""Retrieves a database connection from the connection pool."""
with _module_access_lock:
if _host is None:
setup() # lazy init
if len(_connection_pool) > 0:
return _connection_pool.pop()
return psycopg2.connect(host=_host, port=_port, database=_database, user=_username, password=_password)
def return_connection(connection):
"""Returns the given database connection to the pool."""
try:
connection.rollback() # sanity check
except:
return
with _module_access_lock:
_connection_pool.append(connection)
if __name__ == '__main__':
setup()
|
