pgai 0.1.4 → 0.1.6

checksums.yaml

@@ -1,7 +1,7 @@
 ---
 SHA256:
-  metadata.gz: 29e67c364882e415161c298f0463f6e3850a7ea057b8256a3a4e17643d4a8694
-  data.tar.gz: bdcce4986037ace414698eeae62ce3f73f33084ace78e99e778a2b55975f5ad3
+  metadata.gz: 6ec24e1bc611522aea3c490acbb0f94b02085584544c979cd56ae9a3bd45f61b
+  data.tar.gz: 7351d32aa09341ec0081e0d4c45466d83034d1ec3ec776b74bef7e05c455997e
 SHA512:
-  metadata.gz: 90a083cc81d4c2b55892f9d1137aefa7dbb7bda601b2a058a42e6011a630eac7e4469ec586363a334a95c795955c7809d4b285937c49ba099e4761f16b7fb6a9
-  data.tar.gz: 0c1dcc4947700235508cd9ea3ffbcececa52f32b98c1d04c3f931ae7d97b72ef334c1409b2fd70c86a0867990445dbd4288b8cb9028ce9a910030d40bc51733c
+  metadata.gz: 82761cb5d303b13edc9a8b9541e11a9c5d2c06b3765879ede4d6cf4176a8d11f1bcd24563da88481ff6648b308f265f8cc0a2b56eb01ac2cd5888417ab485062
+  data.tar.gz: 527e1976b374a68409c6ca8591b6ab6cb043ed374f69a3daa87d76116bed7d0d04d518d00d6df38bf1305965f9db71cdf4e8aac9c2e4e01fa9a31aa7a00723be

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,23 +1,16 @@
 # frozen_string_literal: true
 
-require "shellwords"
-require "socket"
-require "json"
-require "securerandom"
-require "net/http"
-require "pathname"
-require "fileutils"
+require 'securerandom'
 
 module Pgai
   class CloneManager
     HOSTNAME = "127.0.0.1"
-    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(environment, config:)
       @environment = environment
       @config = config
+      @port_forward = PortForward.new(config: config, hostname: HOSTNAME)
+      @dblab = Dblab.new(config: config, hostname: HOSTNAME)
     end
 
     def connect
@@ -30,20 +23,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
@@ -61,12 +52,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"),
@@ -80,14 +70,18 @@ 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)
         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
@@ -110,63 +104,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 }'`
-
-      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 start_caffeinate(pid)
+      return if `which caffeinate`.to_s.empty?
 
-    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.4"
+  VERSION = "0.1.6"
 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.4
+  version: 0.1.6
 platform: ruby
 authors:
 - Marius Bobin
 autorequire:
 bindir: bin
 cert_chain: []
-date: 2023-04-12 00:00:00.000000000 Z
+date: 2023-04-27 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: