-
Notifications
You must be signed in to change notification settings - Fork 0
/
db.py
54 lines (43 loc) · 1.52 KB
/
db.py
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
# Copyright (c) 2015 Hong Quach
# This source file is licensed under the "MIT License." Please see the LICENSE
# in this distribution for license terms.
"""This module handles database connection tasks"""
import os
from sqlalchemy import create_engine
from sqlalchemy.engine.url import make_url
import constant
import exc
__author__ = 'htquach'
def get_db_url_from_env():
"""Get the DB URL from the environment variable. Return the DB URL o
:raise exc.ArPrDataDBURLEnvironmentVariableEmptyException:
:rtype : str
"""
db_url = os.getenv(constant.DB_URL_ENV_NAME)
if db_url:
check_db_url(db_url)
return db_url
else:
raise exc.ArPrDataDBURLEnvironmentVariableEmptyException(
"The DB URL environment variable '%s' is empty."
% constant.DB_URL_ENV_NAME)
def check_db_url(db_url):
"""Check for the correct format of the DB URL. Return none if check
passed. Raise exception otherwise.
:param db_url: The DB URL to be checked
:raise exc.ArPrDataBadDBURLFormatException:
:rtype : None
"""
try:
make_url(db_url)
except Exception:
raise exc.ArPrDataBadDBURLFormatException(
"Incorrect DB URL Format. Check http://docs.sqlalchemy.org"
"/en/latest/core/engines.html?highlight=url#database-urls")
def get_engine(db_url):
"""Connect to the database and engine
:param db_url: Database URL to connect to
:return: Database engine
"""
check_db_url(db_url)
return create_engine(db_url)