monq 0.1.0

data/lib/monq/commands/models.rb

@@ -0,0 +1,95 @@
+# frozen_string_literal: true
+
+require "monq/table"
+
+module Monq
+  module Commands
+    module Models
+      HEADERS = ["Table Name", "Model Class", "Abbr"].freeze
+
+      module_function
+
+      def render(envs, env_regexp = nil, collection_regexp = nil, format: "terminal")
+        envs.filter_map { |env| render_env(env, collection_regexp, format) if match_env?(env, env_regexp) }.join("\n")
+      end
+
+      def render_env(env, collection_regexp, format)
+        Table.new("Environment: #{env.name}") do |table|
+          table.headers = HEADERS
+          Array(env.models).each do |klass|
+            collection = collection_name(klass)
+            next unless match_collection?(klass, collection, collection_regexp)
+
+            table.body << [collection, klass.name, abbreviation(env, klass)]
+          end
+        end.to_terminal(format)
+      end
+
+      def match_collection?(klass, collection, collection_regexp)
+        regexp = coerce_regexp(collection_regexp)
+        return true unless regexp
+
+        collection.to_s =~ regexp || klass.name.to_s =~ regexp
+      end
+
+      def coerce_regexp(value)
+        return nil if value.nil? || value.to_s.empty?
+        return value if value.is_a?(Regexp)
+
+        source = value.to_s
+        if source.start_with?("/") && source.end_with?("/") && source.length > 1
+          Regexp.new(source[1..-2])
+        else
+          Regexp.new(source, Regexp::IGNORECASE)
+        end
+      end
+
+      def match_env?(env, env_regexp)
+        return true if env_regexp.nil? || env_regexp.to_s.empty?
+
+        env.name =~ coerce_regexp(env_regexp)
+      end
+
+      def collection_name(klass)
+        options = klass.storage_options
+        options[:collection] || options["collection"]
+      rescue StandardError
+        nil
+      end
+
+      def abbreviation(env, klass)
+        mod = env.namespace_module
+        mod.constants.each do |name|
+          next if name.to_s == klass.name.to_s.split("::").last
+
+          return name.to_s if mod.const_get(name) == klass
+        end
+        ""
+      rescue StandardError
+        ""
+      end
+
+      def register_pry_command!
+        return if Pry.commands["m"]
+
+        Pry.commands.create_command("m", "List loaded models: m [-e env] [-f format] [collection_regexp]") do
+          def options(opt)
+            opt.on "-e", "--env", "Environment name regexp", argument: true, as: String
+            opt.on "-f", "--format", "Table format: terminal, md, org", argument: true, as: String, default: "terminal"
+          end
+
+          def process
+            output.puts Monq::Commands::Models.render(
+              Monq::App.instance.environments,
+              opts[:env],
+              args.first,
+              format: opts[:format] || "terminal"
+            )
+          end
+        end
+
+        Pry.commands.alias_command "l", "m"
+      end
+    end
+  end
+end

data/lib/monq/commands/reconnect.rb

@@ -0,0 +1,45 @@
+# frozen_string_literal: true
+
+module Monq
+  module Commands
+    module Reconnect
+      module_function
+
+      def run(envs, force: false)
+        raise Error, "sandbox is active; run sandbox-quit first" if envs.any?(&:sandbox?)
+
+        envs.each do |env|
+          reconnect_ssh(env, force: force)
+          if force
+            env.reconnect!
+          else
+            env.reconnect! unless env.ping?
+          end
+        end
+      end
+
+      def reconnect_ssh(env, force:)
+        proxy = env.ssh_proxy
+        return unless proxy
+
+        if force
+          proxy.reconnect! if proxy.respond_to?(:reconnect!)
+        elsif proxy.respond_to?(:reconnect)
+          proxy.reconnect
+        end
+      end
+
+      def register_pry_command!
+        return if Pry.commands["reconnect"]
+
+        Pry.commands.block_command("reconnect", "Reconnect dead clients") do
+          Monq::Commands::Reconnect.run(Monq::App.instance.environments, force: false)
+        end
+
+        Pry.commands.block_command("reconnect!", "Force reconnect clients") do
+          Monq::Commands::Reconnect.run(Monq::App.instance.environments, force: true)
+        end
+      end
+    end
+  end
+end

data/lib/monq/commands/redefine.rb

