pgai 0.1.3 → 0.1.5

checksums.yaml

@@ -1,7 +1,7 @@
 ---
 SHA256:
-  metadata.gz: 4ffc13ccc5cc54a852ba48b245a1f4e470696d622fe499aa5fbce8221e77fa5d
-  data.tar.gz: 6e78ffd141aed55d189bab7c986c11434b8d22fbe4c40b032ddc2074875909fb
+  metadata.gz: a31b2fc8c0e25d189cd3f3191a8e1833de2cbe8f21dfe5e4322793fea1aa4efd
+  data.tar.gz: 6061e6095d92912eb7d5a2f64fe95740314494ff23aefab32afe4e9b15d77d3b
 SHA512:
-  metadata.gz: f5457effb587f8f4aeb75b5887cc15c977c736a37459de877ab18c18d8c0fd76e215300f7d46322e56913c398c2e869ec0c461c992aa231ab3b38e3b8bf45b9e
-  data.tar.gz: 4f29dc64c0553db229c009917b2676e9374ddb16b82c79e608588a9be0a655d10e804b2d8ec0566030a6d732a7ccb53c4681e3f53f65a3cb9ba4fcc59e0b3412
+  metadata.gz: ee38e79b9be45e6a848659a830550ce0cf7ad3cf272fc2c6d6e9ddb95214595ae9fd506a4726d3fffe36719e56a4d4bf439ff8543099ee457feed13e5a6ebbc6
+  data.tar.gz: b3d34902baa84fc445e10d1874455b3168fe3896c78fd83d41b84c24d78ee1bdc9ccf02282b984d848cdf2984bc8f4ca136b43ff367c9a0857857d164b554358

data/README.md

@@ -22,10 +22,24 @@ Before usage `pgai config` must be executed and at least an environment must be
 
 An access token will be required and it can be obtained from: https://console.postgres.ai/gitlab/tokens
 
-Configuring the user and identity file for the proxy must be done using the `~/.ssh/config` file. Example:
+Example:
 
+```shell
+pgai config --dbname=gitlabhq_dblab --prefix=<gitlab handle> --proxy=<domain> --exe=<optional /path/to/dblab>
+```
+
+To configure environments:
+
+```shell
+pgai env add --alias ci --id ci-database --port 12345
 ```
-Host <subdomain>.postgres.ai
+
+The environment id, port, and proxy domain can be found by clicking on the `Connect` button on a database instance page.
+
+Configuring the user and identity file for the proxy domain must be done using the `~/.ssh/config` file. Example:
+
+```
+Host <domain>
   IdentityFile ~/.ssh/id_ed25519
   User <username>
 ```
@@ -48,6 +62,13 @@ pgai connect <env alias>
 
 Multiple `connect` commands for an environment will connect to the same clone, it won't start a new one.
 
+### Features
+
+- multiple psql sessions to the same clone
+- multiple environments support
+- automatic port forward management
+- prevents system sleep while psql sessions are active via `caffeinate`
+
 ## Contributing
 
 Bug reports and pull requests are welcome on GitLab at https://gitlab.com/mbobin/pgai.

data/lib/pgai/clone_manager.rb

@@ -1,22 +1,14 @@
 # frozen_string_literal: true
 
-require "shellwords"
-require "socket"
-require "json"
-require "securerandom"
-require "net/http"
-require "pathname"
-require "fileutils"
-
 module Pgai
   class CloneManager
     HOSTNAME = "127.0.0.1"
-    DEFAULT_DBLAB = Pathname.new("~/.dblab/dblab").expand_path
-    DBLAB_BINARY_URL = "https://gitlab.com/api/v4/projects/45068363/jobs/4092515399/artifacts/engine/bin/cli/dblab-%{os}-%{cpu}"
 
     def initialize(environment, config:)
       @environment = environment
       @config = config
+      @port_forward = PortForward.new(config: config, hostname: HOSTNAME)
+      @dblab = Dblab.new(config: config, hostname: HOSTNAME)
     end
 
     def connect
@@ -29,20 +21,18 @@ module Pgai
       configure_enviroment
       return unless find_raw_clone
 
-      dblab("clone", "destroy", clone_id, raw: true)
+      dblab.destroy_clone(id: clone_id)
       config.remove_clone(clone_id)
-      stop_port_forwards
+      port_forward.stop
     end
 
     private
 
-    attr_reader :environment, :config
+    attr_reader :environment, :config, :port_forward, :dblab
 
     def configure_enviroment
-      start_port_forward(enviroment_port)
-
-      dblab("init", "--url", "http://#{HOSTNAME}:#{enviroment_port}", "--token", config.access_token, "--environment-id", environment_id)
-      dblab("config", "switch", environment_id, raw: true)
+      port_forward.start(enviroment_port)
+      dblab.configure_env(port: enviroment_port, token: config.access_token, id: environment_id)
     end
 
     def find_or_create_clone
