Merge pull request influxdata#276 from slomek/create-db-if-not-exist

aviau · aviau · commit 3c16b161cbbf · 2016-01-03T14:00:24.000-05:00
Create db only if doesn't exist yet (Thanks @slomek !)
diff --git a/influxdb/client.py b/influxdb/client.py
@@ -454,13 +454,16 @@ def get_list_database(self):
         """
         return list(self.query("SHOW DATABASES").get_points())
 
-    def create_database(self, dbname):
+    def create_database(self, dbname, if_not_exists=False):
         """Create a new database in InfluxDB.
 
         :param dbname: the name of the database to create
         :type dbname: str
         """
-        self.query("CREATE DATABASE \"%s\"" % dbname)
+        if if_not_exists:
+            self.query("CREATE DATABASE IF NOT EXISTS \"%s\"" % dbname)
+        else:
+            self.query("CREATE DATABASE \"%s\"" % dbname)
 
     def drop_database(self, dbname):
         """Drop a database from InfluxDB.
diff --git a/influxdb/tests/client_test.py b/influxdb/tests/client_test.py
@@ -413,6 +413,19 @@ def test_create_database(self):
                 'create database "new_db"'
             )
 
+    def test_create_database_with_exist_check(self):
+        with requests_mock.Mocker() as m:
+            m.register_uri(
+                requests_mock.GET,
+                "http://localhost:8086/query",
+                text='{"results":[{}]}'
+            )
+            self.cli.create_database('new_db', if_not_exists=True)
+            self.assertEqual(
+                m.last_request.qs['q'][0],
+                'create database if not exists "new_db"'
+            )
+
     def test_create_numeric_named_database(self):
         with requests_mock.Mocker() as m:
             m.register_uri(
diff --git a/influxdb/tests/server_tests/client_test_with_server.py b/influxdb/tests/server_tests/client_test_with_server.py
@@ -131,7 +131,12 @@ def test_create_database(self):
             [{'name': 'new_db_1'}, {'name': 'new_db_2'}]
         )
 
-    def test_create_database_fails(self):
+    def test_create_database_twice_if_not_exist(self):
+        self.assertIsNone(self.cli.create_database('new_db'))
+        self.assertIsNone(
+            self.cli.create_database('new_db', if_not_exists=True))
+
+    def test_create_database_twice_fails(self):
         self.assertIsNone(self.cli.create_database('new_db'))
         with self.assertRaises(InfluxDBClientError) as ctx:
             self.cli.create_database('new_db')
