(PE-39789) Support custom environment

davidmalloncares · davidmalloncares · commit 21e0c8c2dda8 · 2025-11-10T10:11:14.000Z
diff --git a/functions/get_node_group_environment.pp b/functions/get_node_group_environment.pp
@@ -0,0 +1,31 @@
+#
+# @summary check if a custom PE environment is set in pe.conf
+#
+# @param primary the FQDN for the primary, here we will read the pe.conf from
+#
+# @return [String] the desired environment for PE specific node groups
+#
+# @see https://www.puppet.com/docs/pe/latest/upgrade_pe#update_environment
+#
+function peadm::get_node_group_environment(Peadm::SingleTargetSpec $primary) {
+  $peconf = peadm::get_pe_conf(get_target($primary))
+  # if both are set, they need to be set to the same value
+  # if they are not set, we assume that the user runs their infra in production
+  $pe_install = $peconf['pe_install::install::classification::pe_node_group_environment']
+  $puppet_enterprise = $peconf['puppet_enterprise::master::recover_configuration::pe_environment']
+
+  # check if both are equal
+  # This also evaluates to true if both are undef
+  if $pe_install == $puppet_enterprise {
+    # check if the option isn't undef
+    # ToDo: A proper regex for allowed characters in an environment would be nice
+    # https://github.com/puppetlabs/puppet-docs/issues/1158
+    if $pe_install =~ String[1] {
+      return $pe_install
+    } else {
+      return 'production'
+    }
+  } else {
+    fail("pe_install::install::classification::pe_node_group_environment and puppet_enterprise::master::recover_configuration::pe_environment need to be set to the same value, not '${pe_install}' and '${puppet_enterprise}'")
+  }
+}
diff --git a/manifests/setup/node_manager.pp b/manifests/setup/node_manager.pp
@@ -23,6 +23,7 @@
 #   A load balancer address directing traffic to any of the "B" pool
 #   compilers. This is used for DR configuration in large and extra large
 #   architectures.
+# @param node_group_environment the environment that will be assigned to all the PE Infra node groups
 #
 class peadm::setup::node_manager (
   String[1] $primary_host,
@@ -36,6 +37,7 @@
   Optional[String[1]] $compiler_pool_address            = undef,
   Optional[String[1]] $internal_compiler_a_pool_address = $server_a_host,
   Optional[String[1]] $internal_compiler_b_pool_address = $server_b_host,
+  String[1]           $node_group_environment           = 'production',
 ) {
   # "Not-configured" placeholder string. This will be used in places where we
   # cannot set an explicit null, and need to supply some kind of value.
@@ -46,6 +48,7 @@
   # else.
   Node_group {
     purge_behavior => none,
+    environment => $node_group_environment,
   }
 
   ##################################################
diff --git a/plans/add_database.pp b/plans/add_database.pp
@@ -1,3 +1,6 @@
+#
+# @param node_group_environment environment for the PEADM specific node groups, if not set it will be gathered from pe.conf or production
+#
 plan peadm::add_database(
   Peadm::SingleTargetSpec $targets,
   Peadm::SingleTargetSpec $primary_host,
@@ -10,6 +13,7 @@
       'cleanup-db',
   'finalize']] $begin_at_step = undef,
   Optional[Boolean] $is_migration = false,
+  String[1] $node_group_environment = peadm::get_node_group_environment($primary_host),
 ) {
   $primary_target = peadm::get_targets($primary_host, 1)
   $postgresql_target = peadm::get_targets($targets, 1)
@@ -98,7 +102,7 @@
     run_plan('peadm::subplans::component_install', $postgresql_target,
       primary_host       => $primary_target,
       avail_group_letter => $avail_group_letter,
-      role               => 'puppet/puppetdb-database'
+      role               => 'puppet/puppetdb-database',
     )
   }
 
@@ -137,13 +141,15 @@
       run_plan('peadm::util::update_classification', $primary_target,
         postgresql_a_host => $host,
         postgresql_b_host => $host,
-        peadm_config      => $peadm_config
+        peadm_config      => $peadm_config,
+        node_group_environment => $node_group_environment,
       )
     } else {
       run_plan('peadm::util::update_classification', $primary_target,
         postgresql_a_host => $avail_group_letter ? { 'A' => $postgresql_host, default => undef },
         postgresql_b_host => $avail_group_letter ? { 'B' => $postgresql_host, default => undef },
-        peadm_config      => $peadm_config
+        peadm_config      => $peadm_config,
+        node_group_environment => $node_group_environment,
       )
     }
   }
diff --git a/plans/add_replica.pp b/plans/add_replica.pp
@@ -7,6 +7,8 @@
 # @param replica_postgresql_host - The hostname and certname of the host with the replica PE-PosgreSQL database.
 #   Can be a separate host in an XL architecture, or undef in Standard or Large.
 # @param token_file - (optional) the token file in a different location than the default.
+# @param node_group_environment environment for the PEADM specific node groups, if not set it will be gathered from pe.conf or production
+#
 plan peadm::add_replica(
   # Standard or Large
   Peadm::SingleTargetSpec           $primary_host,
@@ -17,6 +19,7 @@
 
   # Common Configuration
   Optional[String] $token_file = undef,
+  String[1] $node_group_environment = peadm::get_node_group_environment($primary_host),
 ) {
   $primary_target             = peadm::get_targets($primary_host, 1)
   $replica_target             = peadm::get_targets($replica_host, 1)
@@ -96,7 +99,8 @@
     server_b_host                    => $replica_avail_group_letter ? { 'B' => $replica_target.peadm::certname(), default => undef },
     internal_compiler_a_pool_address => $replica_avail_group_letter ? { 'A' => $replica_target.peadm::certname(), default => undef },
     internal_compiler_b_pool_address => $replica_avail_group_letter ? { 'B' => $replica_target.peadm::certname(), default => undef },
-    peadm_config                     => $peadm_config
+    peadm_config                     => $peadm_config,
+    node_group_environment => $node_group_environment,
   )
 
   # Source list of files on Primary and synchronize to new Replica
diff --git a/plans/convert.pp b/plans/convert.pp
@@ -5,6 +5,7 @@
 # management using PEAdm.
 #
 # @param begin_at_step The step where the plan should start. If not set, it will start at the beginning
+# @param node_group_environment environment for the PEADM specific node groups, if not set it will be gathered from pe.conf or production
 #
 plan peadm::convert (
   # Standard
@@ -26,6 +27,7 @@
   Array[String]                     $dns_alt_names                    = [],
 
   Optional[Peadm::ConvertSteps] $begin_at_step = undef,
+  String[1] $node_group_environment = peadm::get_node_group_environment($primary_host),
 ) {
   peadm::assert_supported_bolt_version()
 
@@ -316,6 +318,7 @@
           compiler_pool_address            => $compiler_pool_address,
           internal_compiler_a_pool_address => $internal_compiler_a_pool_address,
           internal_compiler_b_pool_address => $internal_compiler_b_pool_address,
+          node_group_environment           => $node_group_environment,
           require                          => Class['peadm::setup::node_manager_yaml'],
         }
 
diff --git a/plans/install.pp b/plans/install.pp
@@ -30,6 +30,7 @@
 # @param uploaddir
 #   Directory the installer tarball will be uploaded to or expected to be in
 #   for offline usage.
