From 4ffdb445ee9de9b5caa3ab71df49b771cacd361c Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Mathieu=20Gagne=CC=81?= <mgagne@iweb.com>
Date: Wed, 3 Dec 2014 18:36:58 -0500
Subject: [PATCH] Add rabbitmq_policy custom type

Include support for RabbitMQ <3.2.0:
  * applyto field was not available before 3.2.0
  * syntax of the set_policy command changed in 3.2.0
---
 README.md                                     |  16 ++-
 .../provider/rabbitmq_policy/rabbitmqctl.rb   | 125 ++++++++++++++++++
 lib/puppet/provider/rabbitmqctl.rb            |  10 ++
 lib/puppet/type/rabbitmq_policy.rb            |  76 +++++++++++
 .../rabbitmq_policy/rabbitmqctl_spec.rb       | 114 ++++++++++++++++
 spec/unit/puppet/type/rabbitmq_policy_spec.rb |  91 +++++++++++++
 6 files changed, 431 insertions(+), 1 deletion(-)
 create mode 100644 lib/puppet/provider/rabbitmq_policy/rabbitmqctl.rb
 create mode 100644 lib/puppet/provider/rabbitmqctl.rb
 create mode 100644 lib/puppet/type/rabbitmq_policy.rb
 create mode 100644 spec/unit/puppet/provider/rabbitmq_policy/rabbitmqctl_spec.rb
 create mode 100644 spec/unit/puppet/type/rabbitmq_policy_spec.rb

diff --git a/README.md b/README.md
index ebf879083..ab54200fa 100644
--- a/README.md
+++ b/README.md
@@ -271,7 +271,7 @@ be changed to latest.
 ####`package_gpg_key`
 
 RPM package GPG key to import. Uses source method. Should be a URL for Debian/RedHat
-OS family, or a file name for RedHat OS family. 
+OS family, or a file name for RedHat OS family.
 Set to http://www.rabbitmq.com/rabbitmq-signing-key-public.asc by default.
 Note, that `key_content`, if specified, would override this parameter for Debian OS family.
 
