/*
 * src/bin/pgcopydb/pgsql.h
 *     Functions for interacting with a postgres server
 */

#ifndef PGSQL_H
#define PGSQL_H


#include <limits.h>
#include <stdbool.h>

#include "postgres.h"
#include "libpq-fe.h"
#include "portability/instr_time.h"

#include "access/xlogdefs.h"

#if PG_MAJORVERSION_NUM >= 15
#include "common/pg_prng.h"
#endif

#include "defaults.h"
#include "parsing_utils.h"
#include "pg_utils.h"


/*
 * OID values from PostgreSQL src/include/catalog/pg_type.h
 */
#define BOOLOID 16
#define NAMEOID 19
#define INT4OID 23
#define INT8OID 20
#define TEXTOID 25
#define LSNOID 3220

#define FLOAT4OID 700
#define FLOAT8OID 701

/*
 * Maximum connection info length as used in walreceiver.h
 */
#define MAXCONNINFO 1024

/*
 * Chunk size for reading and writting large objects
 */
#define LOBBUFSIZE 16 * 1024 * 1024 /* 16 MB */


/*
 * pg_stat_replication.sync_state is one if:
 *   sync, async, quorum, potential
 */
#define PGSR_SYNC_STATE_MAXLENGTH 10

/*
 * We receive a list of "other nodes" from the monitor, and we store that list
 * in local memory. We pre-allocate the memory storage, and limit how many node
 * addresses we can handle because of the pre-allocation strategy.
 */
#define NODE_ARRAY_MAX_COUNT 12


/* abstract representation of a Postgres server that we can connect to */
typedef enum
{
	PGSQL_CONN_SOURCE = 0,
	PGSQL_CONN_TARGET,
} ConnectionType;


/*
 * Retry policy to follow when we fail to connect to a Postgres URI.
 *
 * In almost all the code base the retry mechanism is implemented in the main
 * loop so we want to fail fast and let the main loop handle the connection
 * retry and the different network timeouts that we have, including the network
 * partition detection timeout.
 *
 * When we do retry connecting, we implement an Exponential Backoff with
 * Decorrelated Jitter algorithm as proven useful in the following article:
 *
 *  https://aws.amazon.com/blogs/architecture/exponential-backoff-and-jitter/
 */
typedef struct ConnectionRetryPolicy
{
	int maxT;                   /* maximum time spent retrying (seconds) */
	int maxR;                   /* maximum number of retries, might be zero */
	int maxSleepTime;           /* in millisecond, used to cap sleepTime */
	int baseSleepTime;          /* in millisecond, base time to sleep for */
	int sleepTime;              /* in millisecond, time waited for last round */

	instr_time startTime;       /* time of the first attempt */
	instr_time connectTime;     /* time of successful connection */
	int attempts;               /* how many attempts have been made so far */

#if PG_MAJORVERSION_NUM >= 15
	pg_prng_state prng_state;
#endif
} ConnectionRetryPolicy;

/*
 * Denote if the connetion is going to be used for one, or multiple statements.
 * This is used by psql_* functions to know if a connection is to be closed
 * after successful completion, or if the the connection is to be maintained
 * open for further queries.
 *
 * A common use case for maintaining a connection open, is while wishing to open
 * and maintain a transaction block. Another, is while listening for events.
 */
typedef enum
{
	PGSQL_CONNECTION_SINGLE_STATEMENT = 0,
	PGSQL_CONNECTION_MULTI_STATEMENT
} ConnectionStatementType;

/*
 * Allow higher level code to distinguish between failure to connect to the
 * target Postgres service and failure to run a query or obtain the expected
 * result. To that end we expose PQstatus() of the connection.
 *
 * We don't use the same enum values as in libpq because we want to have the
 * unknown value when we didn't try to connect yet.
 */
typedef enum
{
	PG_CONNECTION_UNKNOWN = 0,
	PG_CONNECTION_OK,
	PG_CONNECTION_BAD
} PGConnStatus;

/*
 * Support for ISOLATION LEVEL in pgsql_set_transaction transaction modes.
 *
 * ISOLATION LEVEL
 * { SERIALIZABLE | REPEATABLE READ | READ COMMITTED | READ UNCOMMITTED }
 */
typedef enum
{
	ISOLATION_SERIALIZABLE = 0,
	ISOLATION_REPEATABLE_READ,
	ISOLATION_READ_COMMITTED,
	ISOLATION_READ_UNCOMMITTED,
} IsolationLevel;

/*
 * As a way to communicate the SQL STATE when an error occurs, every
 * pgsql_execute_with_params context structure must have the same first field,
 * an array of 5 characters (plus '\0' at the end).
 */
#define SQLSTATE_LENGTH 6

