pgai 0.2.4 → 1.0.0

data/lib/pgai/external_command_manager.rb

@@ -0,0 +1,41 @@
+# frozen_string_literal: true
+
+module Pgai
+  class ExternalCommandManager
+    def initialize(env_names, command, logger: nil)
+      @env_names = env_names
+      @command = command
+      @logger = logger || Pgai::Commander.instance.logger
+    end
+
+    def run
+      Signal.trap("INT", "IGNORE")
+
+      prepare(env_names)
+    end
+
+    private
+
+    attr_reader :env_names, :command, :logger
+
+    def prepare(envs, variables = {})
+      if (env_name = envs.pop)
+        Pgai::Commander.instance.with_env(env_name) do |env|
+          Pgai::CloneManager.new(env).prepare do |cached_clone|
+            variables["#{env_name.to_s.upcase}_DATABASE_URL"] = cached_clone.database_url
+            prepare(envs, variables)
+          end
+        end
+      else
+        execute(variables)
+      end
+    end
+
+    def execute(variables)
+      pid = fork do
+        exec(variables, *command)
+      end
+      Process.wait(pid)
+    end
+  end
+end

data/lib/pgai/port/allocator.rb

@@ -0,0 +1,43 @@
+# frozen_string_literal: true
+
+require "socket"
+
+module Pgai
+  module Port
+    class Allocator
+      START_PORT = 5000
+      FINISH_PORT = 9000
+      LOCALHOST = "127.0.0.1"
+
+      def self.pick
+        new.pick
+      end
+
+      def initialize(start: START_PORT, finish: FINISH_PORT)
+        @start = start
+        @finish = finish
+      end
+
+      def pick
+        (start + jitter).upto(finish) do |port|
+          return port if available?(port)
+        end
+      end
+
+      private
+
+      attr_reader :start, :finish
+
+      def jitter
+        rand((finish - start) / 2).floor
+      end
+
+      def available?(port)
+        TCPServer.new(LOCALHOST, port).close
+        true
+      rescue Errno::EADDRINUSE
+        false
+      end
+    end
+  end
+end

data/lib/pgai/port/forwarder.rb

@@ -0,0 +1,83 @@
+# frozen_string_literal: true
+
+require "socket"
+require "net/ssh"
+
+module Pgai
+  module Port
+    class Forwarder
+      LOCALHOST = "127.0.0.1"
+
+      attr_reader :local_port, :remote_host, :remote_port, :logger
+
+      def initialize(local_port, remote_host, remote_port, logger: Pgai::Commander.instance.logger)
+        @local_port = local_port
+        @remote_host = remote_host
+        @remote_port = remote_port
+        @logger = logger
+        @child, @parent = Socket.pair(:UNIX, :DGRAM)
+      end
+
+      def start
+        return if ready?
+
+        debug "starting"
+        start_ssh_connection
+        wait_until_ready
+        debug "ready to accept connections"
+      end
+
+      def stop
+        return unless @pid
+
+        @parent.write("exit")
+        debug "shutting down"
+        Process.wait(@pid)
+        debug "exit"
+        @pid = nil
+      end
+
+      private
+
+      def start_ssh_connection
+        @pid = Process.fork do
+          Signal.trap("INT", "IGNORE")
+          @parent.close
+
+          Net::SSH.start(remote_host) do |ssh|
+            ssh.forward.local(local_port, LOCALHOST, remote_port)
+            ssh.loop(0.1) { keep_running? }
+          end
+        end
+      end
+
+      def keep_running?
+        @child.read_nonblock(10).empty?.tap do |value|
+          debug "shutdown command received"
+        end
+      rescue
+        true
+      end
+
+      def ready?
+        !!TCPSocket.new(LOCALHOST, local_port)
+      rescue Errno::ECONNREFUSED
+        false
+      end
+
+      def wait_until_ready
+        until ready?
+          sleep 0.02
+        end
+      end
+
+      def debug(message)
+        logger.debug "#{logger_inspect} #{message}"
+      end
+
+      def logger_inspect
+        "[#{self.class}(#{local_port}:#{remote_host}:#{remote_port})]:"
+      end
+    end
+  end
+end

data/lib/pgai/port/manager.rb

