pgai 0.2.4 → 1.0.0.alpha2

checksums.yaml

@@ -1,7 +1,7 @@
 ---
 SHA256:
-  metadata.gz: 638625e64ce9fa8f81c98f8a618b74f3d9f1e7ae14632bb4994a581011ad4bf3
-  data.tar.gz: 4703d7e923b428ac933f72335af104b373b991d048abef7516248983d8dcfb25
+  metadata.gz: 8a9ded7a8c272d713343af77a8fa74c139d47f0bb478ce90bb8154205f9bed00
+  data.tar.gz: 5a2d6c82be207fce031a814bebe10e0015b08e79be471e7a526ee0272fb3f7e8
 SHA512:
-  metadata.gz: f05388f1e73e3277a06baea0298bafd175a8e50959e3a0fe9580e099820e84bd13e1665c9f8841b0e46f437fea638728a11c510c332026834fa399144394b022
-  data.tar.gz: dd1c62be57dd6fa39a4b9b685f390072eb6e51ebb9dc7a10627c7366ddf89521e9edca73b17def3e178d36108c889466e6ca95a038cb121635092084c7ff8646
+  metadata.gz: 71ec50b2d9622e2c65d50ceca0f6405133409b0aac40537084f97f86f10a88ddbce2696c775671d62ac08910d0c7d5695941510b22bae9413e0f5ec8281b2020
+  data.tar.gz: 78f9e1794a40b34ea0d2da5c892d4cbfb5e5883cfb42851b370e3e23f7ebab49cb23da2eea140966be4983d1cad780513eb98ba163125df16beabc06f1ba2fe5

data/README.md

@@ -25,35 +25,31 @@ An access token will be required and it can be obtained from: https://console.po
 Example:
 
 ```shell
-pgai config --dbname=gitlabhq_dblab --prefix=<gitlab handle> --proxy=<alias> --exe=<optional /path/to/dblab>
+pgai config --prefix=<gitlab handle>
 ```
 
 To configure environments:
 
 ```shell
-pgai env add --alias ci --id ci-database --port 12345
+pgai env add --alias ci --id ci.lab.internal --port 2345 -n gitlabhq_dblab
 ```
 
-The environment id, port, and proxy domain can be found by clicking on the `Connect` button on a database instance page.
+The environment id and port can be found by clicking on the `Connect` button on a database instance page.
 
-Configuring the proxy attributes must be done via the `~/.ssh/config` file. Example:
+The the id attributes must be configured for SSH port forwarding via the `~/.ssh/config` file. Example:
 
-```
-Host <alias>
-  HostName <domain.postgres.ai>
+```shell
+Host <bastion>
   User <username>
   IdentityFile ~/.ssh/id_ed25519
-```
-
-## dblab binary file
-
-This CLI is built around around the [`dblab`](https://postgres.ai/docs/how-to-guides/cli/cli-install-init) CLI and it looks for it at the following locations:
 
-- at the path specified in the config file if specified when running the `pgai config` command
-- in `$PATH`
-- at `~/.dblab/dblab`
 
-If it's not found at any of the specified locations, a copy will be downloaded to `~/.dblab/dblab`.
+Host *.lab.internal
+  User <username>
+  PreferredAuthentications publickey
+  IdentityFile ~/.ssh/id_ed25519
+  ProxyCommand ssh <bastion> -W %h:%p
+```
 
 ## Usage
 
@@ -63,10 +59,10 @@ pgai connect <env alias>
 
 Multiple `connect` commands for an environment will connect to the same clone, it won't start a new one.
 
-`pgai info <env alias>` prints out a database URL variable that can be used from Rails to integrate with dblab. Example for `CI` database:
+`pgai use --only <env alias> -- command` can be used to run commands with access to a clone. Example for connecting a Rails console to the `CI` database:
 
 ```shell
-CI_DATABASE_URL="postgresql://foo:bar@localhost:9000/foo_test" bin/rails c -e test
+pgai use -o ci -v -- bin/rails c -e test
 ```
 
 ### Features
@@ -76,6 +72,14 @@ CI_DATABASE_URL="postgresql://foo:bar@localhost:9000/foo_test" bin/rails c -e te
 - automatic port forward management
 - prevents system sleep while psql sessions are active via `caffeinate`
 