@@ -0,0 +1,25 @@
+# frozen_string_literal: true
+
+module Monq
+  module Commands
+    module Redefine
+      module_function
+
+      def run(envs)
+        raise Error, "sandbox is active; run sandbox-quit first" if envs.any?(&:sandbox?)
+
+        envs.each(&:redefine!)
+      end
+
+      def register_pry_command!
+        return if Pry.commands["redefine"]
+
+        Pry.commands.block_command("redefine", "Reload model files") do
+          Monq::Commands::Redefine.run(Monq::App.instance.environments)
+        end
+
+        Pry.commands.alias_command "redef", "redefine"
+      end
+    end
+  end
+end

data/lib/monq/commands/sandbox.rb

@@ -0,0 +1,57 @@
+# frozen_string_literal: true
+
+module Monq
+  module Commands
+    module Sandbox
+      module_function
+
+      def enter(envs)
+        return warn("already in sandbox") if envs.any?(&:sandbox?)
+
+        hellos = envs.map { |env| env.client.database.command(hello: 1).documents.first }
+        if hellos.any? { |doc| !transaction_capable?(doc) }
+          raise Error, "current deployment does not support transactions (need replica set or mongos); not entering sandbox"
+        end
+
+        started = []
+        envs.each do |env|
+          session = env.client.start_session
+          session.start_transaction
+          env.bind_session(session)
+          started << [env, session]
+        rescue StandardError
+          started.reverse_each do |e, s|
+            s.abort_transaction
+            s.end_session
+            e.clear_session
+          end
+          raise
+        end
+      end
+
+      def quit(envs)
+        envs.each(&:clear_session)
+      end
+
+      def transaction_capable?(doc)
+        return false if doc.nil?
+
+        set_name = doc["setName"] || doc[:setName]
+        msg = doc["msg"] || doc[:msg]
+        !set_name.nil? || msg == "isdbgrid"
+      end
+
+      def register_pry_command!
+        return if Pry.commands["sandbox-enter"]
+
+        Pry.commands.block_command("sandbox-enter", "Start an all-or-nothing transaction on every environment") do
+          Monq::Commands::Sandbox.enter(Monq::App.instance.environments)
+        end
+
+        Pry.commands.block_command("sandbox-quit", "Abort sandbox transactions and end their sessions") do
+          Monq::Commands::Sandbox.quit(Monq::App.instance.environments)
+        end
+      end
+    end
+  end
+end

data/lib/monq/commands/show_sql.rb

@@ -0,0 +1,92 @@
+# frozen_string_literal: true
+
+require "logger"
+require "json"
+require "monq/multi_io"
+
+module Monq
+  module Commands
+    module ShowSql
+      module Subscriber
+        module_function
+
+        def started(event)
+          write("started", event)
+        end
+
+        def succeeded(event)
+          write("succeeded", event)
+        end
+
+        def failed(event)
+          write("failed", event)
+        end
+
+        def write(phase, event)
+          io = App.log_io
+          return unless io
+
+          payload = {
+            phase: phase,
+            command_name: event.respond_to?(:command_name) ? event.command_name : nil,
+            database: event.respond_to?(:database_name) ? event.database_name : nil,
+            command: jsonable(event.respond_to?(:command) ? event.command : nil),
+            duration: event.respond_to?(:duration) ? event.duration : nil
+          }
+          io.write("#{JSON.generate(payload)}\n")
+        rescue StandardError
+          nil
+        end
+
+        def jsonable(value)
+          JSON.parse(JSON.generate(value))
+        rescue StandardError
+          value.to_s
+        end
+      end
+
+      class << self
+        attr_accessor :subscribed
+
+        def enable
+          App.log_io ||= MultiIO.new
+          App.log_io << $stdout unless App.log_io.include?($stdout)
+          subscribe!
+          mirror_logger!
+        end
+
+        def disable
+          return unless App.log_io.is_a?(MultiIO) && App.log_io.include?($stdout)
+
+          App.log_io.delete($stdout)
+        end
+
+        def subscribe!
+          return if subscribed
+          return unless defined?(Mongo::Monitoring::COMMAND) && Mongo::Monitoring::Global.respond_to?(:subscribe)
+
+          Mongo::Monitoring::Global.subscribe(Mongo::Monitoring::COMMAND, Subscriber)
+          self.subscribed = true
+        end
+
+        def mirror_logger!
+          return unless defined?(Mongo::Logger)
+
+          Mongo::Logger.logger = Logger.new(App.log_io)
+        end
+
+        def register_pry_command!
+          return if Pry.commands["show-sql"]
+
+          Pry.commands.block_command("show-sql", "Print Mongo commands") do
+            Monq::Commands::ShowSql.enable
+          end
+
+          Pry.commands.block_command("hide-sql", "Stop printing Mongo commands") do
+            Monq::Commands::ShowSql.disable
+          end
+        end
+      end
+    end
+  end
+end