@@ -0,0 +1,27 @@
+# frozen_string_literal: true
+
+require "concurrent-ruby"
+
+module Pgai
+  module Port
+    class Manager
+      def initialize(logger: Pgai::Commander.instance.logger)
+        @logger = logger
+        @executor = Concurrent::IndirectImmediateExecutor.new
+      end
+
+      def forward(local_port, host, port)
+        forwarder = start(local_port, host, port)
+        yield
+      ensure
+        forwarder&.stop
+      end
+
+      def start(local_port, host, port)
+        fw = Port::Forwarder.new(local_port, host, port, logger: @logger)
+        @executor.post(fw, &:start)
+        fw
+      end
+    end
+  end
+end

data/lib/pgai/psql_manager.rb

@@ -0,0 +1,37 @@
+# frozen_string_literal: true
+
+module Pgai
+  class PsqlManager
+    def initialize(cached_clone, logger:)
+      @cached_clone = cached_clone
+      @logger = logger
+    end
+
+    def run
+      Signal.trap("INT", "IGNORE")
+      logger.info "Data state at: #{cached_clone.data_state_at}"
+
+      psql
+    end
+
+    private
+
+    attr_reader :cached_clone, :logger
+
+    def psql
+      psql_pid = fork do
+        exec("psql #{cached_clone.connection_string}")
+      end
+
+      start_caffeinate(psql_pid)
+      Process.wait(psql_pid)
+    end
+
+    def start_caffeinate(pid)
+      return if `which caffeinate`.to_s.empty?
+
+      caffeinate_pid = Process.spawn("caffeinate -is -w #{pid}")
+      Process.detach(caffeinate_pid)
+    end
+  end
+end

data/lib/pgai/resources/attributes.rb

@@ -0,0 +1,109 @@
+# frozen_string_literal: true
+
+require "date"
+
+module Pgai
+  module Resources
+    module Attributes
+      def self.included(base)
+        base.extend ClassMethods
+        attr_reader :attributes
+      end
+
+      class Attribute
+        FALSE_VALUES = [
+          "0", "f", "F", "false", "FALSE",
+          "off", "OFF", "NO", "no"
+        ].to_set.freeze
+
+        TYPES = {
+          string: ->(value) { value.to_s },
+          integer: ->(value) { value.to_i },
+          decimal: ->(value) { value.to_f },
+          boolean: ->(value) { !FALSE_VALUES.include?(value.to_s) },
+          datetime: ->(value) { ::Time.at(::Time.new(value.to_s, in: "UTC")) }
+        }
+
+        attr_reader :name
+        attr_reader :type
+        attr_reader :default
+
+        def initialize(name, type, default)
+          @name = name
+          @type = type
+          @default = default
+
+          yield self if block_given?
+        end
+
+        def cast(value)
+          TYPES.fetch(type).call(value)
+        end
+
+        def cast_or_default(user_value)
+          value = default
+          value = cast(user_value) unless user_value.nil?
+          value = value.call if value.respond_to?(:call)
+          value
+        end
+      end
+
+      module ClassMethods
+        def inherited(subclass)
+          attributes.each do |attr|
+            subclass.attribute attr.name, attr.type, default: attr.default
+          end
+        end
+
+        def attributes
+          @attributes ||= []
+        end
+
+        def attribute(name, type, default: nil)
+          Attribute.new(name, type, default) do |attribute|
+            attributes << attribute
+            generate_attribute_methods(attribute)
+          end
+        end
+
+        def generate_attribute_methods(attribute)
+          define_method(attribute.name) do
+            instance_variable_get(:"@#{attribute.name}")
+          end
+
+          define_method("#{attribute.name}=") do |value|
+            instance_variable_set(:"@#{attribute.name}", value)
+            attributes[attribute.name.to_sym] = value
+          end
+        end
+      end
+
+      def initialize(user_attributes = {})
+        @attributes = {}
+
+        self.class.attributes.each do |attribute|
+          assign_attribute(attribute, user_attributes[attribute.name])
+        end
+      end
+
+      def assign_attribute(attribute, user_value)
+        value = attribute.cast_or_default(user_value)
+
+        public_send(:"#{attribute.name}=", value)
+      end
+
+      def assign_attributes(user_attributes)
+        user_attributes.each do |key, value|
+          attribute = self.class.attributes.find { |attr| attr.name == key }
+          raise "shit" unless attribute
+
+          assign_attribute(attribute, value)
+        end
+      end
+
+      def has_attribute?(name)
+        respond_to?(name) && respond_to?(:"#{name}=")
+      end
+    end
+  end
+end

