pg-aws_rds_iam 0.1.0

data/Rakefile

@@ -0,0 +1,39 @@
+# frozen_string_literal: true
+
+require "bundler/gem_tasks"
+require "rake/testtask"
+require "rubocop/rake_task"
+require "yard"
+
+namespace :test do
+  Rake::TestTask.new :acceptance do |t|
+    t.description = "Run acceptance tests"
+    t.libs << "test"
+    t.test_files = ["test/acceptance/test.rb"]
+  end
+
+  Rake::TestTask.new :unit do |t|
+    t.description = "Run unit tests"
+    t.libs << "test"
+    t.test_files = FileList["test/**/*_test.rb"]
+  end
+end
+
+desc "Run all tests"
+task :test => ["test:unit", "test:acceptance"]
+
+RuboCop::RakeTask.new do |t|
+  t.formatters = ["fuubar"]
+end
+
+desc "Generate documentation"
+YARD::Rake::YardocTask.new
+
+namespace :yard do
+  desc "Run documentation server"
+  task :server do
+    exec "bin/yard", "server", "--reload"
+  end
+end
+
+task :default => ["test:unit", :rubocop]

data/bin/bundle

@@ -0,0 +1,105 @@
+#!/usr/bin/env ruby
+# frozen_string_literal: true
+
+#
+# This file was generated by Bundler.
+#
+# The application 'bundle' is installed as part of a gem, and
+# this file is here to facilitate running it.
+#
+
+require "rubygems"
+
+m = Module.new do
+  module_function
+
+  def invoked_as_script?
+    File.expand_path($0) == File.expand_path(__FILE__)
+  end
+
+  def env_var_version
+    ENV["BUNDLER_VERSION"]
+  end
+
+  def cli_arg_version
+    return unless invoked_as_script? # don't want to hijack other binstubs
+    return unless "update".start_with?(ARGV.first || " ") # must be running `bundle update`
+    bundler_version = nil
+    update_index = nil
+    ARGV.each_with_index do |a, i|
+      if update_index && update_index.succ == i && a =~ Gem::Version::ANCHORED_VERSION_PATTERN
+        bundler_version = a
+      end
+      next unless a =~ /\A--bundler(?:[= ](#{Gem::Version::VERSION_PATTERN}))?\z/
+      bundler_version = $1 || ">= 0.a"
+      update_index = i
+    end
+    bundler_version
+  end
+
+  def gemfile
+    gemfile = ENV["BUNDLE_GEMFILE"]
+    return gemfile if gemfile && !gemfile.empty?
+
+    File.expand_path("../../Gemfile", __FILE__)
+  end
+
+  def lockfile
+    lockfile =
+      case File.basename(gemfile)
+      when "gems.rb" then gemfile.sub(/\.rb$/, gemfile)
+      else "#{gemfile}.lock"
+      end
+    File.expand_path(lockfile)
+  end
+
+  def lockfile_version
+    return unless File.file?(lockfile)
+    lockfile_contents = File.read(lockfile)
+    return unless lockfile_contents =~ /\n\nBUNDLED WITH\n\s{2,}(#{Gem::Version::VERSION_PATTERN})\n/
+    Regexp.last_match(1)
+  end
+
+  def bundler_version
+    @bundler_version ||= begin
+      env_var_version || cli_arg_version ||
+        lockfile_version || "#{Gem::Requirement.default}.a"
+    end
+  end
+
+  def load_bundler!
+    ENV["BUNDLE_GEMFILE"] ||= gemfile
+
+    # must dup string for RG < 1.8 compatibility
+    activate_bundler(bundler_version.dup)
+  end
+
+  def activate_bundler(bundler_version)
+    if Gem::Version.correct?(bundler_version) && Gem::Version.new(bundler_version).release < Gem::Version.new("2.0")
+      bundler_version = "< 2"
+    end
+    gem_error = activation_error_handling do
+      gem "bundler", bundler_version
+    end
+    return if gem_error.nil?
+    require_error = activation_error_handling do
+      require "bundler/version"
+    end
+    return if require_error.nil? && Gem::Requirement.new(bundler_version).satisfied_by?(Gem::Version.new(Bundler::VERSION))
+    warn "Activating bundler (#{bundler_version}) failed:\n#{gem_error.message}\n\nTo install the version of bundler this project requires, run `gem install bundler -v '#{bundler_version}'`"
+    exit 42
+  end
+
+  def activation_error_handling
+    yield
+    nil
+  rescue StandardError, LoadError => e
+    e
+  end
+end
+
+m.load_bundler!
+
+if m.invoked_as_script?
+  load Gem.bin_path("bundler", "bundle")
+end

data/bin/console

@@ -0,0 +1,8 @@
+#!/usr/bin/env ruby
+# frozen_string_literal: true
+
+require "bundler/setup"
+require "pg/aws_rds_iam"
+require "pry"
+
+Pry.start

data/bin/install-bundler

@@ -0,0 +1,6 @@
+#!/bin/sh
+set -o errexit
+set -o nounset
+
+bundler_version=$(ruby -e 'puts /\n\nBUNDLED WITH\n\s{2,}(?<version>#{Gem::Version::VERSION_PATTERN})\n/.match(File.read("Gemfile.lock"))[:version]')
+gem install bundler --version "${bundler_version}"

data/bin/rake

@@ -0,0 +1,29 @@
+#!/usr/bin/env ruby
+# frozen_string_literal: true
+
+#
+# This file was generated by Bundler.
+#
+# The application 'rake' is installed as part of a gem, and
+# this file is here to facilitate running it.
+#
+
+require "pathname"
+ENV["BUNDLE_GEMFILE"] ||= File.expand_path("../../Gemfile",
+  Pathname.new(__FILE__).realpath)
+
+bundle_binstub = File.expand_path("../bundle", __FILE__)
+
+if File.file?(bundle_binstub)
+  if File.read(bundle_binstub, 300) =~ /This file was generated by Bundler/
+    load(bundle_binstub)
+  else
+    abort("Your `bin/bundle` was not generated by Bundler, so this binstub cannot run.
+Replace `bin/bundle` by running `bundle binstubs bundler --force`, then run this command again.")
+  end
+end
+
+require "rubygems"
+require "bundler/setup"
+
+load Gem.bin_path("rake", "rake")

data/bin/setup

@@ -0,0 +1,6 @@
+#!/bin/sh
+set -o errexit
+set -o nounset
+
+bin/install-bundler
+bin/bundle install

data/bin/yard

@@ -0,0 +1,29 @@
+#!/usr/bin/env ruby
+# frozen_string_literal: true
+
+#
+# This file was generated by Bundler.
+#
+# The application 'yard' is installed as part of a gem, and
+# this file is here to facilitate running it.
+#
+
+require "pathname"
+ENV["BUNDLE_GEMFILE"] ||= File.expand_path("../../Gemfile",
+  Pathname.new(__FILE__).realpath)
+
+bundle_binstub = File.expand_path("../bundle", __FILE__)
+
+if File.file?(bundle_binstub)
+  if File.read(bundle_binstub, 300) =~ /This file was generated by Bundler/
+    load(bundle_binstub)
+  else
+    abort("Your `bin/bundle` was not generated by Bundler, so this binstub cannot run.
+Replace `bin/bundle` by running `bundle binstubs bundler --force`, then run this command again.")
+  end
+end
+
+require "rubygems"
+require "bundler/setup"
+
+load Gem.bin_path("yard", "yard")

data/lib/pg/aws_rds_iam.rb

@@ -0,0 +1,30 @@
+# frozen_string_literal: true
+
+require "aws-sdk-rds"
+require "pg"
+require "strscan"
+require "uri"
+
+require_relative "aws_rds_iam/auth_token_generator"
+require_relative "aws_rds_iam/auth_token_generator_registry"
+require_relative "aws_rds_iam/auth_token_injector"
+require_relative "aws_rds_iam/connection"
+require_relative "aws_rds_iam/connection_info"
+require_relative "aws_rds_iam/version"
+
+# The top-level [PG](https://deveiate.org/code/pg/PG.html) namespace.
+module PG
+  # The top-level AWS RDS IAM plugin namespace.
+  module AWS_RDS_IAM
+    @auth_token_generators = AuthTokenGeneratorRegistry.new
+
+    # Registry of available {AuthTokenGenerator}s.
+    #
+    # @return [AuthTokenGeneratorRegistry]
+    def self.auth_token_generators
+      @auth_token_generators
+    end
+
+    PG::Connection.singleton_class.prepend Connection
+  end
+end

data/lib/pg/aws_rds_iam/auth_token_generator.rb

@@ -0,0 +1,29 @@
+# frozen_string_literal: true
+
+module PG
+  module AWS_RDS_IAM
+    # Generates short-lived authentication tokens for connecting to Amazon RDS instances.
+    #
+    # @see https://docs.aws.amazon.com/AmazonRDS/latest/UserGuide/UsingWithRDS.IAMDBAuth.html
+    class AuthTokenGenerator
+      # Creates a new authentication token generator.
+      #
+      # @param credentials [Aws::CredentialProvider] the IAM credentials with which to sign the token
+      # @param region [String] the AWS region in which the RDS instances are running
+      def initialize(credentials:, region:)
+        @generator = Aws::RDS::AuthTokenGenerator.new(credentials: credentials)
+        @region = region
+      end
+
+      # Generates an authentication token for connecting to an Amazon RDS instance.
+      #
+      # @param host [String] the host name of the RDS instance that you want to access
+      # @param port [String] the port number used for connecting to your RDS instance
+      # @param user [String] the database account that you want to access
+      # @return [String] the generated authentication token
+      def call(host:, port:, user:)
+        @generator.auth_token(region: @region, endpoint: "#{host}:#{port}", user_name: user)
+      end
+    end
+  end
+end

data/lib/pg/aws_rds_iam/auth_token_generator_registry.rb

@@ -0,0 +1,68 @@
+# frozen_string_literal: true
+
+module PG
+  module AWS_RDS_IAM
+    # Registers {AuthTokenGenerator}s to be used to generate authentication tokens for `PG::Connection`s that have the
+    # `aws_rds_iam_auth_token_generator` connection parameter set to the registered name.
+    class AuthTokenGeneratorRegistry
+      # Creates a new `AuthTokenRegistry`.
+      #
+      # @param default_auth_token_generator_class [Class] the class to register as the default {AuthTokenGenerator}
+      def initialize(default_auth_token_generator_class: AuthTokenGenerator)
+        @default_auth_token_generator_class = default_auth_token_generator_class
+        reset
+      end
+
+      # Registers an {AuthTokenGenerator}.
+      #
+      # @param name [String, Symbol]
+      # @return [void]
+      # @yieldreturn [AuthTokenGenerator]
+      def add(name, &block)
+        @registry[name.to_s] = Memoizer.new(&block)
+      end
+
+      # Looks up an {AuthTokenGenerator} by name.
+      #
+      # @param name [String, Symbol]
+      # @return [AuthTokenGenerator]
+      def fetch(name)
+        @registry.fetch(name.to_s).call
+      end
+
+      # Unregisters an {AuthTokenGenerator}.
+      #
+      # @param name [String, Symbol]
+      # @return [void]
+      def remove(name)
+        @registry.delete name.to_s
+      end
+
+      # Unregisters all {AuthTokenGenerator}s and re-registers the default {AuthTokenGenerator}.
+      #
+      # @return [void]
+      def reset
+        @registry = {}
+
+        add :default do
+          config = Aws::RDS::Client.new.config
+          @default_auth_token_generator_class.new(credentials: config.credentials, region: config.region)
+        end
+      end
+
+      class Memoizer # rubocop:disable Style/Documentation
+        def initialize(&block)
+          @block = block
+        end
+
+        def call
+          return @auth_token_generator if defined?(@auth_token_generator)
+
+          @auth_token_generator = @block.call
+        end
+      end
+
+      private_constant :Memoizer
+    end
+  end
+end

data/lib/pg/aws_rds_iam/auth_token_injector.rb

@@ -0,0 +1,48 @@
+# frozen_string_literal: true
+
+module PG
+  module AWS_RDS_IAM
+    class AuthTokenInjector
+      def self.call(connection_string, auth_token_generators: AWS_RDS_IAM.auth_token_generators)
+        new(connection_string, auth_token_generators: auth_token_generators).call
+      end
+
+      def initialize(connection_string, auth_token_generators:)
+        @connection_string = connection_string
+        @connection_info = ConnectionInfo.new(connection_string)
+        @connection_defaults = PG::Connection.conndefaults_hash
+        @auth_token_generators = auth_token_generators
+      end
+
+      def call
+        return @connection_string unless generate_auth_token?
+
+        @connection_info.password = generate_auth_token
+
+        @connection_info.to_s
+      end
+
+      private
+
+      def generate_auth_token?
+        @connection_info.auth_token_generator_name
+      end
+
+      def generate_auth_token
+        @auth_token_generators
+          .fetch(@connection_info.auth_token_generator_name)
+          .call(
+            user: @connection_info.user || default(:user),
+            host: @connection_info.host || default(:host),
+            port: @connection_info.port || default(:port)
+          )
+      end
+
+      def default(key)
+        @connection_defaults.fetch(key)
+      end
+    end
+
+    private_constant :AuthTokenInjector
+  end
+end

data/lib/pg/aws_rds_iam/connection.rb

@@ -0,0 +1,25 @@
+# frozen_string_literal: true
+
+module PG
+  module AWS_RDS_IAM
+    module Connection
+      def conndefaults
+        super + [{
+          keyword: "aws_rds_iam_auth_token_generator",
+          envvar: nil,
+          compiled: nil,
+          val: nil,
+          label: "AWS-RDS-IAM-auth-token-generator",
+          dispchar: "",
+          dispsize: 64
+        }]
+      end
+
+      def parse_connect_args(*)
+        AuthTokenInjector.call(super)
+      end
+    end
+
+    private_constant :Connection
+  end
+end

data/lib/pg/aws_rds_iam/connection_info.rb

@@ -0,0 +1,21 @@
+# frozen_string_literal: true
+
+require_relative "connection_info/keyword_value_string"
+require_relative "connection_info/parse_error"
+require_relative "connection_info/uri"
+
+module PG
+  module AWS_RDS_IAM
+    module ConnectionInfo
+      def self.new(connection_string)
+        if URI.match?(connection_string)
+          URI.new(connection_string)
+        else
+          KeywordValueString.new(connection_string)
+        end
+      end
+    end
+
+    private_constant :ConnectionInfo
+  end
+end