bool
CancelQuery(PGconn *conn, int timeout)
{
	char errbuf[ERRBUFF_SIZE];
	PGcancel *pgcancel;

	if (wait_connection_availability(conn, timeout) != 1)
		return false;

	pgcancel = PQgetCancel(conn);

	if (pgcancel != NULL)
	{
		/*
		 * PQcancel can only return 0 if socket()/connect()/send()
 		 * fails, in any of those cases we can assume something
		 * bad happened to the connection
		 */
		if (PQcancel(pgcancel, errbuf, ERRBUFF_SIZE) == 0)
		{
			log_warning(_("Can't stop current query: %s\n"), errbuf);
			PQfreeCancel(pgcancel);
			return false;
		}

		PQfreeCancel(pgcancel);
	}

	return true;
}

GSQLCursorState 
pgsql_cursor_stop (GSQLCursor *cursor)
{
	GSQL_TRACE_FUNC;
	
	GSQLSession *session;
	GSQLEPGSQLSession *spec_session;
	GSQLEPGSQLCursor *spec_cursor;
	PGcancel *cancel = NULL;
	gchar buff[256];
	
	g_return_val_if_fail (GSQL_IS_CURSOR (cursor), GSQL_CURSOR_STATE_ERROR);

	session = cursor->session;
	g_return_val_if_fail (GSQL_IS_SESSION (session), 
	                      GSQL_CURSOR_STATE_ERROR);

	spec_session = session->spec;
	spec_cursor = cursor->spec;

	cancel = PQgetCancel (spec_session->pgconn);
	if ( ! PQcancel (cancel, buff, 256) ) {
		pgsql_session_workspace_info (session, buff);
		PQfreeCancel (cancel);
		return GSQL_CURSOR_STATE_ERROR;
	}
	PQfreeCancel (cancel);
	return GSQL_CURSOR_STATE_STOP;
}

/*
 * on_before_exec
 *
 * Set cancel_conn to point to the current database connection.
 */
static void
on_before_exec(PGconn *conn)
{
	PGcancel   *old;

	if (in_cleanup)
		return;	/* forbid cancel during cleanup */

#ifdef WIN32
	EnterCriticalSection(&cancelConnLock);
#endif

	/* Free the old one if we have one */
	old = cancel_conn;

	/* be sure handle_sigint doesn't use pointer while freeing */
	cancel_conn = NULL;

	if (old != NULL)
		PQfreeCancel(old);

	cancel_conn = PQgetCancel(conn);

#ifdef WIN32
	LeaveCriticalSection(&cancelConnLock);
#endif
}

/* MultiClientCancel cancels the running query on the given connection. */
bool
MultiClientCancel(int32 connectionId)
{
	PGconn *connection = NULL;
	PGcancel *cancelObject = NULL;
	int cancelSent = 0;
	bool canceled = true;
	char errorBuffer[STRING_BUFFER_SIZE];

	Assert(connectionId != INVALID_CONNECTION_ID);
	connection = ClientConnectionArray[connectionId];
	Assert(connection != NULL);

	cancelObject = PQgetCancel(connection);

	cancelSent = PQcancel(cancelObject, errorBuffer, sizeof(errorBuffer));
	if (cancelSent == 0)
	{
		ereport(WARNING, (errmsg("could not issue cancel request"),
						  errdetail("Client error: %s", errorBuffer)));

		canceled = false;
	}

	PQfreeCancel(cancelObject);

	return canceled;
}

static void
connection_dealloc(PyObject* obj)
{
    connectionObject *self = (connectionObject *)obj;

    /* Make sure to untrack the connection before calling conn_close, which may
     * allow a different thread to try and dealloc the connection again,
     * resulting in a double-free segfault (ticket #166). */
    PyObject_GC_UnTrack(self);

    conn_close(self);

    if (self->weakreflist) {
        PyObject_ClearWeakRefs(obj);
    }

    conn_notice_clean(self);

    PyMem_Free(self->dsn);
    PyMem_Free(self->encoding);
    if (self->critical) free(self->critical);
    if (self->cancel) PQfreeCancel(self->cancel);

    connection_clear(self);

    pthread_mutex_destroy(&(self->lock));

    Dprintf("connection_dealloc: deleted connection object at %p, refcnt = "
        FORMAT_CODE_PY_SSIZE_T,
        obj, Py_REFCNT(obj)
      );

    Py_TYPE(obj)->tp_free(obj);
}

void pgConn::ResetConnCancel(void)
{
	wxMutexLocker  lock(m_cancelConnMutex);
	PGcancel      *oldCancelConn = m_cancelConn;

	m_cancelConn = NULL;

	if (oldCancelConn != NULL)
		PQfreeCancel(oldCancelConn);
}