data/lib/pgai/resources/local/base_record.rb

@@ -0,0 +1,65 @@
+# frozen_string_literal: true
+
+require "forwardable"
+
+module Pgai
+  module Resources
+    module Local
+      class BaseRecord
+        include Attributes
+        extend Forwardable
+
+        def_delegators :klass, :backend, :record_type
+
+        class << self
+          def all
+            backend.all(record_type).map do |attributes|
+              new(attributes)
+            end
+          end
+
+          def find(id)
+            attributes = backend.find(record_type, id)
+            return unless attributes
+
+            record = new(attributes)
+            yield(record) if block_given?
+            record
+          end
+
+          def delete(id)
+            backend.delete(record_type, id)
+            true
+          end
+
+          def record_type
+            :"#{name.split("::").last.downcase}s"
+          end
+
+          def backend
+            Pgai::Commander.instance.store
+          end
+        end
+
+        def save
+          backend.save(record_type, attributes, key: store_key)
+          true
+        end
+
+        def delete
+          klass.delete(id)
+        end
+
+        def klass
+          self.class
+        end
+
+        private
+
+        def store_key
+          :id
+        end
+      end
+    end
+  end
+end

data/lib/pgai/resources/local/clone.rb

@@ -0,0 +1,39 @@
+# frozen_string_literal: true
+
+module Pgai
+  module Resources
+    module Local
+      class Clone < BaseRecord
+        attribute :id, :string
+        attribute :host, :string
+        attribute :remote_host, :string
+        attribute :port, :integer
+        attribute :password, :string
+        attribute :user, :string
+        attribute :dbname, :string
+        attribute :created_at, :datetime
+        attribute :data_state_at, :datetime
+
+        def connection_string
+          "'host=#{host} port=#{local_port} user=#{user} dbname=#{dbname} password=#{password}'"
+        end
+
+        def database_url
+          "postgresql://#{user}:#{password}@#{host}:#{local_port}/#{dbname}"
+        end
+
+        def with_port_forward(manager = Pgai::Commander.instance.port_manager)
+          manager.forward(local_port, remote_host, port) do
+            yield self
+          end
+        end
+
+        private
+
+        def local_port
+          @local_port ||= Port::Allocator.pick
+        end
+      end
+    end
+  end
+end

data/lib/pgai/resources/local/configuration.rb

@@ -0,0 +1,17 @@
+# frozen_string_literal: true
+
+module Pgai
+  module Resources
+    module Local
+      class Configuration < BaseRecord
+        attribute :id, :string, default: "config"
+        attribute :clone_prefix, :string
+        attribute :access_token, :string
+
+        def self.default
+          find("config")
+        end
+      end
+    end
+  end
+end

data/lib/pgai/resources/local/environment.rb

@@ -0,0 +1,55 @@
+# frozen_string_literal: true
+
+module Pgai
+  module Resources
+    module Local
+      class Environment < BaseRecord
+        attribute :id, :string
+        attribute :alias, :string
+        attribute :port, :integer
+        attribute :dbname, :string
+        attribute :access_token, :string
+        attribute :clone_prefix, :string
+
+        def client
+          @client ||= Pgai::Client.new(
+            token: access_token,
+            host: "http://127.0.0.1:#{local_port}"
+          )
+        end
+
+        def prepare
+          with_port_forward do
+            Resources::Remote::BaseRecord.with_client(client) do
+              yield(self)
+            end
+          end
+        end
+
+        def with_port_forward(manager = Pgai::Commander.instance.port_manager)
+          manager.forward(local_port, id, port) do
+            yield self
+          end
+        end
+
+        def expected_cloning_time
+          @expected_cloning_time ||= client.expected_cloning_time
+        end
+
+        def local_port
+          @local_port ||= Port::Allocator.pick
+        end
+
+        def clone_id
+          @clone_id ||= "#{clone_prefix}_#{self.alias}"
+        end
+
+        private
+
+        def store_key
+          :alias
+        end
+      end
+    end
+  end
+end

