1from __future__ import absolute_import
2from __future__ import unicode_literals
3
4import datetime
5import json
6import logging
7
8from pypuppetdb.errors import APIError
9
10log = logging.getLogger(__name__)
11
12
13class BinaryOperator(object):
14    """
15    This is a parent helper class used to create PuppetDB AST queries
16    for single key-value pairs for the available operators.
17
18    It is possible to directly declare the various types of queries
19    from this class. For instance the code
20    BinaryOperator('=', 'certname', 'node1.example.com') generates the
21    PuppetDB query '["=", "certname", "node1.example.com"]'. It is preferred
22    to use the child classes as they may have restrictions specific
23    to that operator.
24
25    See
26    https://docs.puppet.com/puppetdb/4.0/api/query/v4/ast.html#binary-operators
27    for more information.
28
29    :param operator: The binary query operation performed. There is
30        no value checking on this field.
31    :type operator: :obj:`string`
32    :param field: The PuppetDB endpoint query field. See endpoint
33                  documentation for valid values.
34    :type field: any
35    :param value: The values of the field to match, or not match.
36    :type value: any
37    """
38
39    def __init__(self, operator, field, value):
40        if isinstance(value, datetime.datetime):
41            value = str(value)
42        self.data = [operator, field, value]
43
44    def __repr__(self):
45        return 'Query: {0}'.format(self)
46
47    def __str__(self):
48        return json.dumps(self.json_data())
49
50    def json_data(self):
51        return self.data
52
53
54class BooleanOperator(object):
55    """
56    This is a parent helper class used to create PuppetDB AST queries
57    for available boolean queries.
58
59    It is possible to directly declare a boolean query from this class.
60    For instance the code BooleanOperator("and") will create an empty
61    query '["and",]'. An error will be raised if there are no queries
62    added via :func:`~pypuppetdb.QueryBuilder.BooleanOperator.add`
63
64    See
65    https://docs.puppet.com/puppetdb/4.0/api/query/v4/ast.html#binary-operators
66    for more information.
67
68    :param operator: The boolean query operation to perform.
69    :type operator: :obj:`string`
70    """
71
72    def __init__(self, operator):
73        self.operator = operator
74        self.operations = []
75
76    def add(self, query):
77        if type(query) == list:
78            for i in query:
79                self.add(i)
80        elif type(query) == str:
81            self.operations.append(json.loads(query))
82        elif isinstance(query, (BinaryOperator, InOperator,
83                                BooleanOperator)):
84            self.operations.append(query.json_data())
85        else:
86            raise APIError("Can only accept fixed-string queries, arrays " +
87                           "or operator objects")
88
89    def __repr__(self):
90        return 'Query: {0}'.format(self)
91
92    def __str__(self):
93        return json.dumps(self.json_data())
94
95    def json_data(self):
96        if len(self.operations) == 0:
97            raise APIError("At least one query operation is required")
98        return [self.operator] + self.operations
99
100
101class ExtractOperator(object):
102    """
103    Queries that either do not or cannot require all the key-value pairs
104    from an endpoint can use the Extract Operator as described in
105    https://docs.puppet.com/puppetdb/4.1/api/query/v4/ast.html#projection-operators.
106
107    The syntax of this operator requires a function and/or a list of fields,
108    an optional standard query and an optional group by clause including a
109    list of fields.
110    """
111
112    def __init__(self):
113        self.fields = []
114        self.query = None
115        self.group_by = []
116
117    def add_field(self, field):
118        if isinstance(field, list):
119            for i in field:
120                self.add_field(i)
121        elif isinstance(field, str):
122            self.fields.append(field)
123        elif isinstance(field, FunctionOperator):
124            self.fields.append(field.json_data())
125        else:
126            raise APIError("ExtractOperator.add_field only supports "
127                           "lists and strings")
128
129    def add_query(self, query):
130        if self.query is not None:
131            raise APIError("Only one query is supported by ExtractOperator")
132        elif isinstance(query, str):
133            self.query = json.loads(query)
134        elif isinstance(query, (BinaryOperator, SubqueryOperator,
135                                BooleanOperator)):
136            self.query = query.json_data()
137        else:
138            raise APIError("ExtractOperator.add_query only supports "
139                           "strings, BinaryOperator, BooleanOperator "
140                           "and SubqueryOperator objects")
141
142    def add_group_by(self, field):
143        if isinstance(field, list):
144            for i in field:
145                self.add_group_by(i)
146        elif isinstance(field, str):
147            if len(self.group_by) == 0:
148                self.group_by.append('group_by')
149            self.group_by.append(field)
150        elif isinstance(field, FunctionOperator):
151            if len(self.group_by) == 0:
152                self.group_by.append('group_by')
153            self.group_by.append(field.json_data())
154        else:
155            raise APIError("ExtractOperator.add_group_by only supports "
156                           "lists, strings, and FunctionOperator objects")
157
158    def __repr__(self):
159        return 'Query: {0}'.format(self)
160
161    def __str__(self):
162        return json.dumps(self.json_data())
163
164    def json_data(self):
165        if len(self.fields) == 0:
166            raise APIError("ExtractOperator needs at least one field")
167
168        arr = ['extract', self.fields]
169
170        if self.query is not None:
171            arr.append(self.query)
172        if len(self.group_by) > 0:
173            arr.append(self.group_by)
174
175        return arr
176
177
178class FunctionOperator(object):
179    """
180    Performs an aggregate function on the result of a subquery, full
181    documentation is available at
182    https://docs.puppet.com/puppetdb/4.1/api/query/v4/ast.html#function.
183    This object can only be used in the field list or group by list of
184    an ExtractOperator object.
185
186    :param function: The name of the function to perform.
187    :type function: :obj:`str`
188    :param field: The name of the field to perform the function on. All
189        functions with the exception of count require this value.
190    :type field: :obj:`str`
191    """
192
193    def __init__(self, function, field=None, fmt=None):
194        if function not in ['count', 'avg', 'sum', 'min', 'max', 'to_string']:
195            raise APIError("Unsupport function: {0}".format(function))
196        elif function != "count" and field is None:
197            raise APIError("Function {0} requires a field value".format(
198                function))
199        elif function == 'to_string' and fmt is None:
200            raise APIError("Function {0} requires an extra 'fmt' parameter")
201
202        self.arr = ['function', function]
203
204        if field is not None:
205            self.arr.append(field)
206
207        if function == 'to_string':
208            self.arr.append(fmt)
209
210    def __repr__(self):
211        return 'Query: {0}'.format(self)
212
213    def __str__(self):
214        return json.dumps(self.json_data())
215
216    def json_data(self):
217        return self.arr
218
219
220class SubqueryOperator(object):
221    """
222    Performs a subquery to another puppetDB object, full
223    documentation is available at
224    https://docs.puppet.com/puppetdb/3.2/api/query/v4/operators.html#subquery-operators
225    This object must be used in combination with the InOperator according
226    to documentation.
227
228    :param endpoint: The name of the subquery object
229    :type function: :obj:`str`
230    """
231
232    def __init__(self, endpoint):
233        if endpoint not in ['catalogs', 'edges', 'environments', 'events',
234                            'facts', 'fact_contents', 'fact_paths', 'nodes',
235                            'reports', 'resources']:
236            raise APIError("Unsupported endpoint: {0}".format(endpoint))
237
238        self.query = None
239        self.arr = ['select_{0}'.format(endpoint)]
240
241    def add_query(self, query):
242        if self.query is not None:
243            raise APIError("Only one query is supported by ExtractOperator")
244        else:
245            self.query = True
246            self.arr.append(query.json_data())
247
248    def __repr__(self):
249        return 'Query: {0}'.format(self)
250
251    def __str__(self):
252        return json.dumps(self.json_data())
253
254    def json_data(self):
255        return self.arr
256
257
258class InOperator(object):
259    """
260    Performs boolean compare between a field a subquery result
261    https://docs.puppet.com/puppetdb/3.2/api/query/v4/operators.html#subquery-operators
262    This object must be used in combination with the SubqueryOperator according
263    to documentation.
264
265    :param field: The name of the subquery object
266    :type function: :obj:`str`
267    """
268
269    def __init__(self, field):
270        self.query = None
271        self.arr = ['in', field]
272
273    def add_query(self, query):
274        if self.query is not None:
275            raise APIError("Only one query is supported by ExtractOperator")
276        elif isinstance(query, str):
277            self.query = True
278            self.arr.append(json.loads(query))
279        elif isinstance(query, (ExtractOperator, FromOperator)):
280            self.query = True
281            self.arr.append(query.json_data())
282        else:
283            raise APIError("InOperator.add_query only supports "
284                           "strings, ExtractOperator, and"
285                           "FromOperator objects")
286
287    def add_array(self, values):
288        if self.query is not None:
289            raise APIError("Only one array is supported by the InOperator")
290        elif isinstance(values, list):
291            def depth(a_list):
292                return (isinstance(a_list, list) and len(a_list) != 0) \
293                       and max(map(depth, a_list)) + 1
294
295            if depth(values) == 1:
296                self.query = True
297                self.arr.append(['array', values])
298            else:
299                raise APIError("InOperator.add_array: cannot pass in "
300                               "nested arrays (or empty arrays)")
301        else:
302            raise APIError("InOperator.add_array: Ill-formatted array, "
303                           "must be of the format: "
304                           "['array', [<array values>]]")
305
306    def __repr__(self):
307        return 'Query: {0}'.format(self)
308
309    def __str__(self):
310        return json.dumps(self.json_data())
311
312    def json_data(self):
313        return self.arr
314
315
316class FromOperator(object):
317    """
318    From contextual operator that allows for queries on the root endpoint
319    or subqueries into other entities:
320    https://puppet.com/docs/puppetdb/5.1/api/query/v4/ast.html#from
321
322    Ex.)
323    fr = FromOperator("facts")
324    fr.add_query(EqualsOperator("foo", "bar"))
325    fr.add_order_by(["certname"])
326    fr.add_limit(10)
327
328    note: only supports single entity From operations
329    """
330
331    def __init__(self, endpoint):
332        valid_entities = ["aggregate_event_counts", "catalogs", "edges",
333                          "environments", "event_counts", "events", "facts",
334                          "fact_contents", "fact_names", "fact_paths", "nodes",
335                          "producers", "reports", "resources"]
336
337        if endpoint in valid_entities:
338            self.endpoint = endpoint
339        else:
340            raise APIError("Endpoint is invalid. Must be "
341                           "one of the following : %s"
342                           % valid_entities)
343
344        self.query = None
345        self.order_by = []
346        self.limit = None
347        self.offset = None
348
349    def add_query(self, query):
350        if self.query is not None:
351            raise APIError("Only one main query is supported by FromOperator")
352        elif isinstance(query, (InOperator, ExtractOperator,
353                                BinaryOperator, BooleanOperator,
354                                FunctionOperator)):
355            self.query = query.json_data()
356        else:
357            raise APIError("FromOperator.add_field only supports "
358                           "Operator Objects")
359
360    def add_order_by(self, fields):
361        def depth(a_list):
362            return isinstance(a_list, list) and max(map(depth, a_list)) + 1
363
364        fields_depth = depth(fields)
365
366        if isinstance(fields, list):
367            if fields_depth == 1 or fields_depth == 2:
368                self.order_by = fields
369            else:
370                raise APIError("ExtractOperator.add_order_by only "
371                               "supports lists of fields of depth "
372                               "one or two: [value, <desc/asc>] or "
373                               "[value]")
374        else:
375            raise APIError("ExtractOperator.add_group_by only supports "
376                           "lists of one or more fields")
377
378    def add_limit(self, lim):
379        if isinstance(lim, int):
380            self.limit = lim
381        else:
382            raise APIError("ExtractOperator.add_limit only supports ints")
383
384    def add_offset(self, off):
385        if isinstance(off, int):
386            self.offset = off
387        else:
388            raise APIError("ExtractOperator.add_offset only supports ints")
389
390    def __repr__(self):
391        return 'Query: {0}'.format(self)
392
393    def __str__(self):
394        return json.dumps(self.json_data())
395
396    def json_data(self):
397        if self.query is None:
398            raise APIError("FromOperator needs one main query")
399
400        arr = ['from', self.endpoint, self.query]
401
402        if len(self.order_by) > 0:
403            arr.append(['order_by', self.order_by])
404        if self.limit is not None:
405            arr.append(['limit', self.limit])
406        if self.offset is not None:
407            arr.append(['offset', self.offset])
408
409        return arr
410
411
412class EqualsOperator(BinaryOperator):
413    """
414    Builds an equality filter based on the supplied field-value pair as
415    described
416    https://docs.puppet.com/puppetdb/4.1/api/query/v4/ast.html#equality.
417
418    In order to create the following query:
419
420    ["=", "environment", "production"]
421
422    The following code can be used.
423
424    EqualsOperator('environment', 'production')
425
426    :param field: The PuppetDB endpoint query field. See endpoint
427                  documentation for valid values.
428    :type field: any
429    :param value: The value of the field to match, or not match.
430    :type value: any
431    """
432
433    def __init__(self, field, value):
434        super(EqualsOperator, self).__init__("=", field, value)
435
436
437class GreaterOperator(BinaryOperator):
438    """
439    Builds a greater-than filter based on the supplied field-value pair as
440    described
441    https://docs.puppet.com/puppetdb/4.1/api/query/v4/ast.html#greater-than.
442
443    In order to create the following query:
444
445    [">", "catalog_timestamp", "2016-06-01 00:00:00"]
446
447    The following code can be used.
448
449    GreaterOperator('catalog_timestamp', datetime.datetime(2016, 06, 01))
450
451    :param field: The PuppetDB endpoint query field. See endpoint
452                  documentation for valid values.
453    :type field: any
454    :param value: Matches if the field is greater than this value.
455    :type value: Number, timestamp or array
456    """
457
458    def __init__(self, field, value):
459        super(GreaterOperator, self).__init__(">", field, value)
460
461
462class LessOperator(BinaryOperator):
463    """
464    Builds a less-than filter based on the supplied field-value pair as
465    described
466    https://docs.puppet.com/puppetdb/4.1/api/query/v4/ast.html#less-than.
467
468    In order to create the following query:
469
470    ["<", "catalog_timestamp", "2016-06-01 00:00:00"]
471
472    The following code can be used.
473
474    LessOperator('catalog_timestamp', datetime.datetime(2016, 06, 01))
475
476    :param field: The PuppetDB endpoint query field. See endpoint
477                  documentation for valid values.
478    :type field: any
479    :param value: Matches if the field is less than this value.
480    :type value: Number, timestamp or array
481    """
482
483    def __init__(self, field, value):
484        super(LessOperator, self).__init__("<", field, value)
485
486
487class GreaterEqualOperator(BinaryOperator):
488    """
489    Builds a greater-than or equal-to filter based on the supplied
490    field-value pair as described
491    https://docs.puppet.com/puppetdb/4.1/api/query/v4/ast.html#greater-than-or-equal-to.
492
493    In order to create the following query:
494
495    [">=", "facts_timestamp", "2016-06-01 00:00:00"]
496
497    The following code can be used.
498
499    GreaterEqualOperator('facts_timestamp', datetime.datetime(2016, 06, 01))
500
501    :param field: The PuppetDB endpoint query field. See endpoint
502                  documentation for valid values.
503    :type field: any
504    :param value: Matches if the field is greater than or equal to\
505        this value.
506    :type value: Number, timestamp or array
507    """
508
509    def __init__(self, field, value):
510        super(GreaterEqualOperator, self).__init__(">=", field, value)
511
512
513class LessEqualOperator(BinaryOperator):
514    """
515    Builds a less-than or equal-to filter based on the supplied
516    field-value pair as described
517    https://docs.puppet.com/puppetdb/4.1/api/query/v4/ast.html#less-than-or-equal-to.
518
519    In order to create the following query:
520
521    ["<=", "facts_timestamp", "2016-06-01 00:00:00"]
522
523    The following code can be used.
524
525    LessEqualOperator('facts_timestamp', datetime.datetime(2016, 06, 01))
526
527    :param field: The PuppetDB endpoint query field. See endpoint
528                  documentation for valid values.
529    :type field: any
530    :param value: Matches if the field is less than or equal to\
531        this value.
532    :type value: Number, timestamp or array
533    """
534
535    def __init__(self, field, value):
536        super(LessEqualOperator, self).__init__("<=", field, value)
537
538
539class RegexOperator(BinaryOperator):
540    """
541    Builds a regular expression filter based on the supplied field-value
542    pair as described
543    https://docs.puppet.com/puppetdb/4.1/api/query/v4/ast.html#regexp-match.
544
545    In order to create the following query:
546
547    ["~", "certname", "www\\d+\\.example\\.com"]
548
549    The following code can be used.
550
551    RegexOperator('certname', 'www\\d+\\.example\\.com')
552
553    :param field: The PuppetDB endpoint query field. See endpoint
554                  documentation for valid values.
555    :type field: any
556    :param value: Matches if the field matches this regular expression.
557    :type value: :obj:`string`
558    """
559
560    def __init__(self, field, value):
561        super(RegexOperator, self).__init__("~", field, value)
562
563
564class RegexArrayOperator(BinaryOperator):
565    """
566    Builds a regular expression array filter based on the supplied
567    field-value pair. This query only works on fields with paths as
568    described
569    https://docs.puppet.com/puppetdb/4.1/api/query/v4/ast.html#regexp-array-match.
570
571    In order to create the following query:
572
573    ["~", "path", ["networking", "eth.*", "macaddress"]]
574
575    The following code can be used.
576
577    RegexArrayOperator('path', ["networking", "eth.*", "macaddress"])
578
579    :param field: The PuppetDB endpoint query field. See endpoint
580                  documentation for valid values.
581    :type field: any
582    :param value: Matches if the field matches this regular expression.
583    :type value: :obj:`list`
584    """
585
586    def __init__(self, field, value):
587        super(RegexArrayOperator, self).__init__("~>", field, value)
588
589
590class NullOperator(BinaryOperator):
591    """
592    Builds a null filter based on the field and boolean value pair as
593    described
594    https://docs.puppet.com/puppetdb/4.1/api/query/v4/ast.html#null-is-null.
595    This filter only works on field that may be null. Value may only
596    be True or False.
597
598    In order to create the following query:
599
600    ["null?", "deactivated", true]
601
602    The following code can be used.
603
604    NullOperator('deactivated', True)
605
606    :param field: The PuppetDB endpoint query field. See endpoint
607                  documentation for valid values.
608    :type field: any
609    :param value: Matches if the field value is null (if True) or\
610        not null (if False)
611    :type value: :obj:`bool`
612    """
613
614    def __init__(self, field, value):
615        if type(value) != bool:
616            raise APIError("NullOperator value must be boolean")
617
618        super(NullOperator, self).__init__("null?", field, value)
619
620
621class AndOperator(BooleanOperator):
622    """
623    Builds an AND boolean filter. Only results that match ALL
624    criteria from the included query strings will be returned
625    from PuppetDB. Full documentation is available
626    https://docs.puppet.com/puppetdb/4.1/api/query/v4/ast.html#and
627
628    In order to create the following query:
629
630    ["and",
631      ["=", "catalog_environment", "production"],
632      ["=", "facts_environment", "production"]]
633
634    The following code can be used:
635
636    op = AndOperator()
637    op.add(EqualsOperator("catalog_environment", "production"))
638    op.add(EqualsOperator("facts_environment", "production"))
639    """
640
641    def __init__(self):
642        super(AndOperator, self).__init__("and")
643
644
645class OrOperator(BooleanOperator):
646    """
647    Builds an OR boolean filter. Only results that match ANY
648    criteria from the included query strings will be returned
649    from PuppetDB. Full documentation is available
650    https://docs.puppet.com/puppetdb/4.1/api/query/v4/ast.html#or.
651
652    In order to create the following query:
653
654    ["or", ["=", "name", "hostname"], ["=", "name", "architecture"]]
655
656    The following code can be used:
657
658    op = OrOperator()
659    op.add(EqualsOperator("name", "hostname"))
660    op.add(EqualsOperator("name", "architecture"))
661    """
662
663    def __init__(self):
664        super(OrOperator, self).__init__("or")
665
666
667class NotOperator(BooleanOperator):
668    """
669    Builds a NOT boolean filter. Only results that DO NOT match
670    criteria from the included query strings will be returned
671    from PuppetDB. Full documentation is available
672    https://docs.puppet.com/puppetdb/4.1/api/query/v4/ast.html#not
673
674    Unlike the other Boolean Operator objects this operator only
675    accepts a single query string.
676
677    In order to create the following query:
678
679    ["not", ["=", "osfamily", "RedHat"]]
680
681    The following code can be used.
682
683    op = NotOperator()
684    op.add(EqualsOperator("osfamily", "RedHat"))
685    """
686
687    def __init__(self):
688        super(NotOperator, self).__init__("not")
689
690    def add(self, query):
691        if len(self.operations) > 0:
692            raise APIError("This operator only accept one query string")
693        elif isinstance(query, list) and len(query) > 1:
694            raise APIError("This operator only accept one query string")
695        super(NotOperator, self).add(query)
696