void dbgPgConn::Cancel()
{
	// Attempt to cancel any ongoing query
	if (m_pgConn)
	{
		PGcancel *cancel = PQgetCancel(m_pgConn);
		char errbuf[256];
		PQcancel(cancel, errbuf, sizeof(errbuf));
		PQfreeCancel(cancel);
	}
}

void
CancelQuery(PGconn *conn)
{
	char errbuf[ERRBUFF_SIZE];
	PGcancel *pgcancel;

	pgcancel = PQgetCancel(conn);

	if (!pgcancel || PQcancel(pgcancel, errbuf, ERRBUFF_SIZE) == 0)
		log_warning(_("Can't stop current query: %s\n"), errbuf);

	PQfreeCancel(pgcancel);
}

static void
CancelQuery(void)
{
	char errbuf[256];
	PGcancel *pgcancel;

	pgcancel = PQgetCancel(primaryConn);

	if (!pgcancel || PQcancel(pgcancel, errbuf, 256) == 0)
		fprintf(stderr, "Can't stop current query: %s", errbuf);

	PQfreeCancel(pgcancel);
}

static void
CancelQuery(void)
{
	char errbuf[ERRBUFF_SIZE];
	PGcancel *pgcancel;

	pgcancel = PQgetCancel(primaryConn);

	if (!pgcancel || PQcancel(pgcancel, errbuf, ERRBUFF_SIZE) == 0)
		log_warning("Can't stop current query: %s\n", errbuf);

	PQfreeCancel(pgcancel);
}

/* set up the cancel key of the connection.
 * On success return 0, else set an exception and return -1
 */
RAISES_NEG static int
conn_setup_cancel(connectionObject *self, PGconn *pgconn)
{
    if (self->cancel) {
        PQfreeCancel(self->cancel);
    }

    if (!(self->cancel = PQgetCancel(self->pgconn))) {
        PyErr_SetString(OperationalError, "can't get cancellation key");
        return -1;
    }

    return 0;
}

static inline void free_conn(value v_conn)
{
  PGconn *conn = get_conn(v_conn);
  if (conn) {
    PGcancel *cancel = get_cancel_obj(v_conn);
    set_cancel_obj(v_conn, NULL);
    np_decr_refcount(get_conn_cb(v_conn));
    set_conn_cb(v_conn, NULL);
    set_conn(v_conn, NULL);
    caml_enter_blocking_section();
      PQfreeCancel(cancel);
      PQfinish(conn);
    caml_leave_blocking_section();
  }
}

static void
connection_dealloc(PyObject* obj)
{
    connectionObject *self = (connectionObject *)obj;

    /* Make sure to untrack the connection before calling conn_close, which may
     * allow a different thread to try and dealloc the connection again,
     * resulting in a double-free segfault (ticket #166). */
    PyObject_GC_UnTrack(self);

    /* close the connection only if this is the same process it was created
     * into, otherwise using multiprocessing we may close the connection
     * belonging to another process. */
#ifdef CONN_CHECK_PID
    if (self->procpid == getpid())
#endif
    {
        conn_close(self);
    }

    if (self->weakreflist) {
        PyObject_ClearWeakRefs(obj);
    }

    conn_notice_clean(self);

    PyMem_Free(self->dsn);
    PyMem_Free(self->encoding);
    if (self->error) free(self->error);
    if (self->cancel) PQfreeCancel(self->cancel);
    PQclear(self->pgres);

    connection_clear(self);

    pthread_mutex_destroy(&(self->lock));

    Dprintf("connection_dealloc: deleted connection object at %p, refcnt = "
        FORMAT_CODE_PY_SSIZE_T,
        obj, Py_REFCNT(obj)
      );

    Py_TYPE(obj)->tp_free(obj);
}

/*
 * ResetCancelConn
 *
 * Free the current cancel connection, if any, and set to NULL.
 */
void
ResetCancelConn(void)
{
	PGcancel   *oldCancelConn;

#ifdef WIN32
	EnterCriticalSection(&cancelConnLock);
#endif

	oldCancelConn = cancelConn;
	/* be sure handle_sigint doesn't use pointer while freeing */
	cancelConn = NULL;

	if (oldCancelConn != NULL)
		PQfreeCancel(oldCancelConn);

#ifdef WIN32
	LeaveCriticalSection(&cancelConnLock);
#endif
}