/*
 * That's "x.yy.zz" or "xx.zz" or maybe a debian style version string such as:
 *  "13.8 (Debian 13.8-1.pgdg110+1)"
 *  "16beta1 (Debian 16~beta1-2.pgdg+~20230605.2256.g3f1aaaa)"
 */
#define PG_VERSION_STRING_MAX_LENGTH 128

/* notification processing */
typedef bool (*ProcessNotificationFunction)(int notificationGroupId,
											int64_t notificationNodeId,
											char *channel, char *payload);

typedef struct PGSQL
{
	ConnectionType connectionType;
	ConnectionStatementType connectionStatementType;

	char *connectionString;
	SafeURI safeURI;

	PGconn *connection;
	ConnectionRetryPolicy retryPolicy;
	PGConnStatus status;
	char sqlstate[SQLSTATE_LENGTH];

	char pgversion[PG_VERSION_STRING_MAX_LENGTH];
	int pgversion_num;

	ProcessNotificationFunction notificationProcessFunction;
	int notificationGroupId;
	int64_t notificationNodeId;
	bool notificationReceived;

	bool logSQL;
	bool singleRowMode;

	/*
	 * Keeps track of the last sync time for the pipeline. This is relevant
	 * only for connections in pipeline mode.
	 */
	uint64_t pipelineSyncTime;
} PGSQL;


/*
 * Arrange a generic way to parse PostgreSQL result from a query. Most of the
 * queries we need here return a single row of a single column, so that's what
 * the default context and parsing allows for.
 */

/* callback for parsing query results */
typedef void (ParsePostgresResultCB)(void *context, PGresult *result);

typedef enum
{
	PGSQL_RESULT_BOOL = 1,
	PGSQL_RESULT_INT,
	PGSQL_RESULT_BIGINT,
	PGSQL_RESULT_STRING
} QueryResultType;

#define STR_ERRCODE_CLASS_CONNECTION_EXCEPTION "08"

typedef struct AbstractResultContext
{
	char sqlstate[SQLSTATE_LENGTH];
} AbstractResultContext;

/* data structure for keeping a single-value query result */
typedef struct SingleValueResultContext
{
	char sqlstate[SQLSTATE_LENGTH];
	QueryResultType resultType;
	bool parsedOk;
	bool isNull;
	int ntuples;
	bool boolVal;
	int intVal;
	uint64_t bigint;
	char *strVal;
} SingleValueResultContext;


/* PostgreSQL ("Grand Unified Configuration") setting */
typedef struct GUC
{
	char *name;
	char *value;
} GUC;

bool pgsql_init(PGSQL *pgsql, char *url, ConnectionType connectionType);

PGconn * pgsql_open_connection(PGSQL *pgsql);

void pgsql_set_retry_policy(ConnectionRetryPolicy *retryPolicy,
							int maxT,
							int maxR,
							int maxSleepTime,
							int baseSleepTime);
void pgsql_set_interactive_retry_policy(ConnectionRetryPolicy *retryPolicy);
int pgsql_compute_connection_retry_sleep_time(ConnectionRetryPolicy *retryPolicy);
bool pgsql_retry_policy_expired(ConnectionRetryPolicy *retryPolicy);

bool pgsql_state_is_connection_error(PGSQL *pgsql);

void pgsql_finish(PGSQL *pgsql);
void parseSingleValueResult(void *ctx, PGresult *result);
void fetchedRows(void *ctx, PGresult *result);

bool pgsql_begin(PGSQL *pgsql);
bool pgsql_commit(PGSQL *pgsql);
bool pgsql_rollback(PGSQL *pgsql);

bool pgsql_server_version(PGSQL *pgsql);
bool pgsql_set_client_encoding_for_server(PGSQL *pgsql);

bool pgsql_set_transaction(PGSQL *pgsql,
						   IsolationLevel level,
						   bool readOnly,
						   bool deferrable);

bool pgsql_is_in_recovery(PGSQL *pgsql, bool *is_in_recovery);

bool pgsql_has_database_privilege(PGSQL *pgsql, const char *privilege,
								  bool *granted);

bool pgsql_has_sequence_privilege(PGSQL *pgsql,
								  const char *seqname,
								  const char *privilege,
								  bool *granted);

bool pgsql_has_table_privilege(PGSQL *pgsql,
							   const char *tablename,
							   const char *privilege,
							   bool *granted);

bool pgsql_get_table_relkind(PGSQL *pgsql, const char *qname, char *relkind);

bool pgsql_get_search_path(PGSQL *pgsql, char *search_path, size_t size);
bool pgsql_set_search_path(PGSQL *pgsql, char *search_path, bool local);
bool pgsql_prepend_search_path(PGSQL *pgsql, const char *nsp);

bool pgsql_export_snapshot(PGSQL *pgsql, char *snapshot, size_t size);
bool pgsql_set_snapshot(PGSQL *pgsql, char *snapshot);

