Why Gemfury? Push, build, and install  RubyGems npm packages Python packages Maven artifacts PHP packages Go Modules Debian packages RPM packages NuGet packages

Repository URL to install this package:

Details    
uoy-faculty_aws / lib / faculty_aws / db_connector.rb
Size: Mime:
# frozen_string_literal: true

require 'aws-sdk-rds'
require 'sequel'
require 'logger'

module FacultyAWS
  # Class to deal with database connections
  class DBConnector
    CONFIG_ITEMS = %i[adapter host port database user search_path password sslmode noauto azure url logger].freeze
    BOOLY_ITEMS = %i[noauto azure].freeze
    AWS_CONFIG_MAP = { engine: :adapter, dbname: :database, username: :user }.freeze

    # Mixin for Connection to regenerate auth_token on reconnect
    module ReconnectTokenGenerator
      def server_opts(...)
        super.merge(password: @token_proc.call)
      end

      attr_accessor :token_proc
    end

    def initialize(prefix: 'DB', **overrides)
      @prefix = prefix
      @extensions = []
      @config = {}
      @config.merge! config_from_env
      @config.merge! config_from_secret
      @config.merge! overrides.transform_keys(&:to_sym)
      auto_configure
    end

    def connection
      @connection ||= connect
    end

    def add_extension(*extensions)
      @extensions += extensions
      self
    end

    def add_pooling(validation_timeout: 10)
      @pooled_connection = { validation_timeout: }
      self
    end

    alias pooled_connection connection

    private

    def config_from_env
      CONFIG_ITEMS.each_with_object({}) do |item, config|
        value = ENV.fetch("#{@prefix}_#{item}".upcase, nil)
        value = value.to_s.match?(/^[1ty]/i) if BOOLY_ITEMS.include?(item)
        config[item] = value unless value.nil?
      end
    end

    def config_from_secret
      secret = ENV.fetch("#{@prefix}_SECRET".upcase, nil)
      return {} unless secret

      JSON.parse(secret, symbolize_names: true).transform_keys do |aws_item|
        AWS_CONFIG_MAP[aws_item] || aws_item
      end
    end

    def auto_configure
      return if @config[:noauto]

      @use_auth_token = true unless @config.key? :password
      @config[:sslmode] ||= 'require'
    end

    def token_generator
      @token_generator ||=
        Aws::RDS::AuthTokenGenerator.new(
          credentials:
            Aws::Credentials.new(
              ENV.fetch('AWS_ACCESS_KEY_ID', nil),
              ENV.fetch('AWS_SECRET_ACCESS_KEY', nil),
              ENV.fetch('AWS_SESSION_TOKEN', nil)
            )
        )
    end

    def aws_auth_token
      token_generator.auth_token(
        region: ENV.fetch('AWS_REGION', nil),
        endpoint: "#{@config[:host]}:#{@config[:port] || 5432}",
        user_name: @config[:user]
      )
    end

    def create_logger
      logfile = @config[:logger]
      @config[:logger] = Logger.new(logfile == '-' ? $stdout : logfile)
    end

    def inner_connect
      create_logger if @config[:logger]
      @config[:password] = aws_auth_token if @use_auth_token
      if @config.key?(:url)
        Sequel.connect(@config[:url], **@config)
      else
        Sequel.connect(@config)
      end
    end

    def connect
      inner_connect.tap do |conn|
        conn.extension(*@extensions)
        setup_pooling(conn) if @pooled_connection
        setup_auth_refresh(conn) if @use_auth_token
      end
    end

    def setup_pooling(connection)
      connection.extension :connection_validator
      connection.pool.connection_validation_timeout = @pooled_connection[:validation_timeout]
    end

    def setup_auth_refresh(connection)
      connection.extend ReconnectTokenGenerator
      connection.token_proc = proc { aws_auth_token }
    end
  end
end