|
| 1 | +# frozen_string_literal: true |
| 2 | + |
| 3 | +require "semian/adapter" |
| 4 | +require "activerecord-trilogy-adapter" |
| 5 | +require "active_record/connection_adapters/trilogy_adapter" |
| 6 | + |
| 7 | +module ActiveRecord |
| 8 | + module ConnectionAdapters |
| 9 | + class TrilogyAdapter |
| 10 | + ActiveRecord::ActiveRecordError.include(::Semian::AdapterError) |
| 11 | + |
| 12 | + class SemianError < StatementInvalid |
| 13 | + def initialize(semian_identifier, *args) |
| 14 | + super(*args) |
| 15 | + @semian_identifier = semian_identifier |
| 16 | + end |
| 17 | + end |
| 18 | + |
| 19 | + ResourceBusyError = Class.new(SemianError) |
| 20 | + CircuitOpenError = Class.new(SemianError) |
| 21 | + end |
| 22 | + end |
| 23 | +end |
| 24 | + |
| 25 | +module Semian |
| 26 | + module ActiveRecordTrilogyAdapter |
| 27 | + include Semian::Adapter |
| 28 | + |
| 29 | + ResourceBusyError = ::ActiveRecord::ConnectionAdapters::TrilogyAdapter::ResourceBusyError |
| 30 | + CircuitOpenError = ::ActiveRecord::ConnectionAdapters::TrilogyAdapter::CircuitOpenError |
| 31 | + |
| 32 | + attr_reader :raw_semian_options, :semian_identifier |
| 33 | + |
| 34 | + def initialize(*options) |
| 35 | + *, config = options |
| 36 | + @raw_semian_options = config.delete(:semian) |
| 37 | + @semian_identifier = begin |
| 38 | + name = semian_options && semian_options[:name] |
| 39 | + unless name |
| 40 | + host = config[:host] || "localhost" |
| 41 | + port = config[:port] || 3306 |
| 42 | + name = "#{host}:#{port}" |
| 43 | + end |
| 44 | + :"mysql_#{name}" |
| 45 | + end |
| 46 | + super |
| 47 | + end |
| 48 | + |
| 49 | + def execute(sql, name = nil, async: false, allow_retry: false) |
| 50 | + if query_allowlisted?(sql) |
| 51 | + super(sql, name, async: async, allow_retry: allow_retry) |
| 52 | + else |
| 53 | + acquire_semian_resource(adapter: :trilogy_adapter, scope: :execute) do |
| 54 | + super(sql, name, async: async, allow_retry: allow_retry) |
| 55 | + end |
| 56 | + end |
| 57 | + end |
| 58 | + |
| 59 | + def active? |
| 60 | + acquire_semian_resource(adapter: :trilogy_adapter, scope: :ping) do |
| 61 | + super |
| 62 | + end |
| 63 | + rescue ResourceBusyError, CircuitOpenError |
| 64 | + false |
| 65 | + end |
| 66 | + |
| 67 | + def with_resource_timeout(temp_timeout) |
| 68 | + if connection.nil? |
| 69 | + prev_read_timeout = @config[:read_timeout] || 0 |
| 70 | + @config.merge!(read_timeout: temp_timeout) # Create new client with temp_timeout for read timeout |
| 71 | + else |
| 72 | + prev_read_timeout = connection.read_timeout |
| 73 | + connection.read_timeout = temp_timeout |
| 74 | + end |
| 75 | + yield |
| 76 | + ensure |
| 77 | + @config.merge!(read_timeout: prev_read_timeout) |
| 78 | + connection&.read_timeout = prev_read_timeout |
| 79 | + end |
| 80 | + |
| 81 | + private |
| 82 | + |
| 83 | + def acquire_semian_resource(**) |
| 84 | + super |
| 85 | + rescue ActiveRecord::StatementInvalid => error |
| 86 | + if error.cause.is_a?(Trilogy::TimeoutError) |
| 87 | + semian_resource.mark_failed(error) |
| 88 | + error.semian_identifier = semian_identifier |
| 89 | + end |
| 90 | + raise |
| 91 | + end |
| 92 | + |
| 93 | + def resource_exceptions |
| 94 | + [ActiveRecord::ConnectionNotEstablished] |
| 95 | + end |
| 96 | + |
| 97 | + # TODO: share this with Mysql2 |
| 98 | + QUERY_ALLOWLIST = Regexp.union( |
| 99 | + %r{\A(?:/\*.*?\*/)?\s*ROLLBACK}i, |
| 100 | + %r{\A(?:/\*.*?\*/)?\s*COMMIT}i, |
| 101 | + %r{\A(?:/\*.*?\*/)?\s*RELEASE\s+SAVEPOINT}i, |
| 102 | + ) |
| 103 | + |
| 104 | + def query_allowlisted?(sql, *) |
| 105 | + QUERY_ALLOWLIST.match?(sql) |
| 106 | + rescue ArgumentError |
| 107 | + return false unless sql.valid_encoding? |
| 108 | + |
| 109 | + raise |
| 110 | + end |
| 111 | + |
| 112 | + def connect(*args) |
| 113 | + acquire_semian_resource(adapter: :trilogy_adapter, scope: :connection) do |
| 114 | + super |
| 115 | + end |
| 116 | + end |
| 117 | + end |
| 118 | +end |
| 119 | + |
| 120 | +ActiveRecord::ConnectionAdapters::TrilogyAdapter.prepend(Semian::ActiveRecordTrilogyAdapter) |
0 commit comments