Added create_retention_policy

aviau · aviau · commit 2d02e59206e6 · 2015-03-01T21:21:28.000-05:00
diff --git a/influxdb/client.py b/influxdb/client.py
@@ -334,6 +334,26 @@ def drop_database(self, dbname):
         """
         self.query("DROP DATABASE %s" % dbname)
 
+    def create_retention_policy(self, name, duration, replication, database=None, default=False):
+        """
+        Create a retention policy
+
+        :param duration: The duration. Ex: '1d'
+        :param replication: The replication.
+        :param database: The database. Defaults to current database
+        :param default: (bool) Wether or not to set the policy as default
+        """
+
+        query_string = \
+            "CREATE RETENTION POLICY %s ON %s " \
+            "DURATION %s REPLICATION %s" % \
+            (name, database or self._database, duration, replication)
+
+        if default is True:
+            query_string += " DEFAULT"
+
+        self.query(query_string)
+
     def get_list_retention_policies(self, database=None):
         """
         Get the list of retention policies
diff --git a/tests/influxdb/client_test.py b/tests/influxdb/client_test.py
@@ -372,6 +372,42 @@ def test_get_list_series(self):
                   u'name': u'fsfdsdf', u'replicaN': 2}]
             )
 
+    def test_create_retention_policy_default(self):
+        example_response = u'{"results":[{}]}'
+
+        with requests_mock.Mocker() as m:
+            m.register_uri(
+                requests_mock.GET,
+                "http://localhost:8086/query",
+                text=example_response
+            )
+            self.cli.create_retention_policy(
+                'somename', '1d', 4, default=True, database='db'
+            )
+
+            self.assertEqual(
+                m.last_request.qs['q'][0],
+                'create retention policy somename on db duration 1d replication 4 default'
+            )
+
+    def test_create_retention_policy(self):
+        example_response = u'{"results":[{}]}'
+
+        with requests_mock.Mocker() as m:
+            m.register_uri(
+                requests_mock.GET,
+                "http://localhost:8086/query",
+                text=example_response
+            )
+            self.cli.create_retention_policy(
+                'somename', '1d', 4, database='db'
+            )
+
+            self.assertEqual(
+                m.last_request.qs['q'][0],
+                'create retention policy somename on db duration 1d replication 4'
+            )
+
     def test_get_list_retention_policies(self):
         example_response = \
             u'{"results": [{"rows": [{"values": [["fsfdsdf", "24h0m0s", 2]],' \
