ansible-role-influxdb/library/influxdb_bucket.py

#!/usr/bin/python3

from ansible.module_utils.basic import AnsibleModule
import requests
import json


class InfluxdbBucket:
    def __init__(self, api_url, api_token, name, org):
        self.api_url = api_url
        self.headers = {'Authorization': 'Token {}'.format(api_token)}
        self.name = name
        self.org = org

    def exists(self):
        url = '{}/api/v2/buckets'.format(self.api_url)
        r = requests.get(
            url,
            headers=self.headers,
            params={'name': self.name, 'org': self.org}
        )

        if r.status_code == 404:
            return False
        elif r.status_code == 200:
            data = json.loads(r.text)
            for bucket in data['buckets']:
                if bucket['name'] == self.name:
                    self.id = bucket['id']
                    self.retention = bucket['retentionRules'][0]['everySeconds']
                    self.description = ''
                    if 'description' in bucket:
                        self.description = bucket['description']
                    return True
            return False

        raise Exception(
            'Influxdb',
            'Bad status code {}: {}'.format(r.status_code, r.text)
        )

    def has_changed(self, description, retention):
        if self.description != description:
            return True
        if self.retention != retention:
            return True
        return False

    def get_orgid(self):
        url = '{}/api/v2/orgs'.format(self.api_url)
        r = requests.get(url, headers=self.headers, params={'org': self.org})

        if r.status_code != 200:
            raise Exception(
                'Influxdb',
                'Bad status code {}: {}'.format(r.status_code, r.text)
            )

        data = json.loads(r.text)
        for org in data['orgs']:
            if org['name'] == self.org:
                return org['id']

        raise Exception('Influxdb', 'Don\'t get the orgID')

    def create(self, description, retention):
        url = '{}/api/v2/buckets'.format(self.api_url)
        r = requests.post(
            url,
            headers=self.headers,
            json={
                'name': self.name,
                'description': description,
                'orgID': self.get_orgid(),
                'retentionRules': [{
                    'everySeconds': retention,
                    'shardGroupDurationSeconds': 0,
                    'type': 'expire'
                }]
            }
        )

        if r.status_code != 201:
            raise Exception(
                'Influxdb',
                'Bad status code {}: {}'.format(r.status_code, r.text)
            )

    def delete(self):
        url = '{}/api/v2/buckets/{}'.format(self.api_url, self.id)
        r = requests.delete(url, headers=self.headers)

        if r.status_code != 204:
            raise Exception(
                'Influxdb',
                'Bad status code {}: {}'.format(r.status_code, r.text)
            )

    def update(self, description, retention):
        url = '{}/api/v2/buckets/{}'.format(self.api_url, self.id)
        r = requests.patch(
            url,
            headers=self.headers,
            json={
                'name': self.name,
                'description': description,
                'retentionRules': [{
                    'everySeconds': retention,
                    'shardGroupDurationSeconds': 0,
                    'type': 'expire'
                }]
            }
        )

        if r.status_code != 200:
            raise Exception(
                'Influxdb',
                'Bad status code {}: {}'.format(r.status_code, r.text)
            )


def main():
    fields = {
        'name':        {'type': 'str', 'required': True},
        'org':         {'type': 'str', 'required': True},
        'description': {'type': 'str', 'default': ''},
        'retention':   {'type': 'int', 'default': 0},
        'api_url':     {'type': 'str', 'default': 'http://127.0.0.1:8086'},
        'api_token':   {'type': 'str', 'required': True},
        'state':       {'type': 'str',
                        'default': 'present',
                        'choice': ['present', 'absent']},
    }
    module = AnsibleModule(argument_spec=fields)

    bucket = InfluxdbBucket(
        module.params['api_url'],
        module.params['api_token'],
        module.params['name'],
        module.params['org'],
    )
    changed = False

    if bucket.exists():
        if module.params['state'] == 'absent':
            bucket.delete()
            changed = True
        elif bucket.has_changed(module.params['description'],
                                module.params['retention']):
            bucket.update(
                module.params['description'],
                module.params['retention']
            )
            changed = True
    elif module.params['state'] == 'present':
        bucket.create(module.params['description'], module.params['retention'])
        changed = True

    module.exit_json(changed=changed)


if __name__ == '__main__':
    main()