pgai 0.1.1

checksums.yaml

@@ -0,0 +1,7 @@
+---
+SHA256:
+  metadata.gz: c342a96f17431de20953406ba211771e0b7647f6ec7e59377e50399d89ca63bc
+  data.tar.gz: e1737c3f34f8b2b832c9706019eec6169fb723eb27ffc260acd285a8e24fce1e
+SHA512:
+  metadata.gz: e1981687bdf6204b16e651f6313c3db1eb8e9691817b23f49c4f81611163a9ac3df7588ffe2b58dbae5c4adcd75e924ea6e2c4835fa24c29bc73421d7a07c027
+  data.tar.gz: 7a9a3fe849f9d5cb1d010fba31900f07ff0f10413f222a329a080d3517b44aa6b870f09e7d9995bbff573cf434ba09f4364126df3e8831df5938a69932fbb4e2

data/MIT-LICENSE

@@ -0,0 +1,20 @@
+Copyright (c) 2023 Marius Bobin
+
+Permission is hereby granted, free of charge, to any person obtaining
+a copy of this software and associated documentation files (the
+"Software"), to deal in the Software without restriction, including
+without limitation the rights to use, copy, modify, merge, publish,
+distribute, sublicense, and/or sell copies of the Software, and to
+permit persons to whom the Software is furnished to do so, subject to
+the following conditions:
+
+The above copyright notice and this permission notice shall be
+included in all copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND,
+EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF
+MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND
+NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE
+LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION
+OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION
+WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.

data/README.md