+## Upgrading from version 0
+
+- Update and test the ssh config file
+- `~/.dblab/` and its contents can be removed
+- `~/.config/pgai/config.pstore` should be removed and recreated with `pgai config`
+- The environment ID serves as the proxy host and the database name can be configured per environment.
+- Environments can share the same remote port value.
+
 ## Contributing
 
 Bug reports and pull requests are welcome on GitLab at https://gitlab.com/mbobin/pgai.

data/lib/pgai/cli/base.rb

@@ -2,20 +2,23 @@ require "thor"
 
 module Pgai::Cli
   class Base < Thor
-    def self.exit_on_failure?
-      true
-    end
+    def self.exit_on_failure? = true
 
-    private
+    check_unknown_options!
+
+    class_option :verbose, type: :boolean, aliases: "-v", desc: "Detailed logging"
+    class_option :quiet, type: :boolean, aliases: "-q", desc: "Minimal logging"
 
-    def configuration
-      @configuration ||= load_and_validate_config!
+    def initialize(*)
+      super
+
+      Pgai::Commander.configure(options_with_subcommand_class_options)
     end
 
-    def load_and_validate_config!
-      Pgai::Config.load.tap do |config|
-        raise "Access token is not set" unless config.access_token
-      end
+    private
+
+    def options_with_subcommand_class_options
+      options.merge(@_initializer.last[:class_options] || {})
     end
   end
 end

data/lib/pgai/cli/env.rb

@@ -1,22 +1,42 @@
 module Pgai::Cli
   class Env < Base
     desc "add", "Add new environment"
-    method_option :alias, aliases: "-a", desc: "This will be used internally for the connect command", required: true
-    method_option :id, aliases: "-i", desc: "Environment id from postgres.ai", required: true
-    method_option :port, aliases: "-p", desc: "Environment port from postgres.ai", required: true
+    method_option :alias, aliases: "-a",
+      desc: "This will be used internally for the connect command",
+      required: true
+    method_option :id,
+      aliases: "-i",
+      desc: "Environment id from postgres.ai, example: lab-ci.db-lab.internal",
+      required: true
+    method_option :port,
+      aliases: "-p",
+      desc: "Environment port from postgres.ai",
+      required: true
+    method_option :dbname,
+      aliases: "-n",
+      desc: "Specify database name to connect to by default",
+      required: true,
+      default: "postgres"
     def add
-      configuration.add_env(**options)
+      Pgai::Resources::Local::Environment.new(options).save
     end
 
     desc "remove", "Remove the specified environment from the config"
-    method_option :alias, aliases: "-a", desc: "This will be used internally for the connect command", required: true
+    method_option :alias,
+      aliases: "-a",
+      desc: "The alias used for registering the environment",
+      required: true
     def remove
-      configuration.remove_env(options[:alias])
+      Pgai::Resources::Local::Environment.delete(options[:alias])
     end
 
     desc "list", "List all configured environments"
     def list
-      puts JSON.pretty_generate(configuration.enviroments)
+      data = Pgai::Resources::Local::Environment.all.map do |env|
+        env.attributes.slice(:id, :alias, :port, :dbname)
+      end
+
+      say JSON.pretty_generate(data)
     end
   end
 end

data/lib/pgai/cli/main.rb

@@ -2,60 +2,50 @@ module Pgai::Cli
   class Main < Base
     desc "config", "Configure CLI options"
     method_option :prefix, aliases: "-p", desc: "Specify prefix name to be used for clones", required: true
-    method_option :dbname, aliases: "-n", desc: "Specify database name to connect to by default", required: true, default: "postgres"
-    method_option :proxy, aliases: "-x", desc: "Specify proxy host name", required: true
-    method_option :exe, aliases: "-e", desc: "dblab binary path"
     def config
       token = ask("Access token:", echo: false)
 
-      Pgai::Config.persist(options.merge(token: token))
+      Pgai::Resources::Local::Configuration
+        .new(clone_prefix: options[:prefix], access_token: token)
+        .save
     end
 
     desc "connect database-name", "Create and connect to a thin clone database"
     def connect(name)
-      env = configuration.enviroments.fetch(name)
-
-      Pgai::CloneManager.new(env, config: configuration).connect
+      with_env(name) do |env|
+        Pgai::CloneManager.new(env).connect
+      end
     end
 
     desc "destroy database-name", "Remove the thin clone and all port forwards"
     def destroy(name)