void toQPSqlConnectionSub::nativeCancel()
{
    QVariant v = Connection.driver()->handle();
    if (v.isValid() && v.typeName() == QString("PGconn*"))
    {
#ifdef LIBPQ_DECL_CANCEL
        PGconn *handle = *static_cast<PGconn **>(v.data());
        if (!handle)
            return;

        PGcancel *cancel = PQgetCancel(handle);
        if (!cancel)
            return;

        char *errbuf = new char[1024];
        PQcancel(cancel, errbuf, 1024);
        PQfreeCancel(cancel);
        delete[] errbuf;
#endif
    }
}

/*
 * Close the connection to the database and also cancel off the query if we
 * have one running.
 */
void
DisconnectDatabase(Archive *AHX)
{
	ArchiveHandle *AH = (ArchiveHandle *) AHX;
	PGcancel   *cancel;
	char		errbuf[1];

	if (!AH->connection)
		return;

	if (PQtransactionStatus(AH->connection) == PQTRANS_ACTIVE)
	{
		if ((cancel = PQgetCancel(AH->connection)))
		{
			PQcancel(cancel, errbuf, sizeof(errbuf));
			PQfreeCancel(cancel);
		}
	}

	PQfinish(AH->connection);
	AH->connection = NULL;
}

/*
 * DisconnectDatabase
 *		Disconnect the connection associated with the given slot
 */
static void
DisconnectDatabase(ParallelSlot *slot)
{
	char		errbuf[256];

	if (!slot->connection)
		return;

	if (PQtransactionStatus(slot->connection) == PQTRANS_ACTIVE)
	{
		PGcancel   *cancel;

		if ((cancel = PQgetCancel(slot->connection)))
		{
			PQcancel(cancel, errbuf, sizeof(errbuf));
			PQfreeCancel(cancel);
		}
	}

	PQfinish(slot->connection);
	slot->connection = NULL;
}

/*
 * Send cancel/finish signal to still-running QE through libpq.
 * waitMode is either CANCEL or FINISH.  Returns true if we successfully
 * sent a signal (not necessarily received by the target process).
 */
static DispatchWaitMode
cdbdisp_signalQE(SegmentDatabaseDescriptor * segdbDesc,
				 DispatchWaitMode waitMode)
{
	char errbuf[256];
	PGcancel *cn = PQgetCancel(segdbDesc->conn);
	int	ret = 0;

	if (cn == NULL)
		return DISPATCH_WAIT_NONE;

	/*
	 * PQcancel uses some strcpy/strcat functions; let's
	 * clear this for safety.
	 */
	MemSet(errbuf, 0, sizeof(errbuf));

	if (Debug_cancel_print || DEBUG4 >= log_min_messages)
		write_log("Calling PQcancel for %s", segdbDesc->whoami);

	/*
	 * Send query-finish, unless the client really wants to cancel the
	 * query.  This could happen if cancel comes after we sent finish.
	 */
	if (waitMode == DISPATCH_WAIT_CANCEL)
		ret = PQcancel(cn, errbuf, 256);
	else if (waitMode == DISPATCH_WAIT_FINISH)
		ret = PQrequestFinish(cn, errbuf, 256);
	else
		write_log("unknown waitMode: %d", waitMode);

	if (ret == 0 && (Debug_cancel_print || LOG >= log_min_messages))
		write_log("Unable to cancel: %s", errbuf);

	PQfreeCancel(cn);

	return (ret != 0 ? waitMode : DISPATCH_WAIT_NONE);
}

/*
 * SetCancelConn
 *
 * Set cancelConn to point to the current database connection.
 */
void
SetCancelConn(void)
{
	PGcancel   *oldCancelConn;

#ifdef WIN32
	EnterCriticalSection(&cancelConnLock);
#endif

	/* Free the old one if we have one */
	oldCancelConn = cancelConn;
	/* be sure handle_sigint doesn't use pointer while freeing */
	cancelConn = NULL;

	if (oldCancelConn != NULL)
		PQfreeCancel(oldCancelConn);

	cancelConn = PQgetCancel(pset.db);

#ifdef WIN32
	LeaveCriticalSection(&cancelConnLock);
#endif
}

/*
 * on_after_exec
 *
 * Free the current cancel connection, if any, and set to NULL.
 */
static void
on_after_exec(void)
{
	PGcancel   *old;

	if (in_cleanup)
		return;	/* forbid cancel during cleanup */

#ifdef WIN32
	EnterCriticalSection(&cancelConnLock);
#endif

	old = cancel_conn;

	/* be sure handle_sigint doesn't use pointer while freeing */
	cancel_conn = NULL;

	if (old != NULL)
		PQfreeCancel(old);

#ifdef WIN32
	LeaveCriticalSection(&cancelConnLock);
#endif
}