bool pgsql_execute(PGSQL *pgsql, const char *sql);
bool pgsql_execute_with_params(PGSQL *pgsql, const char *sql, int paramCount,
							   const Oid *paramTypes, const char **paramValues,
							   void *parseContext, ParsePostgresResultCB *parseFun);

bool pgsql_send_with_params(PGSQL *pgsql, const char *sql, int paramCount,
							const Oid *paramTypes, const char **paramValues);

bool pgsql_fetch_results(PGSQL *pgsql, bool *done,
						 void *context, ParsePostgresResultCB *parseFun);

bool pgsql_enable_pipeline_mode(PGSQL *pgsql);
bool pgsql_sync_pipeline(PGSQL *pgsql);

bool pgsql_prepare(PGSQL *pgsql, const char *name, const char *sql,
				   int paramCount, const Oid *paramTypes);

bool pgsql_execute_prepared(PGSQL *pgsql, const char *name,
							int paramCount, const char **paramValues,
							void *context, ParsePostgresResultCB *parseFun);

bool validate_connection_string(const char *connectionString);

bool pgsql_lock_table(PGSQL *pgsql, const char *qname, const char *lockmode);

bool pgsql_truncate(PGSQL *pgsql, const char *qname, char relkind,
					const char *datname);

typedef struct CopyArgs
{
	char *srcQname;
	char *srcAttrList;
	char *srcWhereClause;
	char *dstQname;
	char *dstAttrList;
	char *logCommand;
	char *datname;
	bool truncate;
	bool freeze;
	bool useCopyBinary;
} CopyArgs;


typedef struct CopyStats
{
	uint64_t startTime;
	uint64_t bytesTransmitted;
} CopyStats;

typedef bool (CopyStatsCallback)(void *context, CopyStats *stats);

bool pg_copy(PGSQL *src, PGSQL *dst,
			 CopyArgs *args, CopyStats *stats,
			 void *context, CopyStatsCallback *callback);

bool pg_copy_from_stdin(PGSQL *pgsql, const char *qname);
bool pg_copy_row_from_stdin(PGSQL *pgsql, char *fmt, ...);
bool pg_copy_end(PGSQL *pgsql);

bool pgsql_get_sequence(PGSQL *pgsql, const char *qname,
						int64_t *lastValue,
						bool *isCalled);

bool pgsql_set_gucs(PGSQL *pgsql, GUC *settings);

bool pg_copy_large_object(PGSQL *src,
						  PGSQL *dst,
						  bool dropIfExists,
						  uint32_t oid,
						  uint64_t *bytesTransmitted);

/*
 * Maximum length of serialized pg_lsn value
 * It is taken from postgres file pg_lsn.c.
 * It defines MAXPG_LSNLEN to be 17 and
 * allocates a buffer 1 byte larger. We
 * went for 18 to make buffer allocation simpler.
 */
#define PG_LSN_MAXLENGTH 18

/*
 * TimelineHistoryEntry is taken from Postgres definitions and adapted to
 * client-size code where we don't have all the necessary infrastruture. In
 * particular we don't define a XLogRecPtr data type nor do we define a
 * TimelineID data type.
 *
 * Zero is used indicate an invalid pointer. Bootstrap skips the first possible
 * WAL segment, initializing the first WAL page at WAL segment size, so no XLOG
 * record can begin at zero.
 */
#define InvalidXLogRecPtr 0
#define XLogRecPtrIsInvalid(r) ((r) == InvalidXLogRecPtr)

typedef struct TimelineHistoryEntry
{
	uint32_t tli;
	uint64_t begin;         /* inclusive */
	uint64_t end;           /* exclusive, InvalidXLogRecPtr means infinity */
} TimelineHistoryEntry;


/*
 * The IdentifySystem contains information that is parsed from the
 * IDENTIFY_SYSTEM replication command, and then the TIMELINE_HISTORY result.
 */
typedef struct IdentifySystem
{
	uint64_t identifier;
	uint32_t timeline;
	char xlogpos[PG_LSN_MAXLENGTH];
	char dbname[NAMEDATALEN];
	TimelineHistoryEntry currentTimeline;
	char timelineHistoryFilename[MAXPGPATH];
} IdentifySystem;


/*
 * Logical Decoding support.
 */
typedef enum
{
	STREAM_PLUGIN_UNKNOWN = 0,
	STREAM_PLUGIN_TEST_DECODING,
	STREAM_PLUGIN_WAL2JSON,
	STREAM_PLUGIN_PGOUTPUT
} StreamOutputPlugin;

typedef struct LogicalTrackLSN
{
	XLogRecPtr written_lsn;
	XLogRecPtr flushed_lsn;
	XLogRecPtr applied_lsn;
} LogicalTrackLSN;