-      env = configuration.enviroments.fetch(name)
-
-      Pgai::CloneManager.new(env, config: configuration).cleanup
+      with_env(name) do |env|
+        Pgai::CloneManager.new(env).cleanup
+      end
     end
 
     desc "reset database-name", "Reset clone's state"
     def reset(name)
-      env = configuration.enviroments.fetch(name)
-
-      Pgai::CloneManager.new(env, config: configuration).reset
-    end
-
-    desc "info database-name", "Show clone details"
-    def info(name)
-      env = configuration.enviroments.fetch(name)
-
-      data = Pgai::CloneManager.new(env, config: configuration).info
-      print_table data
+      with_env(name) do |env|
+        Pgai::CloneManager.new(env).reset
+      end
     end
 
     desc "use", "Execute the given command with DATABASE_URLs"
     method_option :only, aliases: "-o", desc: "Export only this database connection", repeatable: true
     def use(*command)
-      envs = options.fetch(:only) { configuration.enviroments.keys }
-
-      vars = envs.each_with_object({}) do |env, acc|
-        environment = configuration.enviroments.fetch(env)
-        clone = Pgai::CloneManager.new(environment, config: configuration).clone
-
-        acc["#{env.to_s.upcase}_DATABASE_URL"] = clone.database_url
-      end
+      envs = options.fetch(:only) { Pgai::Resources::Local::Environment.all.map(&:alias) }
 
-      exec(vars, *command)
+      Pgai::ExternalCommandManager.new(envs, command).run
     end
 
     desc "env", "Manage environments"
     subcommand "env", Pgai::Cli::Env
+
+    private
+
+    def with_env(name, &block)
+      Pgai::Commander.instance.with_env(name, &block)
+    end
   end
 end

data/lib/pgai/client.rb

@@ -0,0 +1,69 @@
+# frozen_string_literal: true
+
+require "excon"
+
+module Pgai
+  class Client
+    def initialize(token:, host: "http://127.0.0.1:2355")
+      @client = Excon.new(host, headers: {"Verification-Token" => token})
+    end
+
+    def version
+      get("healthz").fetch(:version)
+    end
+
+    def expected_cloning_time
+      get("status").dig(:cloning, :expected_cloning_time)
+    end
+
+    def get(path)
+      response = client.get(path: path)
+      parse_reponse(response.body)
+    end
+
+    def post(path, body = nil)
+      response = if body
+        client.post(path: path, body: JSON.dump(body))
+      else
+        client.post(path: path)
+      end
+      return {} if response.body.empty?
+
+      parse_reponse(response.body)
+    end
+
+    def delete(path)
+      response = client.delete(path: path)
+      return {} if response.body.empty?
+
+      parse_reponse(response.body)
+    end
+
+    private
+
+    attr_reader :client
+
+    def parse_reponse(data)
+      data = JSON.parse(data)
+      transform!(data)
+      data
+    end
+
+    def transform_hash!(data)
+      data.transform_keys! { |key| key.to_s.gsub(/(?<=[A-Z])(?=[A-Z][a-z])|(?<=[a-z\d])(?=[A-Z])/, "_").downcase.to_sym }
+      data.each_value { |value| transform!(value) }
+    end
+
+    def transform_array!(data)
+      data.each { transform!(_1) }
+    end
+
+    def transform!(data)
+      if data.is_a?(Hash)
+        transform_hash!(data)
+      elsif data.is_a?(Array)
+        transform_array!(data)
+      end
+    end
+  end
+end

data/lib/pgai/clone_manager.rb

@@ -1,151 +1,81 @@
 # frozen_string_literal: true
 
-require "securerandom"
-
 module Pgai
   class CloneManager
-    HOSTNAME = "127.0.0.1"
-
-    def initialize(environment, config:)
-      @environment = environment
-      @config = config
-      @port_forward = PortForward.new(config: config, hostname: HOSTNAME)
-      @dblab = Dblab.new(config: config, hostname: HOSTNAME)
+    def initialize(env)
+      @env = env
+      @logger = Pgai::Commander.instance.logger
     end
 
     def connect
-      configure_enviroment
+      prepare do |cached_clone|
+        PsqlManager.new(
+          cached_clone,
+          logger: logger
+        ).run
+      end
+    end
 
