Add a query_fact_contents() function

bootc · bootc · commit c298d351fa8f · 2016-06-27T19:52:36.000+01:00
diff --git a/docs/examples.rst b/docs/examples.rst
@@ -10,3 +10,9 @@ Obtain named facts from nodes matching a query (using :mod:`pypuppetdb`):
 
 .. literalinclude:: ../examples/facts.py
     :lines: 21-
+
+Obtain selected structured fact values from nodes matching a query (using
+:mod:`pypuppetdb`):
+
+.. literalinclude:: ../examples/fact_contents.py
+    :lines: 21-
diff --git a/examples/fact_contents.py b/examples/fact_contents.py
@@ -0,0 +1,33 @@
+# -*- coding: utf-8 -*-
+#
+# This file is part of pypuppetdbquery.
+# Copyright © 2016  Chris Boot <bootc@bootc.net>
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""
+Obtain selected structured fact values from nodes matching a query.
+"""
+
+import pypuppetdb
+import pypuppetdbquery
+
+pdb = pypuppetdb.connect()
+
+node_facts = pypuppetdbquery.query_fact_contents(
+    pdb,
+    '(processorcount=4 or processorcount=8) and kernel=Linux',
+    ['system_uptime.days', 'os.lsb.~"dist.*"'])
+
+for node in node_facts:
+    facts = node_facts[node]
+    print(node, facts)
diff --git a/pypuppetdbquery/__init__.py b/pypuppetdbquery/__init__.py
@@ -24,6 +24,7 @@
 
 from collections import defaultdict
 from json import dumps as json_dumps
+from ply.yacc import NullLogger
 from .evaluator import Evaluator
 from .parser import Parser
 
@@ -78,6 +79,11 @@ def query_facts(pdb, s, facts=None, raw=False, lex_options=None,
     values. If `True` it returns raw :class:`pypuppetdb.types.Fact` objects as
     :meth:`pypuppetdb.api.BaseAPI.nodes` does.
 
+    .. note:: This function can return only full facts, not elements of
+        structured facts. For example, only the whole ``os`` fact may be
+        returned but not the ``os.family`` key within the larger structured
+        fact. If you need to do this, look at :func:`query_fact_contents`.
+
     :param pypuppetdb.api.BaseAPI pdb: pypuppetdb connection to query from
     :param str s: The query string (may be empty to query all nodes)
     :param Sequence facts: List of fact names to search for
@@ -114,3 +120,75 @@ def query_facts(pdb, s, facts=None, raw=False, lex_options=None,
     for fact in facts:
         ret[fact.node][fact.name] = fact.value
     return ret
+
+
+def query_fact_contents(pdb, s, facts=None, raw=False, lex_options=None,
+                        yacc_options=None):
+    """
+    Helper to query PuppetDB for fact contents (i.e. within structured facts)
+    on nodes matching a query string.
+
+    Adjusts the query to return only those strucutred fact keys requested in
+    the function call.
+
+    The facts listed in the `facts` list are run through the query parser and
+    treated as "identifier paths". This means the same rules apply as for
+    within the query language, e.g. ``foo.bar`` or ``foo.*`` or even
+    ``foo.~"bar.*"``.
+
+    If `raw` is `False` (the default), the return value is a :class:`dict`
+    with node names as keys containing a :class:`dict` of flattened fact paths
+    to fact values. If `True` it returns raw query output: a list of
+    dictionaries (see the `PuppetDB fact-contents documentation
+    <https://docs.puppet.com/puppetdb/4.1/api/query/v4/fact-contents.html#response-format>`__).
+
+    .. note:: This function can only be used to search deeply within structured
+        facts. It cannot return a whole structured fact, only individual
+        elements within—but you can return all the elements within a structured
+        fact if you want by using a regex match.
+
+    :param pypuppetdb.api.BaseAPI pdb: pypuppetdb connection to query from
+    :param str s: The query string (may be empty to query all nodes)
+    :param Sequence facts: List of fact paths to search for
+    :param bool raw: Whether to skip post-processing the facts into a dict
+        structure grouped by node
+    :param dict lex_options: Options passed to :func:`ply.lex.lex`
+    :param dict yacc_options: Options passed to :func:`ply.yacc.yacc`
+    """
+    query = parse(s, json=False, mode='facts', lex_options=lex_options,
+                  yacc_options=yacc_options)
+
+    if facts:
+        # We need custom optiosn to start with identifier_path, but that then
+        # causes warnings to be issued for unreachable symbols so we silence
+        # those with the NullLogger.
+        yacc_opt_id = dict(yacc_options) if yacc_options else {}
+        yacc_opt_id['errorlog'] = NullLogger()
+        yacc_opt_id['start'] = 'identifier_path'
+
+        parser = Parser(lex_options=lex_options, yacc_options=yacc_opt_id)
+        evaluator = Evaluator()
+
+        factquery = ['or']
+        for fact in facts:
+            ast = parser.parse(fact)
+            factquery.append(evaluator.evaluate(ast, mode='facts'))
+
+        if query:
+            query = ['and', query, factquery]
+        else:
+            query = factquery
+
+    if query is None:
+        return None
+
+    facts = pdb.fact_contents(query=json_dumps(query))
+    if raw:
+        return facts
+
+    ret = defaultdict(dict)
+    for fact in facts:
+        node = fact['certname']
+        name = '.'.join(fact['path'])
+        ret[node][name] = fact['value']
+    return ret
diff --git a/tests/test_frontend.py b/tests/test_frontend.py
@@ -19,7 +19,7 @@
 import mock
 import unittest
 
-from pypuppetdbquery import parse, query_facts
+from pypuppetdbquery import parse, query_facts, query_fact_contents
 
 
 class _FakeNode(object):
@@ -168,3 +168,109 @@ def test_query_facts_in_raw_mode(self):
         node_facts = self._query_facts(mock_pdb, 'foo=bar', raw=True)
 
         self.assertEquals(node_facts, mock_pdb.facts.return_value)
+
+
+class TestFrontendQueryFactContents(unittest.TestCase):
+    """
+    Test cases targetting :func:`pypuppetdbquery.query_fact_contents`.
+    """
+
+    def _query_fact_contents(self, pdb, s, facts=None, raw=False):
+        return query_fact_contents(
+            pdb, s, facts, raw,
+            lex_options={
+                'debug': False,
+                'optimize': False,
+            },
+            yacc_options={
+                'debug': False,
+                'optimize': False,
+                'write_tables': False,
+            })
+
+    def test_with_query_and_facts_list(self):
+        mock_pdb = mock.NonCallableMock()
+        mock_pdb.fact_contents = mock.Mock(return_value=[
+            {
+                'value': 14,
+                'certname': 'alpha',
+                'environment': 'production',
+                'path': ['system_uptime', 'days'],
+                'name': 'system_uptime',
+            },
+        ])
+
+        out = self._query_fact_contents(
+            mock_pdb, 'foo=bar', ['system_uptime.days'])
+
+        mock_pdb.fact_contents.assert_called_once_with(query=json.dumps([
+            'and',
+            ['in', 'certname',
+             ['extract', 'certname',
+              ['select_fact_contents',
+               ['and',
+                ['=', 'path', ['foo']],
+                ['=', 'value', 'bar']]]]],
+            ['or',
+             ['=', 'path',
+              ['system_uptime', 'days']]]]))
+
+        self.assertEquals(out, {
+            'alpha': {'system_uptime.days': 14},
+        })
+
+    def test_without_query(self):
+        mock_pdb = mock.NonCallableMock()
+        mock_pdb.fact_contents = mock.Mock(return_value=[
+            {
+                'value': 14,
+                'certname': 'alpha',
+                'environment': 'production',
+                'path': ['system_uptime', 'days'],
+                'name': 'system_uptime',
+            },
+        ])
+
+        out = self._query_fact_contents(mock_pdb, '', ['system_uptime.days'])
+
+        mock_pdb.fact_contents.assert_called_once_with(query=json.dumps([
+            'or',
+            ['=', 'path',
+             ['system_uptime', 'days']]]))
+
+        self.assertEquals(out, {
+            'alpha': {'system_uptime.days': 14},
+        })
+
+    def test_without_either(self):
+        out = self._query_fact_contents(None, '')
+        self.assertTrue(out is None)
+
+    def test_raw_output(self):
+        mock_pdb = mock.NonCallableMock()
+        mock_pdb.fact_contents = mock.Mock(return_value=[
+            {
+                'value': 14,
+                'certname': 'alpha',
+                'environment': 'production',
+                'path': ['system_uptime', 'days'],
+                'name': 'system_uptime',
+            },
+        ])
+
+        out = self._query_fact_contents(
+            mock_pdb, 'foo=bar', ['system_uptime.days'], True)
+
+        mock_pdb.fact_contents.assert_called_once_with(query=json.dumps([
+            'and',
+            ['in', 'certname',
+             ['extract', 'certname',
+              ['select_fact_contents',
+               ['and',
+                ['=', 'path', ['foo']],
+                ['=', 'value', 'bar']]]]],
+            ['or',
+             ['=', 'path',
+              ['system_uptime', 'days']]]]))
+
+        self.assertEquals(out, mock_pdb.fact_contents.return_value)
