Module: Legion::Data::Connection

Defined in:

lib/legion/data/connection.rb

Defined Under Namespace

Classes: QueryFileLogger, SlowQueryLogger

Constant Summary

ADAPTERS =

%i[sqlite mysql2 postgres].freeze

GENERIC_KEYS =

%i[max_connections pool_timeout preconnect single_threaded test name].freeze

ADAPTER_KEYS =

{
  sqlite:   %i[timeout readonly disable_dqs],
  postgres: %i[connect_timeout sslmode sslrootcert search_path],
  mysql2:   %i[connect_timeout read_timeout write_timeout encoding sql_mode]
}.freeze

ADAPTER_DEFAULTS =

{
  sqlite:   { timeout: 5000, readonly: false, disable_dqs: true },
  postgres: { connect_timeout: 20, sslmode: 'disable' },
  mysql2:   { connect_timeout: 120, encoding: 'utf8mb4' }
}.freeze

QUERY_LOG_DIR =

File.expand_path('~/.legionio/logs').freeze

Class Attribute Summary

• .sequel ⇒ Object

  Returns the value of attribute sequel.

Class Method Summary

• .adapter ⇒ Object
• .connect_with_replicas ⇒ Object
• .creds_builder(final_creds = {}) ⇒ Object
• .merge_tls_creds(creds, adapter:, port:) ⇒ Object
• .pool_stats ⇒ Object
• .read_server ⇒ Object
• .replica_servers ⇒ Object
• .setup ⇒ Object
• .shutdown ⇒ Object
• .stats ⇒ Object

Class Attribute Details

.sequel ⇒ Object

Returns the value of attribute sequel.

# File 'lib/legion/data/connection.rb', line 99

def sequel
  @sequel
end

Class Method Details

.adapter ⇒ Object

# File 'lib/legion/data/connection.rb', line 101

def adapter
  @adapter ||= Legion::Settings[:data][:adapter]&.to_sym || :sqlite
end

.connect_with_replicas ⇒ Object

# File 'lib/legion/data/connection.rb', line 197

def connect_with_replicas
  return unless adapter == :postgres

  replica_url  = Legion::Settings[:data][:read_replica_url]
  replica_list = Array(Legion::Settings[:data][:replicas]).dup

  replica_list.prepend(replica_url) if replica_url && !replica_url.empty?
  replica_list.uniq!
  replica_list.compact!

  return if replica_list.empty?

  @sequel.extension(:server_block)

  replica_list.each_with_index do |url, idx|
    @sequel.add_servers("read_#{idx}": url)
  end

  @replica_servers = replica_list.each_with_index.map { |_, idx| :"read_#{idx}" }
  Legion::Logging.debug "Registered #{@replica_servers.size} read replica(s)" if defined?(Legion::Logging)
end

.creds_builder(final_creds = {}) ⇒ Object

# File 'lib/legion/data/connection.rb', line 255

def creds_builder(final_creds = {})
  final_creds.merge! Legion::Data::Settings.creds(adapter)
  final_creds.merge! Legion::Settings[:data][:creds] if Legion::Settings[:data][:creds].is_a? Hash

  port = final_creds[:port]
  merge_tls_creds(final_creds, adapter: adapter, port: port)

  return final_creds if Legion::Settings[:vault].nil?

  if Legion::Settings[:vault][:connected] && ::Vault.sys.mounts.key?(:database)
    temp_vault_creds = Legion::Crypt.read('database/creds/legion')
    final_creds[:user] = temp_vault_creds[:username]
    final_creds[:password] = temp_vault_creds[:password]
  end

  final_creds
end

.merge_tls_creds(creds, adapter:, port:) ⇒ Object

# File 'lib/legion/data/connection.rb', line 229

