class Sequel::IBMDB::Database

def _execute(conn, sql, opts)
  stmt = log_yield(sql){conn.execute(sql)}
  if block_given?
    begin
      yield(stmt)
    ensure
      stmt.free
    end
  else  
    stmt.affected
  end
end

def alter_table(name, generator=nil, &block)
  res = super
  reorg(name)
  res
end

def begin_transaction(conn, opts={})
  log_yield(TRANSACTION_BEGIN){conn.autocommit = false}
end

def commit_transaction(conn, opts={})
  log_yield(TRANSACTION_COMMIT){conn.commit}
end

def connect(server)
  opts = server_opts(server)
  
  # use uncataloged connection so that host and port can be supported
  connection_string = ( \
      'Driver={IBM DB2 ODBC DRIVER};' \
      "Database=#{opts[:database]};" \
      "Hostname=#{opts[:host]};" \
      "Port=#{opts[:port] || 50000};" \
      'Protocol=TCPIP;' \
      "Uid=#{opts[:user]};" \
      "Pwd=#{opts[:password]};" \
  )
  Connection.new(connection_string)
end

def disconnect_connection(conn)
  conn.close
end

def execute(sql, opts={}, &block)
  if sql.is_a?(Symbol)
    execute_prepared_statement(sql, opts, &block)
  else
    synchronize(opts[:server]){|c| _execute(c, sql, opts, &block)}
  end
rescue Connection::Error => e
  raise_error(e)
end

def execute_insert(sql, opts={})
  synchronize(opts[:server]) do |c|
    if sql.is_a?(Symbol)
      execute_prepared_statement(sql, opts)
    else
      _execute(c, sql, opts)
    end
    _execute(c, "SELECT IDENTITY_VAL_LOCAL() FROM SYSIBM.SYSDUMMY1", opts){|stmt| i = stmt.fetch_array.first.to_i; stmt.free; i}
  end
rescue Connection::Error => e
  raise_error(e)
end

def execute_prepared_statement(ps_name, opts)
  args = opts[:arguments]
  ps = prepared_statements[ps_name]
  sql = ps.prepared_sql
  synchronize(opts[:server]) do |conn|
    unless conn.prepared_statements.fetch(ps_name, []).first == sql
      log_yield("Preparing #{ps_name}: #{sql}"){conn.prepare(sql, ps_name)}
    end
    args = args.map{|v| v.nil? ? nil : prepared_statement_arg(v)}
    stmt = log_yield("Executing #{ps_name}: #{args.inspect}"){conn.execute_prepared(ps_name, *args)}
    if block_given?
      begin
        yield(stmt)
      ensure
        stmt.free
      end
    else  
      stmt.affected
    end
  end
end

def initialize(opts={})
  super
  @conversion_procs = DB2_TYPES.dup
  @conversion_procs[:timestamp] = method(:to_application_timestamp)
end

def metadata_dataset
  ds = super
  ds.convert_smallint_to_bool = false
  ds
end

def prepared_statement_arg(v)
  case v
  when Numeric
    v.to_s
  when Date, Time
    literal(v).gsub("'", '')
  else
    v
  end
end

def remove_transaction(conn, committed)
  conn.autocommit = true
ensure
  super
end

def rollback_transaction(conn, opts={})
  log_yield(TRANSACTION_ROLLBACK){conn.rollback}
end

def schema_column_type(db_type)
  if Sequel::IBMDB.convert_smallint_to_bool && db_type =~ /smallint/i 
    :boolean
  else
    super
  end
end

def table_exists?(name)
  v ||= false # only retry once
  sch, table_name = schema_and_table(name)
  name = SQL::QualifiedIdentifier.new(sch, table_name) if sch
  from(name).first
  true
rescue DatabaseError => e
  if e.to_s =~ /Operation not allowed for reason code "7" on table/ && v == false
    # table probably needs reorg
    reorg(name)
    v = true
    retry 
  end
  false
end