@@ -0,0 +1,41 @@
+# Pgai
+
+Simple CLI wrapper around Postgres.ai CLI to drop directly into a PSQL session.
+
+## Installation
+
+Install the gem by executing:
+
+```shell
+gem install pgai
+```
+
+## Configuration
+
+For an overview of available commands, execute:
+
+```shell
+pgai -h
+```
+
+Before usage `pgai config` must be executed and at least an environment must be configured with `pgai env add`
+
+Postgres.ai does not provide `dblab` binaries for `ARM` processors(yet?), to work around this:
+
+- clone `https://gitlab.com/postgres-ai/database-lab`
+- change [`GOARCH = amd64`](https://gitlab.com/postgres-ai/database-lab/-/blob/0ad2cdd8f1d851e4beda92e0577bb242c6a8dd7f/engine/Makefile#L6) to `GOARCH = arm64`
+- execute [`make build`](https://gitlab.com/postgres-ai/database-lab/-/blob/0ad2cdd8f1d851e4beda92e0577bb242c6a8dd7f/engine/Makefile#L44-L47), we're interested only about building `CLI_BINARY`, the others could be removed.
+- the `dblab` binary should be available in the `bin` directory
+- run `pgai config --exe /path/to/bin/dblab` to set the path if it's not in `$PATH`.
+
+## Usage
+
+```shell
+pgai connect <env alias>
+```
+
+Multiple `connect` commands for an environment will connect to the same clone, it won't start a new one.
+
+## Contributing
+
+Bug reports and pull requests are welcome on GitLab at https://gitlab.com/mbobin/pgai.

data/bin/pgai

@@ -0,0 +1,8 @@
+#!/usr/bin/env ruby
+# frozen_string_literal: true
+
+Thread.report_on_exception = false
+
+require_relative "../lib/pgai"
+
+Pgai::Cli::Main.start(ARGV)

data/lib/pgai/cli/base.rb

@@ -0,0 +1,21 @@
+require "thor"
+
+module Pgai::Cli
+  class Base < Thor
+    def self.exit_on_failure?
+      true
+    end
+
+    private
+
+    def configuration
+      @configuration ||= load_and_validate_config!
+    end
+
+    def load_and_validate_config!
+      Pgai::Config.load.tap do |config|
+        raise "Access token is not set" unless config.access_token
+      end
+    end
+  end
+end

data/lib/pgai/cli/env.rb

@@ -0,0 +1,22 @@
+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
+    def add
+      configuration.add_env(**options)
+    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
+    def remove
+      configuration.remove_env(options[:alias])
+    end
+
+    desc "list", "List all configured environments"
+    def list
+      puts JSON.pretty_generate(configuration.enviroments)
+    end
+  end
+end

data/lib/pgai/cli/main.rb

@@ -0,0 +1,31 @@
+module Pgai::Cli
+  class Main < Base
+    desc "config", "Configure CLI options"
+    method_option :prefix, aliases: "-p", desc: "Specify prefix to be used for clones", required: true
+    method_option :dbname, aliases: "-n", desc: "Specify database name to connect to", required: true
+    method_option :proxy, aliases: "-x", desc: "Specify proxy host name", required: true
+    method_option :exe, aliases: "-e", desc: "dblab binary executable path", required: true
+    def config
+      token = ask("Access token:", echo: false)
+
+      Pgai::Config.persist(options.merge(token: token))
+    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
+    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
+    end
+
+    desc "env", "Manage environments"
+    subcommand "env", Pgai::Cli::Env
+  end
+end

data/lib/pgai/clone_manager.rb

@@ -0,0 +1,143 @@
+# frozen_string_literal: true
+
+require "shellwords"
+require "socket"
+require "json"
+require "securerandom"
+
+module Pgai
+  class CloneManager
+    HOSTNAME = "127.0.0.1"
+
+    def initialize(environment, config:)
+      @environment = environment
+      @config = config
+    end
+
+    def connect
+      configure_enviroment
+
+      psql find_or_create_clone
+    end
+
+    def cleanup
+      configure_enviroment
+      return unless find_raw_clone
+
+      dblab("clone", "destroy", clone_id, raw: true)
+      config.remove_clone(clone_id)
+      stop_port_forwards
+    end
+
+    private
+
+    attr_reader :environment, :config
+
+    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)
+    end
+
+    def find_or_create_clone
+      find_clone || create_clone
+    end
+
+    def find_clone
+      raw_clone = find_raw_clone
+
+      if raw_clone
+        config.find_clone(clone_id) || raise(Pagi::CloneNotFount, "clone not tracked")
+      else
+        config.remove_clone(clone_id) && nil
+      end
+    end
+
+    def find_raw_clone
+      Array(dblab("clone", "list")).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
+
+      attributes = {
+        port: raw_clone.dig("db", "port"),
+        password: clone_password,
+        user: clone_user,
+        host: HOSTNAME,
+        dbname: config.dbname
+      }
+
+      config.persist_clone(clone_id, attributes)
+    end
+
+    def psql(clone)
+      start_port_forward(clone.port)
+
+      psql_pid = fork do
+        wait_for_connections(clone.port)
+        exec("psql #{clone.connection_string}")
+      end
+
+      Process.wait(psql_pid)
+    end
+
+    def environment_id
+      environment.fetch(:id)
+    end
+
+    def enviroment_port
+      environment.fetch(:port)
+    end
+
+    def clone_id
+      @clone_id ||= [config.clone_prefix, environment.fetch(:alias)].join("_")
+    end
+
+    def clone_user
+      @clone_user ||= SecureRandom.hex(8)
+    end
+
+    def clone_password
+      @clone_user ||= SecureRandom.hex(16)
+    end
+
+    def dblab(*args, raw: false)
+      executable = config.path || "dblab"
+
+      output = `#{args.unshift(executable).shelljoin}`
+      return output if raw
+
+      JSON.parse(output) unless output.empty?
+    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 wait_for_connections(port)
+      until port_open?(port)
+        sleep 0.02
+      end
+    end
+  end
+end

data/lib/pgai/config.rb

@@ -0,0 +1,109 @@
+# frozen_string_literal: true
+
+require "pathname"
+require "pstore"
+
+module Pgai
+  class Config
+    Clone = Struct.new(:host, :port, :password, :user, :dbname, keyword_init: true) do
+      def connection_string
+        "'host=#{host} port=#{port} user=#{user} dbname=#{dbname} password=#{password}'"
+      end
+    end
+
+    attr_accessor :clone_prefix
+    attr_accessor :dbname
+    attr_accessor :access_token
+    attr_accessor :proxy
+    attr_accessor :path
+
+    def self.load
+      new { |config| config.load_from_store }
+    end
+
+    def self.persist(data)
+      new { |config| config.write(data) }
+    end
+
+    def initialize
+      yield self if block_given?
+    end
+
+    def load_from_store
+      store.transaction(true) do
+        self.clone_prefix = store[:clone_prefix]
+        self.dbname = store[:dbname]
+        self.access_token = store[:access_token]
+        self.proxy = store[:proxy]
+        self.path = store[:path]
+      end
+    end
+
+    def write(options)
+      store.transaction do
+        store[:clone_prefix] = options[:prefix]
+        store[:dbname] = options[:dbname]
+        store[:access_token] = options[:token]
+        store[:proxy] = options[:proxy]
+        store[:path] = options[:path]
+        store[:clones] = {}
+      end
+    end
+
+    def find_clone(id)
+      store.transaction(true) do
+        clones = store[:clones] || {}
+
+        Clone.new(clones[id]) if clones[id]
+      end
+    end
+
+    def remove_clone(id)
+      store.transaction do
+        store[:clones] ||= {}
+        store[:clones] = store[:clones].except(id)
+      end
+    end
+
+    def persist_clone(id, attributes = {})
+      store.transaction do
+        store[:clones] = {}
+        store[:clones].merge!(id => attributes)
+      end
+
+      Clone.new(attributes)
+    end
+
+    def add_env(options = {})
+      store.transaction do
+        store[:environments] ||= {}
+        store[:environments].merge!(options[:alias] => options)
+      end
+    end
+
+    def remove_env(env_alias)
+      store.transaction do
+        store[:environments] ||= {}
+        store[:environments] = store[:environments].except(env_alias)
+      end
+    end
+
+    def enviroments
+      store.transaction(true) do
+        store[:environments] || {}
+      end
+    end
+
+    private
+
+    def store
+      @store ||= PStore.new(store_path)
+    end
+
+    def store_path
+      @store_path ||= Pathname("~/.config/pgai/config.pstore").expand_path.tap do |path|
+        FileUtils.mkdir_p File.dirname(path)
+      end
+    end
+  end
+end

data/lib/pgai/version.rb

@@ -0,0 +1,3 @@
+module Pgai
+  VERSION = "0.1.1"
+end

data/lib/pgai.rb

@@ -0,0 +1,14 @@
+# frozen_string_literal: true
+
+require_relative "pgai/version"
+require_relative "pgai/config"
+require_relative "pgai/clone_manager"
+require_relative "pgai/cli/base"
+require_relative "pgai/cli/env"
+require_relative "pgai/cli/main"
+
+module Pgai
+  class Error < StandardError; end
+
+  class CloneNotFount < Error; end
+end

metadata

@@ -0,0 +1,116 @@
+--- !ruby/object:Gem::Specification
+name: pgai
+version: !ruby/object:Gem::Version
+  version: 0.1.1
+platform: ruby
+authors:
+- Marius Bobin
+autorequire:
+bindir: bin
+cert_chain: []
+date: 2023-04-10 00:00:00.000000000 Z
+dependencies:
+- !ruby/object:Gem::Dependency
+  name: thor
+  requirement: !ruby/object:Gem::Requirement
+    requirements:
+    - - "~>"
+      - !ruby/object:Gem::Version
+        version: '1.2'
+    - - ">="
+      - !ruby/object:Gem::Version
+        version: 1.2.1
+  type: :runtime
+  prerelease: false
+  version_requirements: !ruby/object:Gem::Requirement
+    requirements:
+    - - "~>"
+      - !ruby/object:Gem::Version
+        version: '1.2'
+    - - ">="
+      - !ruby/object:Gem::Version
+        version: 1.2.1
+- !ruby/object:Gem::Dependency
+  name: rake
+  requirement: !ruby/object:Gem::Requirement
+    requirements:
+    - - "~>"
+      - !ruby/object:Gem::Version
+        version: '13.0'
+  type: :development
+  prerelease: false
+  version_requirements: !ruby/object:Gem::Requirement
+    requirements:
+    - - "~>"
+      - !ruby/object:Gem::Version
+        version: '13.0'
+- !ruby/object:Gem::Dependency
+  name: rspec
+  requirement: !ruby/object:Gem::Requirement
+    requirements:
+    - - "~>"
+      - !ruby/object:Gem::Version
+        version: '3.0'
+  type: :development
+  prerelease: false
+  version_requirements: !ruby/object:Gem::Requirement
+    requirements:
+    - - "~>"
+      - !ruby/object:Gem::Version
+        version: '3.0'
+- !ruby/object:Gem::Dependency
+  name: standard
+  requirement: !ruby/object:Gem::Requirement
+    requirements:
+    - - "~>"
+      - !ruby/object:Gem::Version
+        version: '1.3'
+  type: :development
+  prerelease: false
+  version_requirements: !ruby/object:Gem::Requirement
+    requirements:
+    - - "~>"
+      - !ruby/object:Gem::Version
+        version: '1.3'
+description:
+email:
+- mbobin@gitlab.com
+executables:
+- pgai
+extensions: []
+extra_rdoc_files: []
+files:
+- MIT-LICENSE
+- README.md
+- bin/pgai
+- lib/pgai.rb
+- lib/pgai/cli/base.rb
+- lib/pgai/cli/env.rb
+- lib/pgai/cli/main.rb
+- lib/pgai/clone_manager.rb
+- lib/pgai/config.rb
+- lib/pgai/version.rb
+homepage: https://gitlab.com/mbobin/pgai
+licenses:
+- MIT
+metadata: {}
+post_install_message:
+rdoc_options: []
+require_paths:
+- lib
+required_ruby_version: !ruby/object:Gem::Requirement
+  requirements:
+  - - ">="
+    - !ruby/object:Gem::Version
+      version: 3.0.0
+required_rubygems_version: !ruby/object:Gem::Requirement
+  requirements:
+  - - ">="
+    - !ruby/object:Gem::Version
+      version: '0'
+requirements: []
+rubygems_version: 3.4.1
+signing_key:
+specification_version: 4
+summary: CLI wrapper for postgres.ai thin clones
+test_files: []