+# @param node_group_environment environment for the PEADM specific node groups, if not set it will be gathered from pe.conf or production
 #
 plan peadm::install (
   # Standard
@@ -74,6 +75,7 @@
   Peadm::Download_mode       $download_mode          = 'bolthost',
   Boolean                    $permit_unsafe_versions = false,
   String                     $token_lifetime         = '1y',
+  String[1] $node_group_environment = peadm::get_node_group_environment($primary_host),
 ) {
   # Log parameters for debugging 
   peadm::log_plan_parameters({
@@ -148,6 +150,7 @@
     internal_compiler_b_pool_address => $internal_compiler_b_pool_address,
     deploy_environment               => $deploy_environment,
     ldap_config                      => $ldap_config,
+    node_group_environment           => $node_group_environment,
 
     # Other
     stagingdir                       => $stagingdir,
diff --git a/plans/subplans/configure.pp b/plans/subplans/configure.pp
@@ -19,6 +19,7 @@
 # @param final_agent_state
 #   Configures the state the puppet agent should be in on infrastructure nodes
 #   after PE is configured successfully.
+# @param node_group_environment environment for the PEADM specific node groups, if not set it will be gathered from pe.conf or production
 #
 plan peadm::subplans::configure (
   # Standard
@@ -43,7 +44,8 @@
 
   # Other
   String           $stagingdir                   = '/tmp',
-  Enum['running', 'stopped'] $final_agent_state  = 'running'
+  Enum['running', 'stopped'] $final_agent_state  = 'running',
+  String[1] $node_group_environment = peadm::get_node_group_environment($primary_host),
 ) {
   # TODO: get and validate PE version
 
@@ -106,6 +108,7 @@
       compiler_pool_address            => $compiler_pool_address,
       internal_compiler_a_pool_address => $internal_compiler_a_pool_address,
       internal_compiler_b_pool_address => $internal_compiler_b_pool_address,
+      node_group_environment           => $node_group_environment,
       require                          => Class['peadm::setup::node_manager_yaml'],
     }
   }
diff --git a/plans/upgrade.pp b/plans/upgrade.pp
@@ -31,6 +31,7 @@
 #   Directory the installer tarball will be uploaded to or expected to be in
 #   for offline usage.
 # @param begin_at_step The step where the plan should start. If not set, it will start at the beginning
+# @param node_group_environment environment for the PEADM specific node groups, if not set it will be gathered from pe.conf or production
 #
 plan peadm::upgrade (
   # Standard
@@ -61,6 +62,7 @@
   Boolean                    $permit_unsafe_versions = false,
 
   Optional[Peadm::UpgradeSteps] $begin_at_step = undef,
+  String[1] $node_group_environment = peadm::get_node_group_environment($primary_host),
 ) {
   # Log parameters for debugging
   peadm::log_plan_parameters({
@@ -361,6 +363,7 @@
         compiler_pool_address            => $compiler_pool_address,
         internal_compiler_a_pool_address => $internal_compiler_a_pool_address,
         internal_compiler_b_pool_address => $internal_compiler_b_pool_address,
+        node_group_environment           => $node_group_environment,
         require                          => Class['peadm::setup::node_manager_yaml'],
       }
     }
diff --git a/plans/util/update_classification.pp b/plans/util/update_classification.pp
@@ -2,9 +2,12 @@
 #
 # @summary Configure classification
 #
+# @param node_group_environment environment for the PEADM specific node groups, if not set it will be gathered from pe.conf or production
+#
 plan peadm::util::update_classification (
   # Standard
   Peadm::SingleTargetSpec           $targets,
+  String[1]                         $node_group_environment,
   Optional[Hash]                    $peadm_config = undef,
   Optional[Peadm::SingleTargetSpec] $server_a_host = undef,
   Optional[Peadm::SingleTargetSpec] $server_b_host = undef,
@@ -76,6 +79,7 @@
       compiler_pool_address            => $new['params']['compiler_pool_address'],
       internal_compiler_a_pool_address => $new['params']['internal_compiler_a_pool_address'],
       internal_compiler_b_pool_address => $new['params']['internal_compiler_b_pool_address'],
+      node_group_environment           => $node_group_environment,
       require                          => Class['peadm::setup::node_manager_yaml'],
     }
   }
diff --git a/spec/functions/get_node_group_environment_spec.rb b/spec/functions/get_node_group_environment_spec.rb
@@ -0,0 +1,44 @@
+# frozen_string_literal: true
+
+require 'spec_helper'
+
+describe 'peadm::get_node_group_environment' do
+  include BoltSpec::BoltContext
+  around(:each) do |example|
+    in_bolt_context do
+      example.run
+    end
+  end
+  let(:primary) do
+    ['test-vm.puppet.vm']
+  end
+
+  context 'returns production on empty pe.conf' do
+    it do
+      expect_task('peadm::read_file').with_params('path' => '/etc/puppetlabs/enterprise/conf.d/pe.conf').always_return({ 'content' => '{}' })
+      is_expected.to run.with_params(primary).and_return('production')
+    end
+  end
+
+  context 'returns production on non-empty pe.conf with unrelated settings' do
+    it do
+      expect_task('peadm::read_file').with_params('path' => '/etc/puppetlabs/enterprise/conf.d/pe.conf').always_return({ 'content' => '{"foo": true}' })
+      is_expected.to run.with_params(primary).and_return('production')
+    end
+  end
+  context 'returns environment from pe.conf when set twice correctly' do
+    it do
+      pe = '{"pe_install::install::classification::pe_node_group_environment": "foo", "puppet_enterprise::master::recover_configuration::pe_environment": "foo"}'
+      expect_task('peadm::read_file').with_params('path' => '/etc/puppetlabs/enterprise/conf.d/pe.conf').always_return({ 'content' => pe })
+      is_expected.to run.with_params(primary).and_return('foo')
+    end
+  end
+  context 'fails when both PE options are different' do
+    it do
+      pe = '{"pe_install::install::classification::pe_node_group_environment": "foo", "puppet_enterprise::master::recover_configuration::pe_environment": "bar"}'
+      expect_task('peadm::read_file').with_params('path' => '/etc/puppetlabs/enterprise/conf.d/pe.conf').always_return({ 'content' => pe })
+      is_expected.to run.with_params(primary).and_raise_error(Puppet::PreformattedError,
+%r{pe_install::install::classification::pe_node_group_environment and puppet_enterprise::master::recover_configuration::pe_environment})
+    end
+  end
+end
diff --git a/spec/plans/convert_spec.rb b/spec/plans/convert_spec.rb
@@ -19,7 +19,8 @@
     allow_apply
 
     expect_task('peadm::cert_data').return_for_targets('primary' => trustedjson).be_called_times(2)
-    expect_task('peadm::read_file').always_return({ 'content' => '2021.7.9' })
+    expect_task('peadm::read_file').with_params('path' => '/opt/puppetlabs/server/pe_version').always_return({ 'content' => '2021.7.9' })
+    expect_task('peadm::read_file').with_params('path' => '/etc/puppetlabs/enterprise/conf.d/pe.conf').always_return({ 'content' => '{}' })
     expect_task('peadm::get_group_rules').return_for_targets('primary' => { '_output' => '{"rules": []}' })
     expect_task('peadm::node_group_unpin').with_targets('primary').with_params({ 'node_certnames' => ['pe_compiler_legacy'], 'group_name' => 'PE Master' })
     expect_task('peadm::check_legacy_compilers').with_targets('primary').with_params({ 'legacy_compilers' => 'pe_compiler_legacy' }).return_for_targets('primary' => { '_output' => '' })
diff --git a/spec/plans/install_spec.rb b/spec/plans/install_spec.rb
@@ -8,6 +8,7 @@
       allow_out_message
       expect_plan('peadm::subplans::install')
       expect_plan('peadm::subplans::configure')
+      expect_task('peadm::read_file').with_params('path' => '/etc/puppetlabs/enterprise/conf.d/pe.conf').always_return({ 'content' => '{}' })
       expect(run_plan('peadm::install', 'primary_host' => 'primary', 'console_password' => 'puppetLabs123!', 'version' => '2021.7.9')).to be_ok
     end
   end
diff --git a/spec/plans/upgrade_spec.rb b/spec/plans/upgrade_spec.rb
@@ -35,6 +35,7 @@ def allow_standard_non_returning_calls
       .with_params('path' => '/opt/puppetlabs/server/pe_build')
       .always_return({ 'content' => '2021.7.3' })
 
+    expect_task('peadm::read_file').with_params('path' => '/etc/puppetlabs/enterprise/conf.d/pe.conf').always_return({ 'content' => '{}' })
     expect_task('peadm::cert_data').return_for_targets('primary' => trusted_primary).be_called_times(1)
     expect_task('peadm::check_pe_master_rules').always_return(pe_rule_check)
 
@@ -50,6 +51,7 @@ def allow_standard_non_returning_calls
     expect_task('peadm::read_file')
       .with_params('path' => '/opt/puppetlabs/server/pe_build')
       .always_return({ 'content' => '2021.7.3' })
+    expect_task('peadm::read_file').with_params('path' => '/etc/puppetlabs/enterprise/conf.d/pe.conf').always_return({ 'content' => '{}' })
 
     expect_task('peadm::cert_data').return_for_targets('primary' => trusted_primary,
                                                        'compiler' => trusted_compiler).be_called_times(1)
@@ -109,7 +111,7 @@ def allow_standard_non_returning_calls
     it 'updates pe.conf if r10k_known_hosts is set' do
       expect_task('peadm::read_file')
         .with_params('path' => '/etc/puppetlabs/enterprise/conf.d/pe.conf')
-        .always_return({ 'content' => <<~PECONF })
+        .always_return({ 'content' => <<~PECONF }).be_called_times(2)
           # spec pe.conf
           "puppet_enterprise::puppet_master_host": "%{::trusted.certname}"
         PECONF
@@ -130,6 +132,7 @@ def allow_standard_non_returning_calls
     it 'warns if upgrading to 2023.3+ from 2023.0- without r10k_known_hosts set' do
       # This is fairly horrible, but expect_out_message doesn't take a regex.
       expect_out_message.with_params(r10k_warning)
+      expect_task('peadm::read_file').with_params('path' => '/etc/puppetlabs/enterprise/conf.d/pe.conf').always_return({ 'content' => '{}' })
       expect_task('peadm::check_pe_master_rules').always_return(pe_rule_check)
 
       expect(run_plan('peadm::upgrade',
@@ -143,6 +146,7 @@ def allow_standard_non_returning_calls
 
       it 'does not warn if r10k_known_hosts is not set' do
         expect_out_message.with_params(r10k_warning).not_be_called
+        expect_task('peadm::read_file').with_params('path' => '/etc/puppetlabs/enterprise/conf.d/pe.conf').always_return({ 'content' => '{}' })
         expect_task('peadm::check_pe_master_rules').always_return(pe_rule_check)
 
         expect(run_plan('peadm::upgrade',

Original file line number	Diff line number	Diff line change
`@@ -5,6 +5,7 @@`
`5`	`5`	`# management using PEAdm.`
`6`	`6`	`#`
`7`	`7`	`# @param begin_at_step The step where the plan should start. If not set, it will start at the beginning`
	`8`	`+# @param node_group_environment environment for the PEADM specific node groups, if not set it will be gathered from pe.conf or production`
`8`	`9`	`#`
`9`	`10`	`plan peadm::convert (`
`10`	`11`	`# Standard`
`@@ -26,6 +27,7 @@`
`26`	`27`	`Array[String] $dns_alt_names = [],`
`27`	`28`
`28`	`29`	`Optional[Peadm::ConvertSteps] $begin_at_step = undef,`
	`30`	`+ String[1] $node_group_environment = peadm::get_node_group_environment($primary_host),`
`29`	`31`	`) {`
`30`	`32`	`peadm::assert_supported_bolt_version()`
`31`	`33`
`@@ -316,6 +318,7 @@`
`316`	`318`	`compiler_pool_address => $compiler_pool_address,`
`317`	`319`	`internal_compiler_a_pool_address => $internal_compiler_a_pool_address,`
`318`	`320`	`internal_compiler_b_pool_address => $internal_compiler_b_pool_address,`
	`321`	`+ node_group_environment => $node_group_environment,`
`319`	`322`	`require => Class['peadm::setup::node_manager_yaml'],`
`320`	`323`	`}`
`321`	`324`
Original file line number	Diff line number	Diff line change
`@@ -2,9 +2,12 @@`
`2`	`2`	`#`
`3`	`3`	`# @summary Configure classification`
`4`	`4`	`#`
	`5`	`+# @param node_group_environment environment for the PEADM specific node groups, if not set it will be gathered from pe.conf or production`
	`6`	`+#`
`5`	`7`	`plan peadm::util::update_classification (`
`6`	`8`	`# Standard`
`7`	`9`	`Peadm::SingleTargetSpec $targets,`
	`10`	`+ String[1] $node_group_environment,`
`8`	`11`	`Optional[Hash] $peadm_config = undef,`
`9`	`12`	`Optional[Peadm::SingleTargetSpec] $server_a_host = undef,`
`10`	`13`	`Optional[Peadm::SingleTargetSpec] $server_b_host = undef,`
`@@ -76,6 +79,7 @@`
`76`	`79`	`compiler_pool_address => $new['params']['compiler_pool_address'],`
`77`	`80`	`internal_compiler_a_pool_address => $new['params']['internal_compiler_a_pool_address'],`
`78`	`81`	`internal_compiler_b_pool_address => $new['params']['internal_compiler_b_pool_address'],`
	`82`	`+ node_group_environment => $node_group_environment,`
`79`	`83`	`require => Class['peadm::setup::node_manager_yaml'],`
`80`	`84`	`}`
`81`	`85`	`}`