lib/aptible/cli/helpers/database.rb in aptible-cli-0.6.8 vs lib/aptible/cli/helpers/database.rb in aptible-cli-0.6.9

@@ -42,81 +42,62 @@
           say "=== #{environment.handle}"
           environment.databases.each { |db| say db.handle }
           say ''
         end
 
-        def establish_connection(database, local_port)
-          ENV['ACCESS_TOKEN'] = fetch_token
-          ENV['APTIBLE_DATABASE'] = database.handle
-
-          remote_port = claim_remote_port(database)
-          ENV['TUNNEL_PORT'] = remote_port
-
-          tunnel_args = "-L #{local_port}:localhost:#{remote_port}"
-          command = "ssh #{tunnel_args} #{common_ssh_args(database)}"
-          Kernel.exec(command)
-        end
-
         def clone_database(source, dest_handle)
           op = source.create_operation(type: 'clone', handle: dest_handle)
           poll_for_success(op)
 
           databases_from_handle(dest_handle, source.account).first
         end
 
-        def dump_database(database)
-          execute_local_tunnel(database) do |url|
-            filename = "#{database.handle}.dump"
-            say "Dumping to #{filename}"
-            `pg_dump #{url} > #{filename}`
+        # Creates a local tunnel and yields the helper
+
+        def with_local_tunnel(database, port = 0)
+          env = {
+            'ACCESS_TOKEN' => fetch_token,
+            'APTIBLE_DATABASE' => database.href
+          }
+          command = ['ssh', '-q'] + ssh_args(database)
+          Helpers::Tunnel.new(env, command).tap do |tunnel_helper|
+            tunnel_helper.start(port)
+            yield tunnel_helper
+            tunnel_helper.stop
           end
         end
 
-        # Creates a local tunnel and yields the url to it
-        def execute_local_tunnel(database)
-          local_port = random_local_port
-          pid = fork { establish_connection(database, local_port) }
+        # Creates a local PG tunnel and yields the url to it
 
-          # TODO: Better test for connection readiness
-          sleep 10
+        def with_postgres_tunnel(database)
+          if database.type != 'postgresql'
+            fail Thor::Error, 'This command only works for PostgreSQL'
+          end
 
-          auth = "aptible:#{database.passphrase}"
-          host = "localhost:#{local_port}"
-          yield "postgresql://#{auth}@#{host}/db"
-        ensure
-          Process.kill('HUP', pid) if pid
+          with_local_tunnel(database) do |tunnel_helper|
+            auth = "aptible:#{database.passphrase}"
+            host = "localhost:#{tunnel_helper.port}"
+            yield "postgresql://#{auth}@#{host}/db"
+          end
         end
 
-        def random_local_port
-          # Allocate a dummy server to discover an available port
-          dummy = TCPServer.new('127.0.0.1', 0)
-          port = dummy.addr[1]
-          dummy.close
-          port
-        end
-
         def local_url(database, local_port)
           remote_url = database.connection_url
           uri = URI.parse(remote_url)
 
           "#{uri.scheme}://#{uri.user}:#{uri.password}@" \
           "127.0.0.1:#{local_port}#{uri.path}"
         end
 
-        def claim_remote_port(database)
-          ENV['ACCESS_TOKEN'] = fetch_token
-
-          `ssh #{common_ssh_args(database)} 2>/dev/null`.chomp
-        end
-
-        def common_ssh_args(database)
+        def ssh_args(database)
           host = database.account.bastion_host
           port = database.account.bastion_port
 
-          opts = " -o 'SendEnv=*' -o StrictHostKeyChecking=no " \
-                 '-o UserKnownHostsFile=/dev/null'
-          connection_args = "-p #{port} root@#{host}"
-          "#{opts} #{connection_args}"
+          ['-o', 'SendEnv=APTIBLE_DATABASE',
+           '-o', 'SendEnv=ACCESS_TOKEN',
+           '-o', 'StrictHostKeyChecking=no',
+           '-o', 'UserKnownHostsFile=/dev/null',
+           '-p', port.to_s, "root@#{host}"]
         end
       end
     end
   end
 end