adding vth code as well as aaf and cert config
[it/otf.git] / otf-robot-test-head / app / database.py
diff --git a/otf-robot-test-head/app/database.py b/otf-robot-test-head/app/database.py
new file mode 100644 (file)
index 0000000..ea1b433
--- /dev/null
@@ -0,0 +1,55 @@
+""" Copyright (c) 2019 AT&T Intellectual Property.                             #\r
+#                                                                              #\r
+#   Licensed under the Apache License, Version 2.0 (the "License");            #\r
+#   you may not use this file except in compliance with the License.           #\r
+#   You may obtain a copy of the License at                                    #\r
+#                                                                              #\r
+#       http://www.apache.org/licenses/LICENSE-2.0                             #\r
+#                                                                              #\r
+#   Unless required by applicable law or agreed to in writing, software        #\r
+#   distributed under the License is distributed on an "AS IS" BASIS,          #\r
+#   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.   #\r
+#   See the License for the specific language governing permissions and        #\r
+#   limitations under the License.                                             #\r
+#############################################################################"""\r
+\r
+\r
+import os\r
+from urllib import quote_plus\r
+\r
+from pymongo import MongoClient\r
+from pymongo.errors import ConnectionFailure\r
+\r
+\r
+class DatabaseConfiguration:\r
+    def __init__(self):\r
+        # read environment variables containing information for the MongoDB replica set.\r
+        MONGO_HOST = os.environ['OTF_MONGO_HOSTS']\r
+        MONGO_USERNAME = os.environ['OTF_MONGO_USERNAME']\r
+        MONGO_PASSWORD = os.environ['OTF_MONGO_PASSWORD']\r
+        MONGO_REPLICA_SET = os.environ['OTF_MONGO_REPLICASET']\r
+        MONGO_DATABASE = os.environ['OTF_MONGO_DATABASE']\r
+\r
+        # form the connection string for connection to a MongoDB replica set.\r
+        uri = "mongodb://%s:%s@%s?replicaSet=%s" % (\r
+            quote_plus(MONGO_USERNAME),\r
+            quote_plus(MONGO_PASSWORD),\r
+            MONGO_HOST + MONGO_DATABASE,\r
+            MONGO_REPLICA_SET\r
+        )\r
+\r
+        client = MongoClient(uri)\r
+\r
+        try:\r
+            # The ismaster command is cheap and does not require auth.\r
+            client.admin.command('ismaster')\r
+            print("Established connection to MongoDB.")\r
+            self.database = client[MONGO_DATABASE]\r
+        except ConnectionFailure:\r
+            print("Failed to initialize connection to MongoDB.")\r
+\r
+    def set_database(self, database):\r
+        self.database = database\r
+\r
+    def get_database(self):\r
+        return self.database\r