-      psql find_or_create_clone
+    def prepare
+      find_or_create_clone.with_port_forward do |cached_clone|
+        yield cached_clone
+      end
     end
 
     def cleanup
-      configure_enviroment
-      return unless find_raw_clone
+      return unless clone_already_created?
 
-      dblab.destroy_clone(id: clone_id)
-      config.remove_clone(clone_id)
-      port_forward.stop
+      Resources::Local::Clone.delete(clone_id)
+      Resources::Remote::Clone.find(clone_id).tap do |resource|
+        resource.delete
+        resource.refresh
+        logger.info resource.status_message
+      end
     end
 
     def reset
-      configure_enviroment
-      return unless find_raw_clone
-
-      dblab.reset_clone(id: clone_id)
-    end
-
-    def info
-      configure_enviroment
-
-      clone = find_clone
-      return {} unless clone
+      return unless clone_already_created?
 
-      database_url = "#{enviroment_alias.to_s.upcase}_DATABASE_URL='#{clone.database_url}'"
-
-      {
-        psql_connection_string: clone.connection_string,
-        rails_database_url: database_url,
-        created_at: clone.created_at,
-        data_state_at: clone.data_state_at
-      }
-    end
-
-    def clone
-      configure_enviroment
-      clone = find_or_create_clone
-      port_forward.start(clone.port)
-      clone
+      Resources::Remote::Clone.find(clone_id).tap do |resource|
+        resource.reset
+        logger.info resource.status_message
+      end
     end
 
     private
 
-    attr_reader :environment, :config, :port_forward, :dblab
+    attr_reader :env, :logger
 
-    def configure_enviroment
-      port_forward.start(enviroment_port)
-      dblab.configure_env(port: enviroment_port, token: config.access_token, id: environment_id)
+    def clone_id
+      env.clone_id
     end
 
     def find_or_create_clone
-      find_clone || create_clone
+      find_cached_clone || create_clone
     end
 
-    def find_clone
-      raw_clone = find_raw_clone
-
-      if raw_clone
-        config.find_clone(clone_id) || raise(Pgai::CloneNotFount, "clone not tracked")
+    def find_cached_clone
+      if clone_already_created?
+        Resources::Local::Clone.find(clone_id) || raise_unknown_clone
       else
-        config.remove_clone(clone_id) && nil
+        Resources::Local::Clone.delete(clone_id) && nil
       end
     end
 
-    def find_raw_clone
-      dblab.list_clones.find { |clone| clone["id"] == clone_id }
-    end
-
     def create_clone
-      raw_clone = dblab.create_clone(id: clone_id, user: clone_user, password: clone_password)
-
-      attributes = {
-        port: raw_clone.dig("db", "port"),
-        password: clone_password,
-        user: clone_user,
-        host: HOSTNAME,
-        dbname: config.dbname,
-        created_at: raw_clone["createdAt"],
-        data_state_at: raw_clone.dig("snapshot", "dataStateAt")
-      }
-
-      config.persist_clone(clone_id, attributes)
-    end
-
-    def psql(clone)
-      puts "Data state at: #{clone.data_state_at}"
-      port_forward.start(clone.port)
-
-      psql_pid = fork do
-        exec("psql #{clone.connection_string}")
-      end
-
-      Signal.trap("INT") {}
-      start_caffeinate(psql_pid)
-      Process.wait(psql_pid)
-    ensure
-      port_forward.conditionally_stop(clone.port, [psql_pid])
-      port_forward.conditionally_stop(enviroment_port)
-    end
-
-    def environment_id
-      environment.fetch(:id)
-    end
-
-    def enviroment_port
-      environment.fetch(:port)
+      CreateCloneService.new(clone_id, env: env, logger: logger).execute
     end
 
-    def enviroment_alias
-      environment.fetch(:alias)
+    def raise_unknown_clone
+      raise(Pgai::ResourceNotFound, <<~MESSAGE)
+        Unknown clone. Remove it and retry.
+      MESSAGE
     end
 
