Merge pull request saltstack#29811 from alprs/feature-influxdb_retention

influxdb: add retention policy module functions

salt/modules/influx.py

@@ -419,6 +419,137 @@ def user_remove(name, database=None, user=None, password=None, host=None,
     return client.delete_cluster_admin(name)
 
 
+def retention_policy_get(database,
+                         name,
+                         user=None,
+                         password=None,
+                         host=None,
+                         port=None):
+    '''
+    Get an existing retention policy.
+
+    database
+        The database to operate on.
+
+    name
+        Name of the policy to modify.
+
+    CLI Example:
+
+    .. code-block:: bash
+
+        salt '*' influxdb.retention_policy_get metrics default
+    '''
+    client = _client(user=user, password=password, host=host, port=port)
+
+    for policy in client.get_list_retention_policies(database):
+        if policy['name'] == name:
+            return policy
+
+    return None
+
+
+def retention_policy_exists(database,
+                            name,
+                            user=None,
+                            password=None,
+                            host=None,
+                            port=None):
+    '''
+    Check if a retention policy exists.
+
+    database
+        The database to operate on.
+
+    name
+        Name of the policy to modify.
+
+    CLI Example:
+
+    .. code-block:: bash
+
+        salt '*' influxdb.retention_policy_exists metrics default
+    '''
+    policy = retention_policy_get(name, database, user, password, host, port)
+    return policy is not None
+
+
+def retention_policy_add(database,
+                         name,
+                         duration,
+                         replication,
+                         default=False,
+                         user=None,
+                         password=None,
+                         host=None,
+                         port=None):
+    '''
+    Add a retention policy.
+
+    database
+        The database to operate on.
+
+    name
+        Name of the policy to modify.
+
+    duration
+        How long InfluxDB keeps the data.
+
+    replication
+        How many copies of the data are stored in the cluster.
+
+    default
+        Whether this policy should be the default or not. Default is False.
+
+    CLI Example:
+
+    .. code-block:: bash
+
+        salt '*' influxdb.retention_policy_add metrics default 1d 1
+    '''
+    client = _client(user=user, password=password, host=host, port=port)
+    client.create_retention_policy(name, duration, replication, database, default)
+    return True
+
+
+def retention_policy_alter(database,
+                           name,
+                           duration,
+                           replication,
+                           default=False,
+                           user=None,
+                           password=None,
+                           host=None,
+                           port=None):
+    '''
+    Modify an existing retention policy.
+
+    database
+        The database to operate on.
+
+    name
+        Name of the policy to modify.
+
+    duration
+        How long InfluxDB keeps the data.
+
+    replication
+        How many copies of the data are stored in the cluster.
+
+    default
+        Whether this policy should be the default or not. Default is False.
+
+    CLI Example:
+
+    .. code-block:: bash
+
+        salt '*' influxdb.retention_policy_modify metrics default 1d 1
+    '''
+    client = _client(user=user, password=password, host=host, port=port)
+    client.alter_retention_policy(name, database, duration, replication, default)
+    return True
+
+
 def query(database, query, time_precision='s', chunked=False, user=None,
           password=None, host=None, port=None):
     """

tests/unit/modules/influx_test.py

@@ -249,6 +249,42 @@ def test_query(self):
                                          host='localhost',
                                          port=8000))
 
+    def test_retention_policy_get(self):
+        client = MockInfluxDBClient()
+        policy = {'name': 'foo'}
+        with patch.object(influx, '_client', MagicMock(return_value=client)):
+            client.get_list_retention_policies = MagicMock(return_value=[policy])
+            self.assertEqual(
+                policy,
+                influx.retention_policy_get(database='db', name='foo')
+            )
+
+    def test_retention_policy_add(self):
+        client = MockInfluxDBClient()
+        with patch.object(influx, '_client', MagicMock(return_value=client)):
+            client.create_retention_policy = MagicMock()
+            self.assertTrue(influx.retention_policy_add(
+                database='db',
+                name='name',
+                duration='30d',
+                replication=1,
+            ))
+            client.create_retention_policy.assert_called_once_with(
+                'name', '30d', 1, 'db', False)
+
+    def test_retention_policy_modify(self):
+        client = MockInfluxDBClient()
+        with patch.object(influx, '_client', MagicMock(return_value=client)):
+            client.alter_retention_policy = MagicMock()
+            self.assertTrue(influx.retention_policy_alter(
+                database='db',
+                name='name',
+                duration='30d',
+                replication=1,
+            ))
+            client.alter_retention_policy.assert_called_once_with(
+                'name', 'db', '30d', 1, False)
+
 if __name__ == '__main__':
     from integration import run_tests
     run_tests(InfluxTestCase, needs_daemon=False)