lib/puppet-ghostbuster/puppetdb.rb in puppet-ghostbuster-1.1.0 vs lib/puppet-ghostbuster/puppetdb.rb in puppet-ghostbuster-1.2.0

@@ -11,38 +11,49 @@
     rescue LoadError
       @@puppetdb = "https://#{Puppet[:server]}:8081"
     end
 
     def self.client
-      @@client ||= ::PuppetDB::Client.new({
-                                            server: ENV['PUPPETDB_URL'] || @@puppetdb,
-                                            pem: {
-                                              'key' => ENV['PUPPETDB_KEY_FILE'] || Puppet[:hostprivkey],
-                                              'cert' => ENV['PUPPETDB_CERT_FILE'] || Puppet[:hostcert],
-                                              'ca_file' => ENV['PUPPETDB_CACERT_FILE'] || Puppet[:localcacert],
-                                            },
-                                          }, 4)
+      @@client ||= begin
+        options = {
+          server: ENV['PUPPETDB_URL'] || @@puppetdb,
+        }
+
+        if ENV['PE_TOKEN']
+          token_file = File.expand_path(ENV['PE_TOKEN'])
+          options[:token] = File.exist?(token_file) ? File.read(token_file) : ENV.fetch('PE_TOKEN')
+          options[:cacert] = ENV['PUPPETDB_CACERT_FILE'] || Puppet[:localcacert]
+        else
+          options[:pem] = {
+            'key' => ENV['PUPPETDB_KEY_FILE'] || Puppet[:hostprivkey],
+            'cert' => ENV['PUPPETDB_CERT_FILE'] || Puppet[:hostcert],
+            'ca_file' => ENV['PUPPETDB_CACERT_FILE'] || Puppet[:localcacert],
+          }
+        end
+
+        ::PuppetDB::Client.new(options, 4)
+      end
     end
 
     def client
       self.class.client
     end
 
     def self.classes
       @@classes ||= client.request('',
-                                   'resources[title] { type = "Class" and nodes { deactivated is null } }').data.map do |r|
+                                   'resources[title] { type = "Class" and nodes { deactivated is null } group by title }').data.map do |r|
         r['title']
-      end.uniq
+      end
     end
 
     def classes
       self.class.classes
     end
 
     def self.resources
-      @@resources ||= client.request('', 'resources[type] { nodes { deactivated is null } }').data.map do |r|
+      @@resources ||= client.request('', 'resources[type] { nodes { deactivated is null } group by type }').data.map do |r|
         r['type']
-      end.uniq
+      end
     end
 
     def resources
       self.class.resources
     end