data/lib/pgai/resources/remote/base_record.rb

@@ -0,0 +1,37 @@
+# frozen_string_literal: true
+
+module Pgai
+  module Resources
+    module Remote
+      class BaseRecord
+        include Attributes
+
+        class << self
+          def client
+            @@client ||= nil
+          end
+
+          def client=(value)
+            @@client = value
+          end
+
+          def with_client(new_client)
+            previous = client
+            self.client = new_client
+            yield
+          ensure
+            self.client = previous
+          end
+        end
+
+        def client
+          self.class.client
+        end
+
+        def refresh_attributes(data)
+          assign_attributes(data)
+        end
+      end
+    end
+  end
+end

data/lib/pgai/resources/remote/clone.rb

@@ -0,0 +1,114 @@
+# frozen_string_literal: true
+
+module Pgai
+  module Resources
+    module Remote
+      class Clone < BaseRecord
+        attribute :id, :string
+        attribute :protected, :boolean, default: false
+        attribute :created_at, :datetime
+        attribute :delete_at, :datetime
+
+        attr_accessor :db_object
+        attr_accessor :snapshot
+        attr_accessor :status
+        attr_accessor :metadata
+
+        class << self
+          def find(id)
+            data = client.get(File.join(path, id))
+            raise Pgai::ResourceNotFound unless data.key?(:id)
+
+            new(data.slice(:id, :protected)) do |resource|
+              resource.refresh_attributes(data)
+            end
+          end
+
+          def path
+            "clone"
+          end
+        end
+
+        def save
+          body = {
+            id: id,
+            snapshot: {
+              id: snapshot.id
+            },
+            protected: protected,
+            db: {
+              username: db_object.username,
+              password: db_object.password,
+              restricted: db_object.restricted
+            }
+          }
+
+          data = client.post(self.class.path, body)
+          refresh_attributes(data)
+
+          self
+        end
+
+        def ready?
+          status.code == "OK"
+        end
+
+        def refresh
+          data = client.get(object_path)
+          raise Pgai::ResourceNotFound unless data.key?(:id)
+
+          refresh_attributes(data)
+
+          self
+        end
+
+        def delete
+          data = client.delete(object_path)
+          refresh_attributes(data) unless data.empty?
+
+          self
+        end
+
+        def reset
+          client.post(object_path("reset"))
+
+          self
+        end
+
+        def refresh_attributes(data)
+          assign_attributes(data.slice(:delete_at, :created_at))
+          ensure_metadata.refresh_attributes(data.fetch(:metadata))
+          ensure_status.refresh_attributes(data.fetch(:status))
+          ensure_db_object.refresh_attributes(data.fetch(:db))
+          ensure_snapshot.refresh_attributes(data.fetch(:snapshot))
+        end
+
+        def status_message
+          status&.message
+        end
+
+        private
+
+        def object_path(*fragments)
+          File.join(self.class.path, id, *fragments)
+        end
+
+        def ensure_metadata
+          metadata || self.metadata = CloneMetadata.new
+        end
+
+        def ensure_status
+          status || self.status = CloneStatus.new
+        end
+
+        def ensure_db_object
+          db_object || self.db_object = DbObject.new
+        end
+
+        def ensure_snapshot
+          snapshot || self.snapshot = Snapshot.new
+        end
+      end
+    end
+  end
+end

data/lib/pgai/resources/remote/clone_metadata.rb

@@ -0,0 +1,14 @@
+# frozen_string_literal: true
+
+module Pgai
+  module Resources
+    module Remote
+      class CloneMetadata < BaseRecord
+        attribute :clone_diff_size, :integer
+        attribute :logical_size, :integer
+        attribute :cloning_time, :decimal
+        attribute :max_idle_minutes, :integer
+      end
+    end
+  end
+end

data/lib/pgai/resources/remote/clone_status.rb

@@ -0,0 +1,12 @@
+# frozen_string_literal: true
+
+module Pgai
+  module Resources
+    module Remote
+      class CloneStatus < BaseRecord
+        attribute :code, :string
+        attribute :message, :string
+      end
+    end
+  end
+end