otf-robot-test-head/app/database.py

   1 """ Copyright (c) 2019 AT&T Intellectual Property.                             #\r
   2 #                                                                              #\r
   3 #   Licensed under the Apache License, Version 2.0 (the "License");            #\r
   4 #   you may not use this file except in compliance with the License.           #\r
   5 #   You may obtain a copy of the License at                                    #\r
   6 #                                                                              #\r
   7 #       http://www.apache.org/licenses/LICENSE-2.0                             #\r
   8 #                                                                              #\r
   9 #   Unless required by applicable law or agreed to in writing, software        #\r
  10 #   distributed under the License is distributed on an "AS IS" BASIS,          #\r
  11 #   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.   #\r
  12 #   See the License for the specific language governing permissions and        #\r
  13 #   limitations under the License.                                             #\r
  14 #############################################################################"""\r
  15 \r
  16 \r
  17 import os\r
  18 from urllib import quote_plus\r
  19 \r
  20 from pymongo import MongoClient\r
  21 from pymongo.errors import ConnectionFailure\r
  22 \r
  23 \r
  24 class DatabaseConfiguration:\r
  25     def __init__(self):\r
  26         # read environment variables containing information for the MongoDB replica set.\r
  27         MONGO_HOST = os.environ['OTF_MONGO_HOSTS']\r
  28         MONGO_USERNAME = os.environ['OTF_MONGO_USERNAME']\r
  29         MONGO_PASSWORD = os.environ['OTF_MONGO_PASSWORD']\r
  30         MONGO_REPLICA_SET = os.environ['OTF_MONGO_REPLICASET']\r
  31         MONGO_DATABASE = os.environ['OTF_MONGO_DATABASE']\r
  32 \r
  33         # form the connection string for connection to a MongoDB replica set.\r
  34         uri = "mongodb://%s:%s@%s?replicaSet=%s" % (\r
  35             quote_plus(MONGO_USERNAME),\r
  36             quote_plus(MONGO_PASSWORD),\r
  37             MONGO_HOST + MONGO_DATABASE,\r
  38             MONGO_REPLICA_SET\r
  39         )\r
  40 \r
  41         client = MongoClient(uri)\r
  42 \r
  43         try:\r
  44             # The ismaster command is cheap and does not require auth.\r
  45             client.admin.command('ismaster')\r
  46             print("Established connection to MongoDB.")\r
  47             self.database = client[MONGO_DATABASE]\r
  48         except ConnectionFailure:\r
  49             print("Failed to initialize connection to MongoDB.")\r
  50 \r
  51     def set_database(self, database):\r
  52         self.database = database\r
  53 \r
  54     def get_database(self):\r
  55         return self.database\r