def merge_tls_creds(creds, adapter:, port:)
  return creds if adapter == :sqlite
  return creds unless defined?(Legion::Crypt::TLS)

  tls_settings = data_tls_settings
  return creds unless tls_settings[:enabled] == true

  tls = Legion::Crypt::TLS.resolve(tls_settings, port: port)
  return creds unless tls[:enabled]

  case adapter
  when :postgres
    creds[:sslmode]     = tls[:verify] == :none ? 'require' : 'verify-full'
    creds[:sslrootcert] = tls[:ca] if tls[:ca]
    creds[:sslcert]     = tls[:cert] if tls[:cert]
    creds[:sslkey]      = tls[:key] if tls[:key]
  when :mysql2
    creds[:ssl_mode] = tls[:verify] == :none ? 'required' : 'verify_identity'
    creds[:sslca]    = tls[:ca] if tls[:ca]
    creds[:sslcert]  = tls[:cert] if tls[:cert]
    creds[:sslkey]   = tls[:key] if tls[:key]
  end

  creds
end

.pool_stats ⇒ Object

# File 'lib/legion/data/connection.rb', line 157

def pool_stats
  return {} unless @sequel

  pool = @sequel.pool
  stats = {
    type:     pool.pool_type,
    size:     pool.size,
    max_size: pool.respond_to?(:max_size) ? pool.max_size : nil
  }

  case pool.pool_type
  when :timed_queue, :sharded_timed_queue
    queue_size = pool.instance_variable_get(:@queue)&.size || 0
    stats[:available] = queue_size
    stats[:in_use]    = stats[:size] - queue_size
    stats[:waiting]   = pool.num_waiting
  when :threaded, :sharded_threaded
    avail = pool.instance_variable_get(:@available_connections)
    stats[:available] = avail&.size || 0
    stats[:in_use]    = stats[:size] - stats[:available]
    stats[:waiting]   = pool.num_waiting
  when :single, :sharded_single
    stats[:available] = pool.size
    stats[:in_use]    = 0
    stats[:waiting]   = 0
  end

  stats.compact
rescue StandardError
  {}
end

.read_server ⇒ Object

# File 'lib/legion/data/connection.rb', line 219

def read_server
  return :default if @replica_servers.nil? || @replica_servers.empty?

  :read_0
end

.replica_servers ⇒ Object

# File 'lib/legion/data/connection.rb', line 225

def replica_servers
  @replica_servers || []
end

.setup ⇒ Object

# File 'lib/legion/data/connection.rb', line 105

def setup
  opts = sequel_opts
  @sequel = if adapter == :sqlite
              ::Sequel.connect(opts.merge(adapter: :sqlite, database: sqlite_path))
            else
              begin
                ::Sequel.connect(opts.merge(adapter: adapter, **creds_builder))
              rescue StandardError => e
                raise unless dev_fallback?

                if defined?(Legion::Logging)
                  Legion::Logging.warn(
                    "Shared DB unreachable (#{e.message}), dev_mode fallback to SQLite"
                  )
                end
                @adapter = :sqlite
                sqlite_opts = sequel_opts
                ::Sequel.connect(sqlite_opts.merge(adapter: :sqlite, database: sqlite_path))
              end
            end
  Legion::Settings[:data][:connected] = true
  if defined?(Legion::Logging)
    if adapter == :sqlite
      Legion::Logging.info "Connected to SQLite at #{sqlite_path}"
    else
      creds = Legion::Data::Settings.creds(adapter)
      user = creds[:user] || creds[:username] || 'unknown'
      host = creds[:host] || '127.0.0.1'
      port = creds[:port]
      db   = creds[:database] || creds[:db]
      Legion::Logging.info "Connected to #{adapter}://#{user}@#{host}:#{port}/#{db}"
    end
  end
  configure_extensions
  connect_with_replicas
end

.shutdown ⇒ Object

# File 'lib/legion/data/connection.rb', line 189

def shutdown
  @sequel&.disconnect
  @query_file_logger&.close
  @query_file_logger = nil
  Legion::Settings[:data][:connected] = false
  Legion::Logging.info 'Legion::Data connection closed' if defined?(Legion::Logging)
end

.stats ⇒ Object

# File 'lib/legion/data/connection.rb', line 142

def stats
  return { connected: false } unless @sequel

  data = Legion::Settings[:data]
  {
    connected: data[:connected],
    adapter:   adapter,
    pool:      pool_stats,
    tuning:    tuning_stats(data),
    database:  database_stats
  }
rescue StandardError => e
  { connected: (data[:connected] if data.is_a?(Hash)), adapter: adapter, error: e.message }
end