typedef struct LogicalStreamContext
{
	void *private;

	XLogRecPtr cur_record_lsn;
	int timeline;

	const char *buffer;         /* expose internal buffer */
	int bufferLen;              /* byte length of binary plugin data in buffer */
	StreamOutputPlugin plugin;

	bool forceFeedback;

	TimestampTz now;
	TimestampTz lastFeedbackSync;
	TimestampTz sendTime;
	XLogRecPtr endpos;          /* might be update at runtime */

	LogicalTrackLSN *tracking;  /* expose LogicalStreamClient.current */

	/* This is set to true when the connection running START_REPLICATION
	 * is interrupted and we need to retry.
	 *
	 * This is majorly used to add synthetic ROLLBACK statements to the
	 * JSON file to avoid partial transactions.
	 */
	bool onRetry;
} LogicalStreamContext;


typedef bool (*LogicalStreamReceiver) (LogicalStreamContext *context);

typedef struct LogicalStreamClient
{
	PGSQL pgsql;
	IdentifySystem system;
	char cdcPathDir[MAXPGPATH];

	char slotName[NAMEDATALEN];

	StreamOutputPlugin plugin;
	KeyVal pluginOptions;

	XLogRecPtr startpos;
	XLogRecPtr endpos;

	TimestampTz now;
	TimestampTz last_status;
	TimestampTz last_fsync;

	LogicalTrackLSN current;    /* updated at receive time */
	LogicalTrackLSN feedback;   /* updated at feedback sending time */

	LogicalStreamReceiver writeFunction;
	LogicalStreamReceiver flushFunction;
	LogicalStreamReceiver closeFunction;
	LogicalStreamReceiver feedbackFunction;
	LogicalStreamReceiver keepaliveFunction;

	int fsync_interval;
	int standby_message_timeout;
} LogicalStreamClient;


bool pgsql_init_stream(LogicalStreamClient *client,
					   const char *pguri,
					   const char *cdcPathDir,
					   StreamOutputPlugin plugin,
					   const char *slotName,
					   XLogRecPtr startpos,
					   XLogRecPtr endpos);

StreamOutputPlugin OutputPluginFromString(char *plugin);
char * OutputPluginToString(StreamOutputPlugin plugin);

typedef struct ReplicationSlot
{
	char slotName[BUFSIZE];
	uint64_t lsn;
	char snapshot[BUFSIZE];
	StreamOutputPlugin plugin;
	bool wal2jsonNumericAsString;
	char publicationName[BUFSIZE]; /* pgoutput publication name */
	bool publicationAutoManaged;   /* true = pgcopydb created it, drop on cleanup */
} ReplicationSlot;

bool pgsql_create_logical_replication_slot(LogicalStreamClient *client,
										   ReplicationSlot *slot);

bool pgsql_timestamptz_to_string(TimestampTz ts, char *str, size_t size);

bool pgsql_start_replication(LogicalStreamClient *client);
bool pgsql_stream_logical(LogicalStreamClient *client,
						  LogicalStreamContext *context);

bool pgsql_get_block_size(PGSQL *pgsql, int *blockSize);

bool pgsql_replication_origin_oid(PGSQL *pgsql, char *nodeName, uint32_t *oid);
bool pgsql_replication_origin_create(PGSQL *pgsql, char *nodeName);
bool pgsql_replication_origin_drop(PGSQL *pgsql, char *nodeName);
bool pgsql_replication_origin_session_setup(PGSQL *pgsql, char *nodeName);

bool pgsql_replication_origin_xact_setup(PGSQL *pgsql,
										 char *origin_lsn,
										 char *origin_timestamp);

bool pgsql_replication_origin_advance(PGSQL *pgsql, char *nodeName, char *lsn);

bool pgsql_replication_origin_progress(PGSQL *pgsql,
									   char *nodeName,
									   bool flush,
									   uint64_t *lsn);

bool pgsql_replication_slot_exists(PGSQL *pgsql,
								   const char *slotName,
								   bool *slotExists,
								   uint64_t *lsn);

bool pgsql_drop_replication_slot(PGSQL *pgsql, const char *slotName);

/* Forward declaration; full definition in filtering.h */
struct SourceFilters;

bool pgsql_create_publication(PGSQL *pgsql, const char *pubName,
							  struct SourceFilters *filters);
bool pgsql_drop_publication(PGSQL *pgsql, const char *pubName);

bool pgsql_role_exists(PGSQL *pgsql, const char *roleName, bool *exists);

bool pgsql_configuration_exists(PGSQL *pgsql, const char *setconfig, bool *exists);

bool pgsql_table_exists(PGSQL *pgsql,
						uint32_t oid,
						const char *nspname,
						const char *relname,
						bool *exists);

bool pgsql_current_wal_flush_lsn(PGSQL *pgsql, uint64_t *lsn);

char * pgsql_escape_identifier(PGSQL *pgsql, char *src);
#endif /* PGSQL_H */