@@ -60,12 +50,11 @@ module Pgai
     end
 
     def find_raw_clone
-      Array(dblab("clone", "list")).find { |clone| clone["id"] == clone_id }
+      dblab.list_clones.find { |clone| clone["id"] == clone_id }
     end
 
     def create_clone
-      raw_clone = dblab("clone", "create", "--id", clone_id, "--username", clone_user, "--password", clone_password)
-      raise "Could not create clone" unless raw_clone
+      raw_clone = dblab.create_clone(id: clone_id, user: clone_user, password: clone_password)
 
       attributes = {
         port: raw_clone.dig("db", "port"),
@@ -79,14 +68,16 @@ module Pgai
     end
 
     def psql(clone)
-      start_port_forward(clone.port)
+      port_forward.start(clone.port)
 
       psql_pid = fork do
-        wait_for_connections(clone.port)
+        start_caffeinate(Process.pid)
         exec("psql #{clone.connection_string}")
       end
-
       Process.wait(psql_pid)
+    ensure
+      port_forward.conditionally_stop(clone.port, [psql_pid])
+      port_forward.conditionally_stop(enviroment_port)
     end
 
     def environment_id
@@ -109,63 +100,11 @@ module Pgai
       @clone_user ||= SecureRandom.hex(16)
     end
 
-    def dblab(*args, raw: false)
-      output = `#{args.unshift(dblab_path).shelljoin}`
-      return output if raw
-
-      JSON.parse(output) unless output.empty?
-    end
-
-    def dblab_path
-      return config.path unless config.path.to_s.empty?
-      return "dblab" unless `which dblab`.to_s.empty?
-      return DEFAULT_DBLAB.to_s if DEFAULT_DBLAB.exist?
-
-      download_dblab_to(DEFAULT_DBLAB)
-      File.chmod(0o755, DEFAULT_DBLAB)
-
-      DEFAULT_DBLAB
-    end
-
-    def download_dblab_to(location)
-      platform = Gem::Platform.local
-      uri = URI(DBLAB_BINARY_URL % {os: platform.os, cpu: platform.cpu})
-
-      Net::HTTP.start(uri.host, uri.port, use_ssl: true) do |http|
-        http.request Net::HTTP::Get.new(uri) do |response|
-          FileUtils.mkdir_p File.dirname(location)
-          File.open(location, "w") do |io|
-            response.read_body { |chunk| io.write chunk }
-          end
-        end
-      end
-    end
-
-    def start_port_forward(port)
-      return if port_open?(port)
-
-      system("ssh -fNTML #{port}:#{HOSTNAME}:#{port} #{config.proxy}")
-      wait_for_connections(port)
-    end
-
-    def stop_port_forwards
-      raw_pids = `ps ax | grep 'ssh -fNTML' | grep '#{config.proxy}' | grep -v grep | awk '{ print $1 }'`
+    def start_caffeinate(pid)
+      return if `which caffeinate`.to_s.empty?
 
-      raw_pids.split.map do |pid|
-        Process.kill("HUP", pid.to_i)
-      end
-    end
-
-    def port_open?(port)
-      !!TCPSocket.new(HOSTNAME, port)
-    rescue Errno::ECONNREFUSED
-      false
-    end
-
-    def wait_for_connections(port)
-      until port_open?(port)
-        sleep 0.02
-      end
+      caffeinate_pid = Process.spawn("caffeinate -is -w #{pid}")
+      Process.detach(caffeinate_pid)
     end
   end
 end

data/lib/pgai/dblab.rb

@@ -0,0 +1,78 @@
+# frozen_string_literal: true
+
+require "shellwords"
+require "json"
+require "net/http"
+require "pathname"
+require "fileutils"
+
+module Pgai
+  class Dblab
+    DEFAULT_DBLAB = Pathname.new("~/.dblab/dblab").expand_path
+    DBLAB_RELEASE_CHANNEL = "master"
+    DBLAB_BINARY_URL = "https://storage.googleapis.com/database-lab-cli/#{DBLAB_RELEASE_CHANNEL}/dblab-%{os}-%{cpu}"
+
+    def initialize(config:, hostname:)
+      @config = config
+      @hostname = hostname
+    end
+
+    def configure_env(port:, token:, id:)
+      dblab("init", "--url", "http://#{hostname}:#{port}", "--token", token, "--environment-id", id, silence: true)
+      dblab("config", "switch", id, raw: true)
+    end
+
+    def list_clones
+      Array(dblab("clone", "list"))
+    end
+
+    def create_clone(id:, user:, password:)
+      data = dblab("clone", "create", "--id", id, "--username", user, "--password", password)
+      raise "Could not create clone" unless data
+
+      data
+    end
+
+    def destroy_clone(id:)
+      dblab("clone", "destroy", id, raw: true)
+    end
+
+    private
+
+    attr_reader :config, :hostname
+
+    def dblab(*args, raw: false, silence: false)
+      redirect = "2>/dev/null" if silence
+
+      output = `#{args.unshift(dblab_path).shelljoin} #{redirect}`
+      return output if raw
+
+      JSON.parse(output) unless output.empty?
+    end
+
+    def dblab_path
+      return config.path unless config.path.to_s.empty?
+      return "dblab" unless `which dblab`.to_s.empty?
+      return DEFAULT_DBLAB.to_s if DEFAULT_DBLAB.exist?
+
+      download_dblab_to(DEFAULT_DBLAB)
+      File.chmod(0o755, DEFAULT_DBLAB)
+
+      DEFAULT_DBLAB
+    end
+
+    def download_dblab_to(location)
+      platform = Gem::Platform.local
+      uri = URI(DBLAB_BINARY_URL % {os: platform.os, cpu: platform.cpu})
+
+      Net::HTTP.start(uri.host, uri.port, use_ssl: true) do |http|
+        http.request Net::HTTP::Get.new(uri) do |response|
+          FileUtils.mkdir_p File.dirname(location)
+          File.open(location, "w") do |io|
+            response.read_body { |chunk| io.write chunk }
+          end
+        end
+      end
+    end
+  end
+end

data/lib/pgai/port_forward.rb

@@ -0,0 +1,66 @@
+# frozen_string_literal: true
+
+require "shellwords"
+require "socket"
+
+module Pgai
+  class PortForward
+    def initialize(config:, hostname:)
+      @config = config
+      @hostname = hostname
+    end
+
+    def start(port)
+      return if ready?(port)
+
+      system("ssh -fNTML #{port.to_i}:#{escape hostname}:#{port.to_i} #{escape config.proxy}")
+      wait_until_ready(port)
+    end
+
+    def stop(port = nil)
+      port_forward_pids(port).each do |pid|
+        Process.kill("HUP", pid.to_i)
+      end
+    end
+
+    def conditionally_stop(port, ignored_pids = [])
+      pids = lsof(port, ignored_pids + [Process.pid])
+      pf_pids = port_forward_pids(port)
+
+      stop(port) if (pids - pf_pids).empty?
+    end
+
+    private
+
+    attr_reader :config, :hostname
+
+    def ready?(port)
+      !!TCPSocket.new(hostname, port)
+    rescue Errno::ECONNREFUSED
+      false
+    end
+
+    def wait_until_ready(port)
+      until ready?(port)
+        sleep 0.02
+      end
+    end
+
+    def port_forward_pids(port_filter = nil)
+      ssh_filter = escape "ssh -fNTML #{port_filter&.to_i}".strip
+      host_filter = escape config.proxy
+
+      `ps ax | grep #{ssh_filter} | grep #{host_filter} | grep -v grep | awk '{ print $1 }'`.split.map(&:to_i)
+    end
+
+    def lsof(port, exclude_pids = [])
+      ignored_pids = exclude_pids.map { |pid| "-p^#{pid.to_i}" }.join(" ")
+
+      `lsof -i:#{port.to_i} #{ignored_pids} -t`.split.map(&:to_i)
+    end
+
+    def escape(...)
+      Shellwords.escape(...)
+    end
+  end
+end

data/lib/pgai/version.rb

@@ -1,3 +1,3 @@
 module Pgai
-  VERSION = "0.1.3"
+  VERSION = "0.1.5"
 end

data/lib/pgai.rb

@@ -2,6 +2,8 @@
 
 require_relative "pgai/version"
 require_relative "pgai/config"
+require_relative "pgai/dblab"
+require_relative "pgai/port_forward"
 require_relative "pgai/clone_manager"
 require_relative "pgai/cli/base"
 require_relative "pgai/cli/env"

metadata

@@ -1,14 +1,14 @@
 --- !ruby/object:Gem::Specification
 name: pgai
 version: !ruby/object:Gem::Version
-  version: 0.1.3
+  version: 0.1.5
 platform: ruby
 authors:
 - Marius Bobin
 autorequire:
 bindir: bin
 cert_chain: []
-date: 2023-04-11 00:00:00.000000000 Z
+date: 2023-04-25 00:00:00.000000000 Z
 dependencies:
 - !ruby/object:Gem::Dependency
   name: thor
@@ -89,6 +89,8 @@ files:
 - lib/pgai/cli/main.rb
 - lib/pgai/clone_manager.rb
 - lib/pgai/config.rb
+- lib/pgai/dblab.rb
+- lib/pgai/port_forward.rb
 - lib/pgai/version.rb
 homepage: https://gitlab.com/mbobin/pgai
 licenses: