Add onlyif parameter to postgresql_psql to only run command if onlyif returns true

Add option to only attempt table grant if table already exists

Make this slightly more generic

Author: kimor79

lib/puppet/type/postgresql_psql.rb

@@ -43,6 +43,28 @@ def matches(value)
     end
   end
 
+  newparam(:onlyif) do
+    desc "An optional SQL command to execute prior to the main :command; " +
+        "this is generally intended to be used for idempotency, to check " +
+        "for the existence of an object in the database to determine whether " +
+        "or not the main SQL command needs to be executed at all."
+
+    # Return true if a matching row is found
+    def matches(value)
+      if Puppet::PUPPETVERSION.to_f < 4
+        output, status = provider.run_unless_sql_command(value)
+      else
+        output = provider.run_unless_sql_command(value)
+        status = output.exitcode
+      end
+      self.fail("Error evaluating 'onlyif' clause, returned #{status}: '#{output}'") unless status == 0
+
+      result_count = output.strip.to_i
+      self.debug("Found #{result_count} row(s) executing 'onlyif' clause")
+      result_count > 0
+    end
+  end
+
   newparam(:db) do
     desc "The name of the database to execute the SQL command against."
   end
@@ -97,7 +119,9 @@ def matches(value)
   end
 
   def should_run_sql(refreshing = false)
+    onlyif_param = @parameters[:onlyif]
     unless_param = @parameters[:unless]
+    return false if !onlyif_param.nil? && !onlyif_param.value.nil? && !onlyif_param.matches(onlyif_param.value)
     return false if !unless_param.nil? && !unless_param.value.nil? && unless_param.matches(unless_param.value)
     return false if !refreshing && @parameters[:refreshonly].value == :true
     true

manifests/server/grant.pp

@@ -2,12 +2,13 @@
 define postgresql::server::grant (
   $role,
   $db,
-  $privilege   = undef,
-  $object_type = 'database',
-  $object_name = undef,
-  $psql_db     = $postgresql::server::default_database,
-  $psql_user   = $postgresql::server::user,
-  $port        = $postgresql::server::port
+  $privilege     = undef,
+  $object_type   = 'database',
+  $object_name   = undef,
+  $psql_db       = $postgresql::server::default_database,
+  $psql_user     = $postgresql::server::user,
+  $port          = $postgresql::server::port,
+  $onlyif_exists = false,
 ) {
   $group     = $postgresql::server::group
   $psql_path = $postgresql::server::psql_path
@@ -18,6 +19,8 @@
     $_object_name = $object_name
   }
 
+  validate_bool($onlyif_exists)
+
   ## Munge the input values
   $_object_type = upcase($object_type)
   $_privilege   = upcase($privilege)
@@ -59,6 +62,7 @@
         'ALL','ALL PRIVILEGES')
       $unless_function = 'has_database_privilege'
       $on_db = $psql_db
+      $onlyif_function = undef
     }
     'SCHEMA': {
       $unless_privilege = $_privilege ? {
@@ -69,6 +73,7 @@
       validate_string($_privilege, 'CREATE', 'USAGE', 'ALL', 'ALL PRIVILEGES')
       $unless_function = 'has_schema_privilege'
       $on_db = $db
+      $onlyif_function = undef
     }
     'TABLE': {
       $unless_privilege = $_privilege ? {
@@ -79,12 +84,17 @@
         'TRUNCATE','REFERENCES','TRIGGER','ALL','ALL PRIVILEGES')
       $unless_function = 'has_table_privilege'
       $on_db = $db
+      $onlyif_function = $onlyif_exists ? {
+        true    => 'table_exists',
+        default => undef,
+      }
     }
     'ALL TABLES IN SCHEMA': {
       validate_string($_privilege,'SELECT','INSERT','UPDATE','DELETE',
         'TRUNCATE','REFERENCES','TRIGGER','ALL','ALL PRIVILEGES')
       $unless_function = 'custom'
       $on_db = $db
+      $onlyif_function = undef
 
       $schema = $object_name
 
@@ -150,6 +160,11 @@
                   '${_granted_object}', '${unless_privilege}')",
   }
 
+  $_onlyif = $onlyif_function ? {
+    'table_exists' => "SELECT true FROM pg_tables WHERE tablename = '${_togrant_object}'",
+    default        => undef,
+  }
+
   $grant_cmd = "GRANT ${_privilege} ON ${_object_type} \"${_togrant_object}\" TO
       \"${role}\""
   postgresql_psql { "grant:${name}":
@@ -160,6 +175,7 @@
     psql_group => $group,
     psql_path  => $psql_path,
     unless     => $_unless,
+    onlyif     => $_onlyif,
     require    => Class['postgresql::server']
   }
 

manifests/server/table_grant.pp

@@ -5,18 +5,20 @@
   $table,
   $db,
   $role,
-  $port      = $postgresql::server::port,
-  $psql_db   = undef,
-  $psql_user = undef
+  $port          = $postgresql::server::port,
+  $psql_db       = undef,
+  $psql_user     = undef,
+  $onlyif_exists = false,
 ) {
   postgresql::server::grant { "table:${name}":
-    role        => $role,
-    db          => $db,
-    port        => $port,
-    privilege   => $privilege,
-    object_type => 'TABLE',
-    object_name => $table,
-    psql_db     => $psql_db,
-    psql_user   => $psql_user,
+    role          => $role,
+    db            => $db,
+    port          => $port,
+    privilege     => $privilege,
+    object_type   => 'TABLE',
+    object_name   => $table,
+    psql_db       => $psql_db,
+    psql_user     => $psql_user,
+    onlyif_exists => $onlyif_exists,
   }
 }

spec/acceptance/postgresql_psql_spec.rb

@@ -114,6 +114,36 @@ class { 'postgresql::server': } ->
     end
   end
 
+  it 'should not run some SQL when the onlyif query returns no rows' do
+    pp = <<-EOS
+      class { 'postgresql::server': } ->
+      postgresql_psql { 'foobar':
+        db        => 'postgres',
+        psql_user => 'postgres',
+        command   => 'select 1',
+        onlyif    => 'select 1 where 1=2',
+      }
+    EOS
+
+    apply_manifest(pp, :catch_failures => true)
+    apply_manifest(pp, :expect_changes  => true)
+  end
+
+  it 'should run SQL when the onlyif query returns rows' do
+    pp = <<-EOS
+      class { 'postgresql::server': } ->
+      postgresql_psql { 'foobar':
+        db        => 'postgres',
+        psql_user => 'postgres',
+        command   => 'select * from pg_database limit 1',
+        onlfy    => 'select 1 where 1=1',
+      }
+    EOS
+
+    apply_manifest(pp, :catch_failures => true)
+    apply_manifest(pp, :catch_changes => true)
+  end
+
   context 'with secure password passing by environment' do
     it 'should run SQL that contanins password passed by environment' do
       select = "select \\'$PASS_TO_EMBED\\'"