@@ -418,6 +418,20 @@ rabbitmq_user_permissions { 'dan@myhost':
 }
 ```
 
+### rabbitmq\_policy
+
+```puppet
+rabbitmq_policy { 'ha-all@myhost':
+  pattern    => '.*',
+  priority   => 0,
+  applyto    => 'all',
+  definition => {
+    'ha-mode'      => 'all',
+    'ha-sync-mode' => 'automatic'
+  }
+}
+```
+
 ### rabbitmq\_plugin
 
 query all currently enabled plugins `$ puppet resource rabbitmq_plugin`
diff --git a/lib/puppet/provider/rabbitmq_policy/rabbitmqctl.rb b/lib/puppet/provider/rabbitmq_policy/rabbitmqctl.rb
new file mode 100644
index 000000000..04f5a780f
--- /dev/null
+++ b/lib/puppet/provider/rabbitmq_policy/rabbitmqctl.rb
@@ -0,0 +1,125 @@
+require 'json'
+require 'puppet/util/package'
+
+require File.expand_path(File.join(File.dirname(__FILE__), '..', 'rabbitmqctl'))
+Puppet::Type.type(:rabbitmq_policy).provide(:rabbitmqctl, :parent => Puppet::Provider::Rabbitmqctl) do
+
+  defaultfor :feature => :posix
+
+  # cache policies
+  def self.policies(name, vhost)
+    @policies = {} unless @policies
+    unless @policies[vhost]
+      @policies[vhost] = {}
+      rabbitmqctl('list_policies', '-q', '-p', vhost).split(/\n/).each do |line|
+        # rabbitmq<3.2 does not support the applyto field
+        # 1 2      3?  4  5                                            6
+        # / ha-all all .* {"ha-mode":"all","ha-sync-mode":"automatic"} 0
+        if line =~ /^(\S+)\s+(\S+)\s+(all|exchanges|queues)?\s*(\S+)\s+(\S+)\s+(\d+)$/
+          applyto = $3 || 'all'
+          @policies[vhost][$2] = {
+            :applyto    => applyto,
+            :pattern    => $4,
+            :definition => JSON.parse($5),
+            :priority   => $6}
+        else
+          raise Puppet::Error, "cannot parse line from list_policies:#{line}"
+        end
+      end
+    end
+    @policies[vhost][name]
+  end
+
+  def policies(name, vhost)
+    self.class.policies(vhost, name)
+  end
+
+  def should_policy
+    if @should_policy
+      @should_policy
+    else
+      @should_policy = resource[:name].rpartition('@').first
+    end
+  end
+
+  def should_vhost
+    if @should_vhost
+      @should_vhost
+    else
+      @should_vhost = resource[:name].rpartition('@').last
+    end
+  end
+
+  def create
+    set_policy
+  end
+
+  def destroy
+    rabbitmqctl('clear_policy', '-p', should_vhost, should_policy)
+  end
+
+  def exists?
+    policies(should_vhost, should_policy)
+  end
+
+  def pattern
+    policies(should_vhost, should_policy)[:pattern]
+  end
+
+  def pattern=(pattern)
+    set_policy
+  end
+
+  def applyto
+    policies(should_vhost, should_policy)[:applyto]
+  end
+
+  def applyto=(applyto)
+    set_policy
+  end
+
+  def definition
+    policies(should_vhost, should_policy)[:definition]
+  end
+
+  def definition=(definition)
+    set_policy
+  end
+
+  def priority
+    policies(should_vhost, should_policy)[:priority]
+  end
+
+  def priority=(priority)
+    set_policy
+  end
+
+  def set_policy
+    unless @set_policy
+      @set_policy = true
+      resource[:applyto]    ||= applyto
+      resource[:definition] ||= definition
+      resource[:pattern]    ||= pattern
+      resource[:priority]   ||= priority
+      # rabbitmq>=3.2.0
+      if Puppet::Util::Package.versioncmp(self.class.rabbitmq_version, '3.2.0') >= 0
+        rabbitmqctl('set_policy',
+          '-p', should_vhost,
+          '--priority', resource[:priority],
+          '--apply-to', resource[:applyto].to_s,
+          should_policy,
+          resource[:pattern],
+          resource[:definition].to_json
+        )
+      else
+        rabbitmqctl('set_policy',
+          '-p', should_vhost,
+          should_policy,
+          resource[:pattern],
+          resource[:definition].to_json,
+          resource[:priority]
+        )
+      end
+    end
+  end
+end
diff --git a/lib/puppet/provider/rabbitmqctl.rb b/lib/puppet/provider/rabbitmqctl.rb
new file mode 100644
index 000000000..4afdfd168
--- /dev/null
+++ b/lib/puppet/provider/rabbitmqctl.rb
@@ -0,0 +1,10 @@
+class Puppet::Provider::Rabbitmqctl < Puppet::Provider
+  initvars
+  commands :rabbitmqctl => 'rabbitmqctl'
+
+  def self.rabbitmq_version
+    output = rabbitmqctl('-q', 'status')
+    version = output.match(/\{rabbit,"RabbitMQ","([\d\.]+)"\}/)
+    version[1] if version
+  end
+end
diff --git a/lib/puppet/type/rabbitmq_policy.rb b/lib/puppet/type/rabbitmq_policy.rb
new file mode 100644
index 000000000..c147c8039
--- /dev/null
+++ b/lib/puppet/type/rabbitmq_policy.rb
@@ -0,0 +1,76 @@
+Puppet::Type.newtype(:rabbitmq_policy) do
+  desc 'Type for managing rabbitmq policies'
+
+  ensurable do
+    defaultto(:present)
+    newvalue(:present) do
+      provider.create
+    end
+    newvalue(:absent) do
+      provider.destroy
+    end
+  end
+
+  autorequire(:service) { 'rabbitmq-server' }
+
+  validate do
+    fail('pattern parameter is required.') if self[:ensure] == :present and self[:pattern].nil?
+    fail('definition parameter is required.') if self[:ensure] == :present and self[:definition].nil?
+  end
+
+  newparam(:name, :namevar => true) do
+    desc 'combination of policy@vhost to create policy for'
+    newvalues(/^\S+@\S+$/)
+  end
+
+  newproperty(:pattern) do
+    desc 'policy pattern'
+    validate do |value|
+      resource.validate_pattern(value)
+    end
+  end
+
+  newproperty(:applyto) do
+    desc 'policy apply to'
+    newvalue(:all)
+    newvalue(:exchanges)
+    newvalue(:queues)
+    defaultto :all
+  end
+
+  newproperty(:definition) do
+    desc 'policy definition'
+    validate do |value|
+      resource.validate_definition(value)
+    end
+  end
+
+  newproperty(:priority) do
+    desc 'policy priority'
+    newvalues(/^\d+$/)
+    defaultto 0
+  end
+
+  autorequire(:rabbitmq_vhost) do
+    [self[:name].split('@')[1]]
+  end
+
+  def validate_pattern(value)
+    begin
+      Regexp.new(value)
+    rescue RegexpError
+      raise ArgumentError, "Invalid regexp #{value}"
+    end
+  end
+
+  def validate_definition(definition)
+    unless [Hash].include?(definition.class)
+      raise ArgumentError, "Invalid definition"
+    end
+    definition.each do |k,v|
+      unless [String].include?(v.class)
+        raise ArgumentError, "Invalid definition"
+      end
+    end
+  end
+end
diff --git a/spec/unit/puppet/provider/rabbitmq_policy/rabbitmqctl_spec.rb b/spec/unit/puppet/provider/rabbitmq_policy/rabbitmqctl_spec.rb
new file mode 100644
index 000000000..cddb6c0c0
--- /dev/null
+++ b/spec/unit/puppet/provider/rabbitmq_policy/rabbitmqctl_spec.rb
@@ -0,0 +1,114 @@
+require 'puppet'
+require 'mocha'
+
+RSpec.configure do |config|
+  config.mock_with :mocha
+end
+
+describe Puppet::Type.type(:rabbitmq_policy).provider(:rabbitmqctl) do
+
+  let(:resource) do
+    Puppet::Type.type(:rabbitmq_policy).new(
+      :name       => 'ha-all@/',
+      :pattern    => '.*',
+      :definition => {
+        'ha-mode' => 'all'
+      },
+      :provider => described_class.name
+    )
+  end
+
+  let(:provider) { resource.provider }
+
+  after(:each) do
+    described_class.instance_variable_set(:@policies, nil)
+  end
+
+  it 'should accept @ in policy name' do
+    resource = Puppet::Type.type(:rabbitmq_policy).new(
+      :name       => 'ha@home@/',
+      :pattern    => '.*',
+      :definition => {
+        'ha-mode' => 'all'
+      },
+      :provider => described_class.name
+    )
+    provider = described_class.new(resource)
+    provider.should_policy.should == 'ha@home'
+    provider.should_vhost.should == '/'
+  end
+
+  it 'should fail with invalid output from list' do
+    provider.class.expects(:rabbitmqctl).with('list_policies', '-q', '-p', '/').returns 'foobar'
+    expect { provider.exists? }.to raise_error(Puppet::Error, /cannot parse line from list_policies/)
+  end
+
+  it 'should match policies from list (>=3.2.0)' do
+    provider.class.expects(:rabbitmqctl).with('list_policies', '-q', '-p', '/').returns <<-EOT
+/ ha-all all .* {"ha-mode":"all","ha-sync-mode":"automatic"} 0
+/ test exchanges .* {"ha-mode":"all"} 0
+EOT
+    provider.exists?.should == {
+      :applyto    => 'all',
+      :pattern    => '.*',
+      :priority   => '0',
+      :definition => {
+        'ha-mode'      => 'all',
+        'ha-sync-mode' => 'automatic'}
+      }
+  end
+
+  it 'should match policies from list (<3.2.0)' do
+    provider.class.expects(:rabbitmqctl).with('list_policies', '-q', '-p', '/').returns <<-EOT
+/ ha-all .* {"ha-mode":"all","ha-sync-mode":"automatic"} 0
+/ test .* {"ha-mode":"all"} 0
+EOT
+    provider.exists?.should == {
+      :applyto    => 'all',
+      :pattern    => '.*',
+      :priority   => '0',
+      :definition => {
+        'ha-mode'      => 'all',
+        'ha-sync-mode' => 'automatic'}
+      }
+  end
+
+  it 'should not match an empty list' do
+    provider.class.expects(:rabbitmqctl).with('list_policies', '-q', '-p', '/').returns ''
+    provider.exists?.should == nil
+  end
+
+  it 'should destroy policy' do
+    provider.expects(:rabbitmqctl).with('clear_policy', '-p', '/', 'ha-all')
+    provider.destroy
+  end
+
+  it 'should only call set_policy once (<3.2.0)' do
+    provider.class.expects(:rabbitmq_version).returns '3.1.0'
+    provider.resource[:priority] = '10'
+    provider.resource[:applyto] = 'exchanges'
+    provider.expects(:rabbitmqctl).with('set_policy',
+      '-p', '/',
+      'ha-all',
+      '.*',
+      '{"ha-mode":"all"}',
+      '10').once
+    provider.priority = '10'
+    provider.applyto = 'exchanges'
+  end
+
+  it 'should only call set_policy once (>=3.2.0)' do
+    provider.class.expects(:rabbitmq_version).returns '3.2.0'
+    provider.resource[:priority] = '10'
+    provider.resource[:applyto] = 'exchanges'
+    provider.expects(:rabbitmqctl).with('set_policy',
+      '-p', '/',
+      '--priority', '10',
+      '--apply-to', 'exchanges',
+      'ha-all',
+      '.*',
+      '{"ha-mode":"all"}').once
+    provider.priority = '10'
+    provider.applyto = 'exchanges'
+  end
+end
diff --git a/spec/unit/puppet/type/rabbitmq_policy_spec.rb b/spec/unit/puppet/type/rabbitmq_policy_spec.rb
new file mode 100644
index 000000000..2a8064a1a
--- /dev/null
+++ b/spec/unit/puppet/type/rabbitmq_policy_spec.rb
@@ -0,0 +1,91 @@
+require 'puppet'
+require 'puppet/type/rabbitmq_policy'
+
+describe Puppet::Type.type(:rabbitmq_policy) do
+
+  before do
+    @policy = Puppet::Type.type(:rabbitmq_policy).new(
+      :name       => 'ha-all@/',
+      :pattern    => '.*',
+      :definition => {
+        'ha-mode' => 'all'
+      })
+  end
+
+  it 'should accept a valid name' do
+    @policy[:name] = 'ha-all@/'
+    @policy[:name].should == 'ha-all@/'
+  end
+
+  it 'should require a name' do
+    expect {
+      Puppet::Type.type(:rabbitmq_policy).new({})
+    }.to raise_error(Puppet::Error, 'Title or name must be provided')
+  end
+
+  it 'should fail when name does not have a @' do
+    expect {
+      @policy[:name] = 'ha-all'
+    }.to raise_error(Puppet::Error, /Valid values match/)
+  end
+
+  it 'should accept a valid regex for pattern' do
+    @policy[:pattern] = '.*?'
+    @policy[:pattern].should == '.*?'
+  end
+
+  it 'should accept an empty string for pattern' do
+    @policy[:pattern] = ''
+    @policy[:pattern].should == ''
+  end
+
+  it 'should not accept invalid regex for pattern' do
+    expect {
+      @policy[:pattern] = '*'
+    }.to raise_error(Puppet::Error, /Invalid regexp/)
+  end
+
+  it 'should accept valid value for applyto' do
+    [:all, :exchanges, :queues].each do |v|
+      @policy[:applyto] = v
+      @policy[:applyto].should == v
+    end
+  end
+
+  it 'should not accept invalid value for applyto' do
+    expect {
+      @policy[:applyto] = 'me'
+    }.to raise_error(Puppet::Error, /Invalid value/)
+  end
+
+  it 'should accept a valid hash for definition' do
+    definition = {'ha-mode' => 'all', 'ha-sync-mode' => 'automatic'}
+    @policy[:definition] = definition
+    @policy[:definition].should == definition
+  end
+
+  it 'should not accept invalid hash for definition' do
+    expect {
+      @policy[:definition] = 'ha-mode'
+    }.to raise_error(Puppet::Error, /Invalid definition/)
+
+    expect {
+      @policy[:definition] = {'ha-mode' => ['a', 'b']}
+    }.to raise_error(Puppet::Error, /Invalid definition/)
+  end
+
+  it 'should accept valid value for priority' do
+    [0, 10, '0', '10'].each do |v|
+      @policy[:priority] = v
+      @policy[:priority].should == v
+    end
+  end
+
+  it 'should not accept invalid value for priority' do
+    ['-1', -1, '1.0', 1.0, 'abc', ''].each do |v|
+      expect {
+        @policy[:priority] = v
+      }.to raise_error(Puppet::Error, /Invalid value/)
+    end
+  end
+end