summaryrefslogtreecommitdiff
path: root/db/__init__.py
diff options
context:
space:
mode:
Diffstat (limited to 'db/__init__.py')
-rw-r--r--db/__init__.py60
1 files changed, 60 insertions, 0 deletions
diff --git a/db/__init__.py b/db/__init__.py
new file mode 100644
index 0000000..e02a1fe
--- /dev/null
+++ b/db/__init__.py
@@ -0,0 +1,60 @@
+# coding=utf-8
+
+import os
+
+import psycopg2
+
+__all__ = ['get_connection', 'return_connection', 'setup', 'models']
+
+_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."""
+ 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."""
+ _connection_pool.append(connection)
+
+
+if __name__ == '__main__':
+ setup()