#!/usr/bin/env ruby

require_relative '../builder/builder.rb'

module Prick::SubCommand
  # Drop database owner
  def self.drop_owner(username)
    # No database at this point so no session is running and no cascade. A nop
    # if the user is the current user or owns objects in other databases
    if username != ENV['USER']
      super_conn.role.drop(username, fail: false, silent: true)
    end
  end

  # Drop all users associated with the given database except the owner. We
  # assume that users can only be connected to the project database and that
  # they don't own objects in other databases
  def self.drop_users(database)
    constrain database, String
    conn = super_conn
    if conn.rdbms.exist? database
      users = conn.role.list(database: database).reject { _1 == conn.rdbms.owner(database) }
      conn.session.terminate(database, users)
      if conn.rdbms.exist?(database)
        # Connect to database to make cascade work. Run in a block to close connection afterwards
        PgConn.new(database) { |conn| conn.role.drop(users, cascade: true) }
      else
        conn.rdbms.role.drop(users)
      end
    else
      # We don't terminate sessions because we assume one-database-users
      # FIXME: users is undefined
      raise NotImpementedError
      conn.role.drop(users, cascade: true) # Fails if the users owns objects in other databases
    end
  end

  # Drop a database
  def self.drop_database(database)
    State.connection.session.terminate(database, nil)
    State.connection { |conn| conn.rdbms.drop database }
  end

  # Drop data not in snapshot
  def self.drop_data(database, schemas = nil)
    conn = Prick.state.conn
    conn.session.triggers(false) {
      if schemas.nil?
        builder = Prick::Build::Builder.new(conn, Prick.state.schema_dir)
        pool = builder.pool
        schemas = pool.refresh_schemas
      end
      schema_list = conn.quote_values(schemas)
      schema_expr = schemas ? "true = true" : "schema_name in (#{schema_list})"
      tables = conn.tuples %(
        select schema_name, table_name, max_id
        from prick.snapshots
        where #{schema_expr}
      )
      tables.each { |schema, table, max_id|
        uid = "#{schema}.#{table}"
        if max_id
          conn.exec "delete from #{uid} where id > #{max_id}"
          conn.schema.set_serial(schema, table, max_id) if conn.schema.sequence(schema, table)
        else
          conn.exec "delete from #{uid}"
          conn.schema.set_serial(schema, table, nil) if conn.schema.sequence(schema, table)
        end
      }
    }
  end

  # Empty the database
  def self.drop_schema(database, schemas = [])
    constrain database, String
    constrain schemas, [String]
    if schemas.empty?
      State.connection { |conn| conn.rdbms.empty!(database, exclude: "prick") }
    else
      Prick.state.connection { |conn|
        schemas.each { |schema| conn.schema.drop(schema, cascade: true) }
      }
    end
  end
end