data/lib/monq/commands.rb

@@ -0,0 +1,25 @@
+# frozen_string_literal: true
+
+require_relative "commands/info"
+require_relative "commands/models"
+require_relative "commands/show_sql"
+require_relative "commands/reconnect"
+require_relative "commands/redefine"
+require_relative "commands/sandbox"
+
+module Monq
+  module Commands
+    module_function
+
+    def register_pry_commands!
+      return unless defined?(Pry)
+
+      Info.register_pry_command!
+      Models.register_pry_command!
+      ShowSql.register_pry_command!
+      Reconnect.register_pry_command!
+      Redefine.register_pry_command!
+      Sandbox.register_pry_command!
+    end
+  end
+end

data/lib/monq/config.rb

@@ -0,0 +1,216 @@
+# frozen_string_literal: true
+
+require "yaml"
+require "ostruct"
+
+module Monq
+  module Config
+    RESERVED = %w[default arql].freeze
+
+    module_function
+
+    def load_file(path)
+      raw = YAML.safe_load(File.read(path), aliases: true)
+      (raw || {}).transform_keys(&:to_s)
+    rescue ArgumentError
+      raw = YAML.safe_load(File.read(path))
+      (raw || {}).transform_keys(&:to_s)
+    end
+
+    def select_environments(yaml, names, options)
+      yaml = stringify(yaml)
+      if names.nil? || names.empty?
+        return { "default" => cli_env(options) } if sufficient_cli?(options)
+
+        raise Error, "No -e given and CLI has no connection. Pass -e or --uri, or create ~/.monq.d/init.yml"
+      end
+
+      names.each do |name|
+        raise Error, "#{name} is a reserved environment name" if RESERVED.include?(name)
+        raise Error, "Environment #{name} not in config. Known: #{known(yaml).join(', ')}" unless yaml.key?(name)
+      end
+
+      selected = names.to_h { |n| [n, deep_sym(yaml[n])] }
+      apply_cli!(selected, options)
+      assign_namespaces!(selected)
+      assert_unique_namespaces!(selected)
+      selected
+    end
+
+    def known(yaml)
+      yaml.keys - RESERVED
+    end
+
+    def sufficient_cli?(options)
+      opt(options, :uri).to_s != "" || (opt(options, :host).to_s != "" && opt(options, :database).to_s != "")
+    end
+
+    def cli_env(options)
+      {
+        uri: opt(options, :uri),
+        host: opt(options, :host),
+        port: opt(options, :port),
+        database: opt(options, :database),
+        username: opt(options, :username),
+        password: opt(options, :password),
+        auth_source: opt(options, :auth_source)
+      }.compact
+    end
+
+    def apply_cli!(selected, options)
+      overrides = connection_overrides(options)
+      if selected.size > 1 && overrides.any?
+        raise Error, "Connection CLI flags are not allowed with multiple -e environments"
+      end
+      return if selected.size != 1
+
+      env = selected.values.first
+      overrides.each { |k, v| env[k] = v }
+      merge_ssh!(env, options)
+    end
+
+    def connection_overrides(options)
+      {
+        uri: opt(options, :uri),
+        host: opt(options, :host),
+        port: opt(options, :port),
+        database: opt(options, :database),
+        username: opt(options, :username),
+        password: opt(options, :password),
+        auth_source: opt(options, :auth_source)
+      }.compact.reject { |_, v| v.respond_to?(:empty?) && v.empty? }
+    end
+
+    def merge_ssh!(env, options)
+      ssh = opt(options, :ssh)
+      return if ssh.nil? || ssh.empty?
+
+      env[:ssh] = (env[:ssh] || {}).merge(ssh.transform_keys(&:to_sym))
+    end
+
+    def assign_namespaces!(selected)
+      selected.each do |name, env|
+        env[:namespace] ||= camelize(name)
+      end
+    end
+
+    def assert_unique_namespaces!(selected)
+      groups = selected.group_by { |_, env| env[:namespace] }
+      dup = groups.find { |_, pairs| pairs.size > 1 }
+      return unless dup
+
+      raise Error, "Environments share namespace #{dup.first}: #{dup.last.map(&:first).join(', ')}"
+    end
+
+    def camelize(name)
+      name.split(/[^A-Za-z0-9]+/).reject(&:empty?).map { |p| p[0].upcase + p[1..] }.join
+    end
+
+    def stringify(hash)
+      hash.to_h.transform_keys(&:to_s).transform_values { |v| v.is_a?(Hash) ? stringify(v) : v }
+    end
+
+    def deep_sym(obj)
+      case obj
+      when Hash then obj.to_h { |k, v| [k.to_sym, deep_sym(v)] }
+      else obj
+      end
+    end
+
+    def prefix_for(env_name)
+      env_name.split(/[._-]/, 2).first
+    end
+
+    def client_name_for(env_name)
+      env_name.gsub(/[^A-Za-z0-9]/, "_").to_sym
+    end
+
+    def opt(options, key)
+      if options.respond_to?(key)
+        options.public_send(key)
+      elsif options.respond_to?(:[])
+        options[key] || options[key.to_s]
+      end
+    end
+
+    def client_options(env)
+      env = deep_sym(env)
+      ssh = env[:ssh]
+      if ssh && env[:uri].to_s.start_with?("mongodb+srv:")
+        raise Error, "mongodb+srv with ssh is not supported; use a direct host"
+      end
+
+      built = env[:uri].to_s != "" ? from_uri(env) : from_fields(env)
+      return built unless ssh && ssh[:local_port]
+
+      {
+        hosts: ["127.0.0.1:#{ssh[:local_port]}"],
+        database: built[:database],
+        options: without_connect(without_replica_set(built[:options].merge(direct_connection: true)))
+      }
+    end
+
+    def from_uri(env)
+      require "mongo"
+      parsed = Mongo::URI.get(env[:uri])
+      database = env[:database].to_s
+      database = parsed.client_options["database"].to_s if database.empty?
+      raise Error, "database is required (URI path is empty)" if database.empty?
+
+      options = symbolize_options(parsed.uri_options)
+      creds = parsed.credentials
+      options[:user] = creds[:user] if creds[:user]
+      options[:password] = creds[:password] if creds[:password]
+      options = without_replica_set(options) if env[:ssh] && env[:ssh][:local_port]
+      {
+        hosts: parsed.servers,
+        database: database,
+        options: options
+      }
+    end
+
+    def from_fields(env)
+      raise Error, "host and database are required" if env[:host].to_s.empty? || env[:database].to_s.empty?
+
+      port = env[:port] || 27017
+      options = {}
+      options[:user] = env[:username] if env[:username]
+      options[:password] = env[:password] if env[:password]
+      options[:auth_source] = env[:auth_source] if env[:auth_source]
+      options[:replica_set] = env[:replica_set] if env[:replica_set] && !(env[:ssh] && env[:ssh][:local_port])
+      options[:ssl] = env[:ssl] unless env[:ssl].nil?
+      { hosts: ["#{env[:host]}:#{port}"], database: env[:database], options: options }
+    end
+
+    def mask_uri(uri)
+      uri.to_s.sub(%r{://([^:/@]+):[^@]+@}, '://\1:***@')
+    end
+
+    def ssh_forward_target(env)
+      env = deep_sym(env)
+      if env[:uri].to_s != ""
+        require "mongo"
+        host, port = Mongo::URI.get(env[:uri]).servers.first.split(":")
+        { host: host, port: (port || 27017).to_i }
+      else
+        { host: env[:host], port: (env[:port] || 27017).to_i }
+      end
+    end
+
+    def symbolize_options(options)
+      options.to_h { |key, value| [key.to_sym, value] }
+    end
+
+    def without_replica_set(options)
+      options.reject { |key, _| replica_set_key?(key) }
+    end
+
+    def without_connect(options)
+      options.except(:connect, "connect")
+    end
+
+    def replica_set_key?(key)
+      %w[replica_set replicaset].include?(key.to_s.downcase.tr("-", "_"))
+    end
+  end
+end