-    def clone_id
-      @clone_id ||= [config.clone_prefix, enviroment_alias].join("_")
-    end
-
-    def clone_user
-      @clone_user ||= SecureRandom.hex(8)
-    end
-
-    def clone_password
-      @clone_password ||= SecureRandom.hex(16)
-    end
-
-    def start_caffeinate(pid)
-      return if `which caffeinate`.to_s.empty?
-
-      caffeinate_pid = Process.spawn("caffeinate -is -w #{pid}")
-      Process.detach(caffeinate_pid)
+    def clone_already_created?
+      !!Resources::Remote::Clone.find(clone_id)
+    rescue Pgai::ResourceNotFound
+      false
     end
   end
 end

data/lib/pgai/commander.rb

@@ -0,0 +1,74 @@
+# frozen_string_literal: true
+
+require "tty-logger"
+require "singleton"
+require "forwardable"
+
+module Pgai
+  class Commander
+    include Singleton
+    extend Forwardable
+
+    attr_reader :verbosity
+    def_delegators :logger, *TTY::Logger::LOG_TYPES.keys
+
+    def self.configure(options)
+      instance.tap do |commander|
+        if options[:verbose]
+          commander.verbosity = :debug
+        end
+
+        if options[:quiet]
+          commander.verbosity = :error
+        end
+      end
+    end
+
+    def initialize
+      self.verbosity = :info
+    end
+
+    def store
+      @store ||= Pgai::Store.new
+    end
+
+    def config
+      @config ||= Pgai::Resources::Local::Configuration.default
+    end
+
+    def configure
+      yield self
+    end
+
+    def verbosity=(verbosity)
+      @verbosity = verbosity
+
+      TTY::Logger.configure do |config|
+        config.level = verbosity
+      end
+    end
+
+    def logger
+      @logger ||= TTY::Logger.new
+    end
+
+    def port_manager
+      @port_manager ||= Port::Manager.new(logger: logger)
+    end
+
+    def check_access_token_presence!
+      raise "Access token is not set" unless config&.access_token
+    end
+
+    def with_env(env_name, &block)
+      check_access_token_presence!
+
+      env = Resources::Local::Environment.find(env_name) do |env|
+        env.access_token = config.access_token
+        env.clone_prefix = config.clone_prefix
+      end
+
+      env.prepare(&block)
+    end
+  end
+end

data/lib/pgai/create_clone_service.rb

@@ -0,0 +1,75 @@
+# frozen_string_literal: true
+
+require "tty-progressbar"
+
+module Pgai
+  class CreateCloneService
+    HOSTNAME = "127.0.0.1"
+    CLONE_REFRESH_INTERVAL = 0.2
+
+    attr_reader :id, :env, :logger
+
+    def initialize(id, env:, logger:)
+      @id = id
+      @env = env
+      @logger = logger
+    end
+
+    def execute
+      create_clone_resource
+      logger.debug(clone_resource.status_message)
+      progress_bar.advance
+      wait_for_clone_to_be_ready
+      logger.debug(clone_resource.status_message)
+
+      Resources::Local::Clone.new(cached_clone_attributes).tap(&:save)
+    end
+
+    def create_clone_resource
+      @clone_resource = Resources::Remote::Clone.new(id: id).tap do |resource|
+        resource.db_object = Resources::Remote::DbObject.new(db_name: env.dbname)
+        resource.snapshot = Resources::Remote::Snapshot.latest
+        resource.save
+      end
+    end
+
+    def cached_clone_attributes
+      {
+        id: id,
+        port: clone_resource.db_object.port,
+        password: clone_resource.db_object.password,
+        user: clone_resource.db_object.username,
+        host: HOSTNAME,
+        remote_host: env.id,
+        dbname: env.dbname,
+        created_at: clone_resource.created_at,
+        data_state_at: clone_resource.snapshot.data_state_at
+      }
+    end
+
+    def clone_resource
+      @clone_resource ||= create_clone_resource
+    end
+
+    def wait_for_clone_to_be_ready
+      loop do
+        clone_resource.refresh
+        progress_bar.advance
+        break if clone_resource.ready?
+        sleep(CLONE_REFRESH_INTERVAL)
+      end
+      progress_bar.finish
+    end
+
+    def progress_bar
+      @progress_bar ||= TTY::ProgressBar.new(
+        "Preparing clone ... [:bar]",
+        total: progress_bar_steps
+      )
+    end
+
+    def progress_bar_steps
+      (env.expected_cloning_time / CLONE_REFRESH_INTERVAL).ceil + 1
+    end
+  end
+end