The personal URL shortener
You can not select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
 
 

76 lines
2.5 KiB

import typing
import logging
import requests
from cloudant.client import CouchDB
from shortenit.exceptions import DBConnectionFailed
class DB:
"""
Database object class
"""
def __init__(self, config: dict) -> typing.NoReturn:
"""
Initialize the Database object.
:param config: The Database configuration.
"""
self.logger = logging.getLogger(self.__class__.__name__)
self.username = config['username']
self.password = config['password']
self.url = config['url']
self.client = None
self.session = None
def initialize_shortenit(self) -> typing.NoReturn:
"""
Method to initialize the database for shortenit.
This will check if all the needed tables already exist in the database.
Otherwise, it will create the database tables.
"""
try:
self.data_db = self.client['data']
except KeyError:
self.logger.warn(
"The 'data' database was not found, creating...")
self.data_db = self.client.create_database('data')
if self.data_db.exists():
self.logger.info(
"The 'data' database was successfully created.")
try:
self.pointers_db = self.client['pointers']
except KeyError:
self.logger.warn(
"The 'pointers' database was not found, creating...")
self.pointers_db = self.client.create_database('pointers')
if self.pointers_db.exists():
self.logger.info(
"The 'pointers' database was successfully created.")
def __enter__(self) -> CouchDB:
"""
Method used when entering the database context.
:returns: The CouchDB object.
"""
try:
self.client = CouchDB(self.username, self.password,
url=self.url, connect=True)
except requests.exceptions.ConnectionError as e:
self.logger.fatal("Failed to connect to database, is it on?")
self.logger.fatal("%s", e)
raise DBConnectionFailed
except requests.exceptions.HTTPError as e:
self.logger.fatal("Failed to authenticate to database.")
self.logger.fatal("%s", e)
raise DBConnectionFailed
self.session = self.client.session()
return self
def __exit__(self, *args) -> typing.NoReturn:
"""
Method used when exiting the database context.
"""
self.client.disconnect()