Added functions to get the number of rows for a table or a query

odbc_fdw--0.1.0--0.2.0.sql

@@ -34,3 +34,11 @@ CREATE TYPE __tabledata AS (schema text, name text);
 CREATE OR REPLACE FUNCTION ODBCTablesList(text, integer DEFAULT 0) RETURNS SETOF __tabledata
 AS 'MODULE_PATHNAME', 'odbc_tables_list'
 LANGUAGE C STRICT;
+
+CREATE OR REPLACE FUNCTION ODBCTableSize(text, text) RETURNS INTEGER
+AS 'MODULE_PATHNAME', 'odbc_table_size'
+LANGUAGE C STRICT;
+
+CREATE OR REPLACE FUNCTION ODBCQuerySize(text, text) RETURNS INTEGER
+AS 'MODULE_PATHNAME', 'odbc_query_size'
+LANGUAGE C STRICT;

odbc_fdw--0.2.0--0.1.0.sql

@@ -30,5 +30,7 @@ ALTER FOREIGN DATA WRAPPER odbc_fdw
   VALIDATOR odbc_fdw_validator;
 
 DROP FUNCTION IF EXISTS ODBCTablesList(text, integer);
+DROP FUNCTION IF EXISTS ODBCTableSize(text, text);
+DROP FUNCTION IF EXISTS ODBCQuerySize(text, text);
 
 DROP TYPE __tabledata;

odbc_fdw--0.2.0.sql

@@ -34,3 +34,11 @@ CREATE TYPE __tabledata AS (schema text, name text);
 CREATE FUNCTION ODBCTablesList(text, integer DEFAULT 0) RETURNS SETOF __tabledata
 AS 'MODULE_PATHNAME', 'odbc_tables_list'
 LANGUAGE C STRICT;
+
+CREATE FUNCTION ODBCTableSize(text, text) RETURNS INTEGER
+AS 'MODULE_PATHNAME', 'odbc_table_size'
+LANGUAGE C STRICT;
+
+CREATE FUNCTION ODBCQuerySize(text, text) RETURNS INTEGER
+AS 'MODULE_PATHNAME', 'odbc_query_size'
+LANGUAGE C STRICT;

odbc_fdw.c

@@ -44,6 +44,9 @@
 #include "utils/array.h"
 #include "utils/builtins.h"
 #include "utils/rel.h"
+#include "nodes/nodes.h"
+#include "nodes/makefuncs.h"
+#include "nodes/pg_list.h"
 
 #include "optimizer/pathnode.h"
 #include "optimizer/restrictinfo.h"
@@ -141,10 +144,14 @@ typedef enum { TEXT_CONVERSION, HEX_CONVERSION, BIN_CONVERSION, BOOL_CONVERSION
 extern Datum odbc_fdw_handler(PG_FUNCTION_ARGS);
 extern Datum odbc_fdw_validator(PG_FUNCTION_ARGS);
 extern Datum odbc_tables_list(PG_FUNCTION_ARGS);
+extern Datum odbc_table_size(PG_FUNCTION_ARGS);
+extern Datum odbc_query_size(PG_FUNCTION_ARGS);
 
 PG_FUNCTION_INFO_V1(odbc_fdw_handler);
 PG_FUNCTION_INFO_V1(odbc_fdw_validator);
 PG_FUNCTION_INFO_V1(odbc_tables_list);
+PG_FUNCTION_INFO_V1(odbc_table_size);
+PG_FUNCTION_INFO_V1(odbc_query_size);
 
 /*
  * FDW callback routines
@@ -179,6 +186,7 @@ static void check_return(SQLRETURN ret, char *msg, SQLHANDLE handle, SQLSMALLINT
 static void odbcConnStr(StringInfoData *conn_str, odbcFdwOptions* options);
 static char* get_schema_name(odbcFdwOptions *options);
 static inline bool is_blank_string(const char *s);
+static Oid oid_from_server_name(char *serverName);
 
 /*
  * Check if string pointer is NULL or points to empty string
@@ -857,7 +865,7 @@ odbcGetTableSize(odbcFdwOptions* options, unsigned int *size)
 	}
 	else
 	{
-		elog(DEBUG1, "Oops!");
+		elog(WARNING, "Error getting the table %s size", options->table);
 	}
 
 	/* Free handles, and disconnect */
@@ -880,26 +888,6 @@ odbcGetTableSize(odbcFdwOptions* options, unsigned int *size)
 		SQLDisconnect(dbc);
 }
 
-/*
- * Get the list of tables for the current datasource
- */
-typedef struct {
-  SQLSMALLINT TargetType;
-  SQLPOINTER TargetValuePtr;
-  SQLINTEGER BufferLength;
-  SQLLEN StrLen_or_Ind;
-} DataBinding;
-
-typedef struct {
-  Oid serverOid;
-  DataBinding* tableResult;
-  SQLHSTMT stmt;
-  SQLCHAR schema;
-  SQLCHAR name;
-  SQLUINTEGER rowLimit;
-  SQLUINTEGER currentRow;
-} TableDataCtx;
-
 static int strtoint(const char *nptr, char **endptr, int base)
 {
   long val = strtol(nptr, endptr, base);
@@ -942,6 +930,67 @@ static Oid oid_from_server_name(char *serverName)
   return serverOid;
 }
 
+Datum
+odbc_table_size(PG_FUNCTION_ARGS)
+{
+  char *serverName = text_to_cstring(PG_GETARG_TEXT_PP(0));
+  char *tableName = text_to_cstring(PG_GETARG_TEXT_PP(1));
+  char *defname = "table";
+  int *tableSize;
+  List *tableOptions = NIL;
+  Node *val = (Node *) makeString(tableName);
+  DefElem *elem = (DefElem *) makeDefElem(defname, val);
+
+  tableOptions = lappend(tableOptions, elem);
+  Oid serverOid = oid_from_server_name(serverName);
+  odbcFdwOptions options;
+  odbcGetOptions(serverOid, tableOptions, &options);
+  odbcGetTableSize(&options, tableSize);
+
+  PG_RETURN_INT32(*tableSize);
+}
+
+Datum
+odbc_query_size(PG_FUNCTION_ARGS)
+{
+  char *serverName = text_to_cstring(PG_GETARG_TEXT_PP(0));
+  char *sqlQuery = text_to_cstring(PG_GETARG_TEXT_PP(1));
+  char *defname = "sql_query";
+  int *querySize;
+  List *queryOptions = NIL;
+  Node *val = (Node *) makeString(sqlQuery);
+  DefElem *elem = (DefElem *) makeDefElem(defname, val);
+
+  queryOptions = lappend(queryOptions, elem);
+  Oid serverOid = oid_from_server_name(serverName);
+  odbcFdwOptions options;
+  odbcGetOptions(serverOid, queryOptions, &options);
+  odbcGetTableSize(&options, querySize);
+
+  PG_RETURN_INT32(*querySize);
+}
+
+/*
+ * Get the list of tables for the current datasource
+ */
+typedef struct {
+  SQLSMALLINT TargetType;
+  SQLPOINTER TargetValuePtr;
+  SQLINTEGER BufferLength;
+  SQLLEN StrLen_or_Ind;
+} DataBinding;
+
+typedef struct {
+  Oid serverOid;
+  DataBinding* tableResult;
+  SQLHSTMT stmt;
+  SQLCHAR schema;
+  SQLCHAR name;
+  SQLUINTEGER rowLimit;
+  SQLUINTEGER currentRow;
+} TableDataCtx;
+
+
 Datum odbc_tables_list(PG_FUNCTION_ARGS)
 {
 	SQLHENV env;