gistfile1.txt

# /etc/puppetlabs/code/environments/production/modules/mymodule/lib/puppet/functions/mymodule/upcase.rb
require 'cassandra'
require 'json'
require 'yaml'

class MetadataServiceSessionAdapter < Puppet::Pops::Adaptable::Adapter
  attr_accessor :session
end

Puppet::Functions.create_function(:'puppet_metadata_service::data_hash', Puppet::Functions::InternalFunction) do

  dispatch :hiera_data do
    scope_param
    param 'Hash', :options
    param 'Puppet::LookupContext', :context
  end

  def hiera_data(scope, options, context)
    uri = options['uri']

    begin
      context.explain { "[puppet_metadata_service::data_hash] Attempting to use cached session from adapter" }
      session = MetadataServiceSessionAdapter.get(scope.compiler).session
      context.explain { "[puppet_metadata_service::data_hash] Obtained cached session from adapter" }
    rescue
      context.explain { "[puppet_metadata_service::data_hash] Failed to obtain cached session from adapter" }
      context.explain { "[puppet_metadata_service::data_hash] Creating new session and saving to adapter" }
      MetadataServiceSessionAdapter.adapt(scope.compiler).session = Cassandra.cluster(hosts: [Socket.gethostname]).connect('puppet')
      retry
    end

    data = session.execute(
      "SELECT key,value FROM hieradata where level=%s" % "$$#{uri}$$").rows.map { |row|
        { row['key'] => row['value'] }
      }.reduce({}, :merge
    )

    if data.empty?
      context.not_found
    else
      return data
    end
  end
end