#include "DbTest.h"
#include "PgTestEnvironment.h"
#include "catalog/PgDatabaseSizes.h"
#include "Pgsql_Result.h"

using namespace pgtest;

namespace {

Oid oidOfDatabase(Pgsql::Connection &c, const QString &db)
{
	Pgsql::Result r = c.query(QString("SELECT oid FROM pg_database WHERE datname = '%1'").arg(db));
	if (!r || r.rows() != 1)
		return InvalidOid;
	return QString::fromUtf8(r.val(0, 0)).toUInt();
}

const PgDatabaseSize *findSize(const PgDatabaseSizes &sizes, Oid oid)
{
	for (const auto &s : sizes) {
		if (s.oid == oid)
			return &s;
	}
	return nullptr;
}

} // namespace

class DatabaseSizes : public DbTest {};

TEST_F(DatabaseSizes, ConnectableDatabaseReturnsSize)
{
	auto plain = plainConn();
	Oid oid = oidOfDatabase(*plain, dbName());
	ASSERT_NE(oid, InvalidOid);

	PgDatabaseSizes sizes = loadDatabaseSizes(*plain);
	const PgDatabaseSize *s = findSize(sizes, oid);
	ASSERT_NE(s, nullptr);
	EXPECT_GE(s->totalBytes, 0);
}

TEST_F(DatabaseSizes, RevokedConnectReturnsMinusOne)
{
	Oid oid = oidOfDatabase(conn(), dbName());
	ASSERT_NE(oid, InvalidOid);
	conn().query(QString("REVOKE CONNECT ON DATABASE %1 FROM PUBLIC").arg(quoteIdent(dbName())));

	auto plain = std::make_unique<Pgsql::Connection>();
	plain->connect(urlForDbAsPlainRole("postgres"));

	PgDatabaseSizes sizes = loadDatabaseSizes(*plain);
	const PgDatabaseSize *s = findSize(sizes, oid);
	ASSERT_NE(s, nullptr);
	EXPECT_EQ(s->totalBytes, -1);
}

TEST_F(DatabaseSizes, SuperuserSeesAllDatabases)
{
	PgDatabaseSizes sizes = loadDatabaseSizes(conn());

	Pgsql::Result r = conn().query("SELECT count(*) FROM pg_database");
	ASSERT_EQ(r.rows(), 1);
	EXPECT_EQ(static_cast<size_t>(QString::fromUtf8(r.val(0, 0)).toULongLong()), sizes.size());

	for (const auto &s : sizes)
		EXPECT_GE(s.totalBytes, 0) << "database oid " << s.oid;
}