Merge pull request #382 from rundeck-plugins/enh/run-2670_ansible-mod…

…el-source-yaml-parameter

RUN-2670: Enh: Added yaml data size parameter at model source

src/main/groovy/com/rundeck/plugins/ansible/ansible/AnsibleDescribable.java

@@ -158,6 +158,8 @@ public static String[] getValues() {
 
     public static final String ANSIBLE_ENCRYPT_EXTRA_VARS = "ansible-encrypt-extra-vars";
 
+    String ANSIBLE_YAML_DATA_SIZE = "ansible-yaml-data-size";
+
     public static Property PLAYBOOK_PATH_PROP = PropertyUtil.string(
     			ANSIBLE_PLAYBOOK_PATH,
               "Playbook",
@@ -527,4 +529,13 @@ public static String[] getValues() {
             .title("Encrypt Extra Vars.")
             .description("Encrypt the value of the extra vars keys.")
             .build();
+
+    Property YAML_DATA_SIZE_PROP = PropertyBuilder.builder()
+            .integer(ANSIBLE_YAML_DATA_SIZE)
+            .required(false)
+            .title("Inventory Yaml Data Size")
+            .description("Set the MB size (Default value is 10)"+
+                    " therefore, the plugin can process the yaml data response coming from Ansible."+
+                    " (This only applies when Gather Facts = No)")
+            .build();
 }

src/main/groovy/com/rundeck/plugins/ansible/plugin/AnsibleResourceModelSource.java

@@ -25,12 +25,14 @@
 import com.rundeck.plugins.ansible.ansible.AnsibleRunner;
 import com.rundeck.plugins.ansible.ansible.InventoryList;
 import com.rundeck.plugins.ansible.util.VaultPrompt;
+import lombok.Setter;
 import org.rundeck.app.spi.Services;
 import org.rundeck.storage.api.PathUtil;
 import org.rundeck.storage.api.StorageException;
 import org.yaml.snakeyaml.LoaderOptions;
 import org.yaml.snakeyaml.Yaml;
 import org.yaml.snakeyaml.constructor.SafeConstructor;
+import org.yaml.snakeyaml.error.YAMLException;
 
 import java.io.BufferedReader;
 import java.io.ByteArrayOutputStream;
@@ -45,6 +47,7 @@
 import java.nio.file.SimpleFileVisitor;
 import java.nio.file.attribute.BasicFileAttributes;
 import java.util.ArrayList;
+import java.util.Collections;
 import java.util.HashMap;
 import java.util.HashSet;
 import java.util.List;
@@ -60,8 +63,10 @@ public class AnsibleResourceModelSource implements ResourceModelSource, ProxyRun
   public static final String HOST_TPL_J2 = "host-tpl.j2";
   public static final String GATHER_HOSTS_YML = "gather-hosts.yml";
 
+  @Setter
   private Framework framework;
 
+  @Setter
   private Services services;
 
   private String project;
@@ -72,6 +77,8 @@ public class AnsibleResourceModelSource implements ResourceModelSource, ProxyRun
 
   private String inventory;
   private boolean gatherFacts;
+  @Setter
+  private Integer yamlDataSize;
   private boolean ignoreErrors = false;
   private String limit;
   private String ignoreTagPrefix;
@@ -118,17 +125,14 @@ public class AnsibleResourceModelSource implements ResourceModelSource, ProxyRun
 
   protected boolean encryptExtraVars = false;
 
+  @Setter
   private AnsibleInventoryList.AnsibleInventoryListBuilder ansibleInventoryListBuilder = null;
 
   public AnsibleResourceModelSource(final Framework framework) {
       this.framework = framework;
   }
 
-  public void setAnsibleInventoryListBuilder(AnsibleInventoryList.AnsibleInventoryListBuilder builder) {
-    this.ansibleInventoryListBuilder = builder;
-  }
-
-  private static String resolveProperty(
+    private static String resolveProperty(
             final String attribute,
             final String defaultValue,
             final Properties configuration,
@@ -142,18 +146,32 @@ private static String resolveProperty(
         }
   }
 
+  private static Integer resolveIntProperty(
+          final String attribute,
+          final Integer defaultValue,
+          final Properties configuration,
+          final Map<String, Map<String, String>> dataContext) throws ConfigurationException {
+      final String strValue = resolveProperty(attribute, null, configuration, dataContext);
+      if (null != strValue) {
+          try {
+              return Integer.parseInt(strValue);
+          } catch (NumberFormatException e) {
+              throw new ConfigurationException("Can't parse attribute :" + attribute +
+                      ", value: " + strValue +
+                      " Expected Integer. : " + e.getMessage(), e);
+          }
+      }
+      return defaultValue;
+  }
+
   private static Boolean skipVar(final String hostVar, final List<String> varList) {
     for (final String specialVarString : varList) {
       if (hostVar.startsWith(specialVarString)) return true;
     }
     return false;
   }
 
-  public void setServices(Services services) {
-    this.services = services;
-  }
-
-  public void configure(Properties configuration) throws ConfigurationException {
+    public void configure(Properties configuration) throws ConfigurationException {
 
     project = configuration.getProperty("project");
     configDataContext = new HashMap<String, Map<String, String>>();
@@ -167,6 +185,8 @@ public void configure(Properties configuration) throws ConfigurationException {
     gatherFacts = "true".equals(resolveProperty(AnsibleDescribable.ANSIBLE_GATHER_FACTS,null,configuration,executionDataContext));
     ignoreErrors = "true".equals(resolveProperty(AnsibleDescribable.ANSIBLE_IGNORE_ERRORS,null,configuration,executionDataContext));
 
+    yamlDataSize = resolveIntProperty(AnsibleDescribable.ANSIBLE_YAML_DATA_SIZE,10, configuration, executionDataContext);
+
     limit = (String) resolveProperty(AnsibleDescribable.ANSIBLE_LIMIT,null,configuration,executionDataContext);
     ignoreTagPrefix = (String) resolveProperty(AnsibleDescribable.ANSIBLE_IGNORE_TAGS,null,configuration,executionDataContext);
 
@@ -670,14 +690,22 @@ public FileVisitResult postVisitDirectory(Path dir, IOException exc) throws IOEx
    */
   public void ansibleInventoryList(NodeSetImpl nodes, AnsibleRunner.AnsibleRunnerBuilder runnerBuilder) throws ResourceModelSourceException {
 
+    int codePointLimit = yamlDataSize * 1024 * 1024;
+
     LoaderOptions snakeOptions = new LoaderOptions();
     // max inventory file size allowed to 10mb
-    snakeOptions.setCodePointLimit(10_485_760);
+    snakeOptions.setCodePointLimit(codePointLimit);
     Yaml yaml = new Yaml(new SafeConstructor(snakeOptions));
 
     String listResp = getNodesFromInventory(runnerBuilder);
 
-    Map<String, Object> allInventory = yaml.load(listResp);
+    Map<String, Object> allInventory;
+    try {
+      allInventory = yaml.load(listResp);
+    } catch (YAMLException e) {
+      throw new ResourceModelSourceException("Cannot load yaml data coming from Ansible: " + e.getMessage(), e);
+    }
+
     Map<String, Object> all = InventoryList.getValue(allInventory, ALL);
     Map<String, Object> children = InventoryList.getValue(all, CHILDREN);
 

src/main/groovy/com/rundeck/plugins/ansible/plugin/AnsibleResourceModelSourceFactory.java

@@ -35,6 +35,7 @@ public AnsibleResourceModelSourceFactory(final Framework framework) {
         builder.property(INVENTORY_PROP);
         builder.property(CONFIG_FILE_PATH);
         builder.property(GATHER_FACTS_PROP);
+        builder.property(YAML_DATA_SIZE_PROP);
         builder.property(IGNORE_ERRORS_PROP);
         builder.property(LIMIT_PROP);
         builder.property(DISABLE_LIMIT_PROP);

src/test/groovy/com/rundeck/plugins/ansible/plugin/AnsibleResourceModelSourceSpec.groovy

@@ -1,20 +1,20 @@
 package com.rundeck.plugins.ansible.plugin
 
-import com.dtolabs.rundeck.core.storage.StorageTree
-import com.dtolabs.rundeck.core.storage.keys.KeyStorageTree
-import org.rundeck.app.spi.Services
-
-import static com.rundeck.plugins.ansible.ansible.AnsibleInventoryList.AnsibleInventoryListBuilder
-
 import com.dtolabs.rundeck.core.common.Framework
 import com.dtolabs.rundeck.core.common.INodeEntry
 import com.dtolabs.rundeck.core.common.INodeSet
 import com.dtolabs.rundeck.core.resources.ResourceModelSource
+import com.dtolabs.rundeck.core.resources.ResourceModelSourceException
+import com.dtolabs.rundeck.core.storage.keys.KeyStorageTree
 import com.rundeck.plugins.ansible.ansible.AnsibleDescribable
 import com.rundeck.plugins.ansible.ansible.AnsibleInventoryList
+import org.rundeck.app.spi.Services
 import org.yaml.snakeyaml.Yaml
+import org.yaml.snakeyaml.error.YAMLException
 import spock.lang.Specification
 
+import static com.rundeck.plugins.ansible.ansible.AnsibleInventoryList.AnsibleInventoryListBuilder
+
 /**
  * AnsibleResourceModelSource test
  */
@@ -94,4 +94,87 @@ class AnsibleResourceModelSourceSpec extends Specification {
         'NODE_3' | 'ansible_os_family' | 'ansible_os_name' | 'ansible_kernel' | 'ansible_architecture' | 'ansible_user_id'  | 'ansible_distribution'
     }
 
+    void "ansible yaml data size parameter without an Exception"() {
+        given:
+        Framework framework = Mock(Framework) {
+            getBaseDir() >> Mock(File) {
+                getAbsolutePath() >> '/tmp'
+            }
+        }
+        ResourceModelSource plugin = new AnsibleResourceModelSource(framework)
+        Properties config = new Properties()
+        config.put('project', 'project_1')
+        config.put(AnsibleDescribable.ANSIBLE_GATHER_FACTS, 'false')
+        plugin.configure(config)
+        Services services = Mock(Services) {
+            getService(KeyStorageTree.class) >> Mock(KeyStorageTree)
+        }
+        plugin.setServices(services)
+        plugin.yamlDataSize = 2
+
+        when: "small inventory"
+        AnsibleInventoryListBuilder inventoryListBuilder = mockInventoryList(qtyNodes)
+        plugin.ansibleInventoryListBuilder = inventoryListBuilder
+        INodeSet nodes = plugin.getNodes()
+
+        then: "non exception is thrown because data can be handled"
+        notThrown(YAMLException)
+        nodes.size() == qtyNodes
+
+        where:
+        qtyNodes | _
+        2_0000   | _
+        3_0000   | _
+    }
+
+    void "ansible yaml data size parameter with an Exception"() {
+        given:
+        Framework framework = Mock(Framework) {
+            getBaseDir() >> Mock(File) {
+                getAbsolutePath() >> '/tmp'
+            }
+        }
+        ResourceModelSource plugin = new AnsibleResourceModelSource(framework)
+        Properties config = new Properties()
+        config.put('project', 'project_1')
+        config.put(AnsibleDescribable.ANSIBLE_GATHER_FACTS, 'false')
+        plugin.configure(config)
+        Services services = Mock(Services) {
+            getService(KeyStorageTree.class) >> Mock(KeyStorageTree)
+        }
+        plugin.setServices(services)
+        plugin.yamlDataSize = 2
+
+        when: "huge inventory"
+        AnsibleInventoryListBuilder inventoryListBuilder = mockInventoryList(100_000)
+        plugin.ansibleInventoryListBuilder = inventoryListBuilder
+        plugin.getNodes()
+
+        then: "throws an exception because data is too big to be precessed"
+        thrown(ResourceModelSourceException)
+    }
+
+    private AnsibleInventoryListBuilder mockInventoryList(int qtyNodes) {
+        return Mock(AnsibleInventoryListBuilder) {
+            build() >> Mock(AnsibleInventoryList) {
+                getNodeList() >> createNodes(qtyNodes)
+            }
+        }
+    }
+
+    private static String createNodes(int qty) {
+        Yaml yaml = new Yaml()
+        Map<String, Object> host = [:]
+        for (int i=1; i <= qty; i++) {
+            String nodeName = "node-$i"
+            String hostValue = "any-name-$i"
+            host.put((nodeName), ['hostname' : (hostValue)])
+        }
+        def hosts = ['hosts' : host]
+        def groups = ['ungrouped' : hosts]
+        def children = ['children' : groups]
+        def all = ['all' : children]
+        return yaml.dump(all)
+    }
+
 }