Merge pull request #22 from adammck/counted_instances

Add support for grouped count=n resources

README.md

@@ -80,11 +80,15 @@ It's just a Go app, so the usual:
 To test against an example statefile, run:
 
 	terraform-inventory --list fixtures/example.tfstate
-	terraform-inventory --host=web-aws fixtures/example.tfstate
+	terraform-inventory --host=52.7.58.202 fixtures/example.tfstate
 
 To update the fixtures, populate `fixtures/secrets.tfvars` with your DO and AWS
-account details, and run `fixtures/update`. You almost certainly don't need to
-do this.
+account details, and run `fixtures/update`. To run a tiny Ansible playbook on
+the example resourecs, run:
+
+	TF_STATE=fixtures/example.tfstate ansible-playbook --inventory-file=terraform-inventory fixtures/playbook.yml
+
+You almost certainly don't need to do any of this. Use the tests instead.
 
 
 ## License

cli.go

@@ -9,18 +9,28 @@ import (
 func cmdList(stdout io.Writer, stderr io.Writer, s *state) int {
 	groups := make(map[string][]string, 0)
 
-	// add each instance as a pseudo-group, so they can be provisioned
+	// Add each instance name as a pseudo-group, so they can be provisioned
 	// individually where necessary.
-	for name, res := range s.resources() {
-		groups[name] = []string{res.Address()}
+	for _, res := range s.resources() {
+		_, ok := groups[res.Name]
+		if !ok {
+			groups[res.Name] = []string{}
+		}
+
+		// Add the instance by name. There can be many instances with the same name,
+		// created using the count parameter.
+		groups[res.Name] = append(groups[res.Name], res.Address())
+
+		// Add the instance by its full name, including the counter.
+		groups[res.NameWithCounter()] = []string{res.Address()}
 	}
 
 	return output(stdout, stderr, groups)
 }
 
 func cmdHost(stdout io.Writer, stderr io.Writer, s *state, hostname string) int {
-	for name, res := range s.resources() {
-		if hostname == name {
+	for _, res := range s.resources() {
+		if hostname == res.Address() {
 			return output(stdout, stderr, res.Attributes())
 		}
 	}

fixtures/example.tf

@@ -29,6 +29,7 @@ resource "aws_instance" "web-aws" {
   subnet_id = "${var.aws_subnet_id}"
   associate_public_ip_address = true
   key_name = "terraform-inventory"
+  count = 2
   root_block_device = {
     delete_on_termination = true
   }

fixtures/example.tfstate

@@ -8,7 +8,7 @@
             ],
             "outputs": {},
             "resources": {
-                "aws_instance.web-aws": {
+                "aws_instance.web-aws.0": {
                     "type": "aws_instance",
                     "primary": {
                         "id": "i-e8f3a238",
@@ -32,6 +32,45 @@
                             "root_block_device.0.volume_size": "8",
                             "root_block_device.0.volume_type": "standard",
                             "security_groups.#": "0",
+                            "source_dest_check": "true",
+                            "subnet_id": "subnet-59f9b32e",
+                            "tags.#": "0",
+                            "tenancy": "default",
+                            "vpc_security_group_ids.#": "1",
+                            "vpc_security_group_ids.2076429742": "sg-b42329d0"
+                        },
+                        "meta": {
+                            "schema_version": "1"
+                        }
+                    }
+                },
+                "aws_instance.web-aws.1": {
+                    "type": "aws_instance",
+                    "primary": {
+                        "id": "i-f747c141",
+                        "attributes": {
+                            "ami": "ami-96a818fe",
+                            "associate_public_ip_address": "true",
+                            "availability_zone": "us-east-1d",
+                            "ebs_block_device.#": "0",
+                            "ebs_optimized": "false",
+                            "ephemeral_block_device.#": "0",
+                            "iam_instance_profile": "",
+                            "id": "i-f747c141",
+                            "instance_type": "t2.micro",
+                            "key_name": "terraform-inventory",
+                            "monitoring": "false",
+                            "private_dns": "ip-10-0-0-10.ec2.internal",
+                            "private_ip": "10.0.0.10",
+                            "public_dns": "",
+                            "public_ip": "52.91.51.56",
+                            "root_block_device.#": "1",
+                            "root_block_device.0.delete_on_termination": "true",
+                            "root_block_device.0.iops": "0",
+                            "root_block_device.0.volume_size": "8",
+                            "root_block_device.0.volume_type": "standard",
+                            "security_groups.#": "0",
+                            "source_dest_check": "true",
                             "subnet_id": "subnet-59f9b32e",
                             "tags.#": "0",
                             "tenancy": "default",

parser.go

@@ -2,9 +2,11 @@ package main
 
 import (
 	"encoding/json"
+	"fmt"
 	"io"
 	"io/ioutil"
 	"regexp"
+	"strconv"
 )
 
 type state struct {
@@ -15,6 +17,7 @@ type state struct {
 // state file. This allows us to support multiple types of resource without too
 // much fuss.
 var keyNames []string
+var nameParser *regexp.Regexp
 
 func init() {
 	keyNames = []string{
@@ -26,6 +29,9 @@ func init() {
 		"access_ip_v4", // OpenStack
 		"floating_ip",  // OpenStack
 	}
+
+	// type.name.0
+	nameParser = regexp.MustCompile(`^(\w+)\.([\w\-]+)(?:\.(\d+))?$`)
 }
 
 // read populates the state object from a statefile.
@@ -49,35 +55,73 @@ func (s *state) read(stateFile io.Reader) error {
 // resources returns a map of name to resourceState, for any supported resources
 // found in the statefile.
 func (s *state) resources() map[string]resourceState {
-	typeRemover := regexp.MustCompile(`^[\w_]+\.`)
 	inst := make(map[string]resourceState)
 
 	for _, m := range s.Modules {
 		for k, r := range m.Resources {
 			if r.isSupported() {
-				name := typeRemover.ReplaceAllString(k, "")
-				inst[name] = r
+
+				_, name, counter := parseName(k)
+				//fmt.Println(resType, name, counter)
+				r.Name = name
+				r.Counter = counter
+				inst[k] = r
 			}
 		}
 	}
 
 	return inst
 }
 
+func parseName(name string) (string, string, int) {
+	m := nameParser.FindStringSubmatch(name)
+
+	// This should not happen unless our regex changes.
+	// TODO: Warn instead of silently ignore error?
+	if len(m) != 4 {
+		//fmt.Printf("len=%d\n", len(m))
+		return "", "", 0
+	}
+
+	var c int
+	var err error
+	if m[3] != "" {
+		c, err = strconv.Atoi(m[3])
+		if err != nil {
+			fmt.Printf("err: %s\n", err)
+			// ???
+		}
+	}
+
+	return m[1], m[2], c
+}
+
 type moduleState struct {
 	Resources map[string]resourceState `json:"resources"`
 }
 
 type resourceState struct {
+
+	// Populated from statefile
 	Type    string        `json:"type"`
 	Primary instanceState `json:"primary"`
+
+	// Extracted from key name, and injected by resources method
+	Name    string
+	Counter int
 }
 
 // isSupported returns true if terraform-inventory supports this resource.
 func (s resourceState) isSupported() bool {
 	return s.Address() != ""
 }
 
+// NameWithCounter returns the resource name with its counter. For resources
+// created without a 'count=' attribute, this will always be zero.
+func (s resourceState) NameWithCounter() string {
+	return fmt.Sprintf("%s.%d", s.Name, s.Counter)
+}
+
 // Address returns the IP address of this resource.
 func (s resourceState) Address() string {
 	for _, key := range keyNames {

parser_test.go

@@ -1,9 +1,12 @@
 package main
 
 import (
-	"github.com/stretchr/testify/assert"
+	"bytes"
+	"encoding/json"
 	"strings"
 	"testing"
+
+	"github.com/stretchr/testify/assert"
 )
 
 const exampleStateFile = `
@@ -17,7 +20,7 @@ const exampleStateFile = `
 			],
 			"outputs": {},
 			"resources": {
-				"aws_instance.one": {
+				"aws_instance.one.0": {
 					"type": "aws_instance",
 					"primary": {
 						"id": "i-aaaaaaaa",
@@ -27,6 +30,16 @@ const exampleStateFile = `
 						}
 					}
 				},
+				"aws_instance.one.1": {
+					"type": "aws_instance",
+					"primary": {
+						"id": "i-a1a1a1a1",
+						"attributes": {
+							"id": "i-a1a1a1a1",
+							"private_ip": "10.0.1.1"
+						}
+					}
+				},
 				"aws_instance.two": {
 					"type": "aws_instance",
 					"primary": {
@@ -101,12 +114,75 @@ const exampleStateFile = `
 }
 `
 
+const expectedListOutput = `
+{
+	"one":   ["10.0.0.1", "10.0.1.1"],
+	"two":   ["50.0.0.1"],
+	"three": ["192.168.0.3"],
+	"four":  ["10.2.1.5"],
+	"five":  ["10.20.30.40"],
+	"six":   ["10.120.0.226"],
+
+	"one.0":   ["10.0.0.1"],
+	"one.1":   ["10.0.1.1"],
+	"two.0":   ["50.0.0.1"],
+	"three.0": ["192.168.0.3"],
+	"four.0":  ["10.2.1.5"],
+	"five.0":  ["10.20.30.40"],
+	"six.0":   ["10.120.0.226"]
+}
+`
+
+const expectedHostOneOutput = `
+{
+	"id":"i-aaaaaaaa",
+	"private_ip":"10.0.0.1"
+}
+`
+
+func TestListCommand(t *testing.T) {
+	var s state
+	r := strings.NewReader(exampleStateFile)
+	err := s.read(r)
+	assert.Nil(t, err)
+
+	// Run the command, capture the output
+	var stdout, stderr bytes.Buffer
+	exitCode := cmdList(&stdout, &stderr, &s)
+	assert.Equal(t, 0, exitCode)
+	assert.Equal(t, "", stderr.String())
+
+	var exp, act interface{}
+	json.Unmarshal([]byte(expectedListOutput), &exp)
+	json.Unmarshal([]byte(stdout.String()), &act)
+	assert.Equal(t, exp, act)
+}
+
+func TestHostCommand(t *testing.T) {
+	var s state
+	r := strings.NewReader(exampleStateFile)
+	err := s.read(r)
+	assert.Nil(t, err)
+
+	// Run the command, capture the output
+	var stdout, stderr bytes.Buffer
+	exitCode := cmdHost(&stdout, &stderr, &s, "10.0.0.1")
+	assert.Equal(t, 0, exitCode)
+	assert.Equal(t, "", stderr.String())
+
+	var exp, act interface{}
+	json.Unmarshal([]byte(expectedHostOneOutput), &exp)
+	json.Unmarshal([]byte(stdout.String()), &act)
+	assert.Equal(t, exp, act)
+}
+
 func TestStateRead(t *testing.T) {
 	var s state
 	r := strings.NewReader(exampleStateFile)
 	err := s.read(r)
 	assert.Nil(t, err)
-	assert.Equal(t, "aws_instance", s.Modules[0].Resources["aws_instance.one"].Type)
+	assert.Equal(t, "aws_instance", s.Modules[0].Resources["aws_instance.one.0"].Type)
+	assert.Equal(t, "aws_instance", s.Modules[0].Resources["aws_instance.two"].Type)
 }
 
 func TestResources(t *testing.T) {
@@ -117,13 +193,14 @@ func TestResources(t *testing.T) {
 	assert.Nil(t, err)
 
 	inst := s.resources()
-	assert.Equal(t, 6, len(inst))
-	assert.Equal(t, "aws_instance", inst["one"].Type)
-	assert.Equal(t, "aws_instance", inst["two"].Type)
-	assert.Equal(t, "digitalocean_droplet", inst["three"].Type)
-	assert.Equal(t, "cloudstack_instance", inst["four"].Type)
-	assert.Equal(t, "vsphere_virtual_machine", inst["five"].Type)
-	assert.Equal(t, "openstack_compute_instance_v2", inst["six"].Type)
+	assert.Equal(t, 7, len(inst))
+	assert.Equal(t, "aws_instance", inst["aws_instance.one.0"].Type)
+	assert.Equal(t, "aws_instance", inst["aws_instance.one.1"].Type)
+	assert.Equal(t, "aws_instance", inst["aws_instance.two"].Type)
+	assert.Equal(t, "digitalocean_droplet", inst["digitalocean_droplet.three"].Type)
+	assert.Equal(t, "cloudstack_instance", inst["cloudstack_instance.four"].Type)
+	assert.Equal(t, "vsphere_virtual_machine", inst["vsphere_virtual_machine.five"].Type)
+	assert.Equal(t, "openstack_compute_instance_v2", inst["openstack_compute_instance_v2.six"].Type)
 }
 
 func TestAddress(t *testing.T) {
@@ -134,13 +211,14 @@ func TestAddress(t *testing.T) {
 	assert.Nil(t, err)
 
 	inst := s.resources()
-	assert.Equal(t, 6, len(inst))
-	assert.Equal(t, "10.0.0.1", inst["one"].Address())
-	assert.Equal(t, "50.0.0.1", inst["two"].Address())
-	assert.Equal(t, "192.168.0.3", inst["three"].Address())
-	assert.Equal(t, "10.2.1.5", inst["four"].Address())
-	assert.Equal(t, "10.20.30.40", inst["five"].Address())
-	assert.Equal(t, "10.120.0.226", inst["six"].Address())
+	assert.Equal(t, 7, len(inst))
+	assert.Equal(t, "10.0.0.1", inst["aws_instance.one.0"].Address())
+	assert.Equal(t, "10.0.1.1", inst["aws_instance.one.1"].Address())
+	assert.Equal(t, "50.0.0.1", inst["aws_instance.two"].Address())
+	assert.Equal(t, "192.168.0.3", inst["digitalocean_droplet.three"].Address())
+	assert.Equal(t, "10.2.1.5", inst["cloudstack_instance.four"].Address())
+	assert.Equal(t, "10.20.30.40", inst["vsphere_virtual_machine.five"].Address())
+	assert.Equal(t, "10.120.0.226", inst["openstack_compute_instance_v2.six"].Address())
 }
 
 func TestIsSupported(t *testing.T) {