Support TSQL syntax for geospatial functions

contrib/babelfishpg_tsql/src/multidb.c

@@ -1277,12 +1277,6 @@ get_physical_schema_name_by_mode(char *db_name, const char *schema_name, Migrati
 
 			snprintf(result, (MAX_BBF_NAMEDATALEND), "%s_%s", db_name, name);
 		}
-		else if (!DbidIsValid(get_db_id(db_name)))
-		{
-			ereport(ERROR,
-					(errcode(ERRCODE_UNDEFINED_DATABASE),
-					 errmsg("database \"%s\" does not exist. Make sure that the name is entered correctly.", db_name)));
-		}
 		else
 		{
 			/* all schema names are not prepended with db name on single-db */

contrib/babelfishpg_tsql/src/pl_comp.c

@@ -27,6 +27,7 @@
 #include "catalog/pg_type.h"
 #include "funcapi.h"
 #include "nodes/makefuncs.h"
+#include "parser/parse_expr.h"
 #include "parser/parse_relation.h"
 #include "parser/parse_type.h"
 #include "utils/builtins.h"
@@ -44,6 +45,8 @@
 #include "codegen.h"
 #include "iterative_exec.h"
 #include "multidb.h"
+#include "session.h"
+#include "catalog.h"
 
 /* ----------
  * Our own local and global variables
@@ -104,6 +107,12 @@ static const ExceptionLabelMap exception_label_map[] = {
 
 static int	cur_handle_id = 1;
 
+/*
+ * Global variable used to prevent recurring check for TSQL functional call for Geospatial types
+ * Eg: table.col.stx  (here stx will be considered as a Geospatial function)
+ */
+static bool is_geospatial_function = false;
+
 /* ----------
  * static prototypes
  * ----------
@@ -119,6 +128,7 @@ static void add_dummy_return(PLtsql_function *function);
 static void add_decl_table(PLtsql_function *function, int tbl_dno, char *tbl_typ);
 static Node *pltsql_pre_column_ref(ParseState *pstate, ColumnRef *cref);
 static Node *pltsql_post_column_ref(ParseState *pstate, ColumnRef *cref, Node *var);
+static List *pltsql_post_func_ref(ParseState *pstate, FuncCall *fn, List *fargs);
 static void pltsql_post_expand_star(ParseState *pstate, ColumnRef *cref, List *l);
 static Node *pltsql_param_ref(ParseState *pstate, ParamRef *pref);
 static Node *resolve_column_ref(ParseState *pstate, PLtsql_expr *expr,
@@ -142,6 +152,9 @@ static void pltsql_HashTableInsert(PLtsql_function *function,
 								   PLtsql_func_hashkey *func_key);
 static void pltsql_HashTableDelete(PLtsql_function *function);
 static void delete_function(PLtsql_function *func);
+static Node *resolve_geospatial_col_ref(ParseState *pstate, ColumnRef *cref);
+static List *resolve_geospatial_func_ref(ParseState *pstate, FuncCall *fn, List *fargs);
+static char *extract_logical_schema_name(char *db_name, char *schema_name, MigrationMode mode);
 
 extern Portal ActivePortal;
 extern bool pltsql_function_parse_error_transpose(const char *prosrc);
@@ -1479,6 +1492,7 @@ pltsql_parser_setup(struct ParseState *pstate, PLtsql_expr *expr)
 {
 	pstate->p_pre_columnref_hook = pltsql_pre_column_ref;
 	pstate->p_post_columnref_hook = pltsql_post_column_ref;
+	pstate->p_post_funcref_hook = pltsql_post_func_ref;
 	pstate->p_post_expand_star_hook = pltsql_post_expand_star;
 	pstate->p_paramref_hook = pltsql_param_ref;
 	/* no need to use p_coerce_param_hook */
@@ -1507,6 +1521,7 @@ pltsql_post_column_ref(ParseState *pstate, ColumnRef *cref, Node *var)
 {
 	PLtsql_expr *expr = (PLtsql_expr *) pstate->p_ref_hook_state;
 	Node	   *myvar;
+	Node       *geovar;
 
 	if (expr->func->resolve_option == PLTSQL_RESOLVE_VARIABLE)
 		return NULL;			/* we already found there's no match */
@@ -1539,6 +1554,16 @@ pltsql_post_column_ref(ParseState *pstate, ColumnRef *cref, Node *var)
 				 errdetail("It could refer to either a PL/tsql variable or a table column."),
 				 parser_errposition(pstate, cref->location)));
 	}
+	/*
+	 * We check if the current colRef represents a geospatial function
+	 * Eg: SELECT col.STX from t1;
+	 * where col is a column of table t1
+	 */
+	else if (var == NULL)
+	{
+		if ((geovar = resolve_geospatial_col_ref(pstate, cref)) != NULL)
+			return geovar;
+	}
 
 	return myvar;
 }
@@ -1642,6 +1667,200 @@ pltsql_param_ref(ParseState *pstate, ParamRef *pref)
 	return make_datum_param(expr, nse->itemno, pref->location);
 }
 
+/*
+ * pltsql_post_func_ref		parser callback for FuncRefs to check and modify Geospatial function call
+ * We are returning a List * of function args because in the cases we are concerned with geospatial types
+ * we are just trying to modify the arg list for functions, which is then passed to ParseFuncOrColumn in transformFuncCall.
+ */
+static List *
+pltsql_post_func_ref(ParseState *pstate, FuncCall *fn, List *fargs)
+{
+	List *ret = resolve_geospatial_func_ref(pstate, fn, fargs);
+	if (ret)
+		return ret;
+
+	return fargs;
+}
+
+/*
+ * This function reverts schema name, which is already rewritten by rewrite_column_refs or rewrite_plain_name functions
+ * This is required so that parser identifies the name for Geospatial Function call
+ * Eg: select t2.geom.STX from t2; (where 'geom' is geometry column in table t2)
+ * In this Case, 't2.geom.STX' will rewritten as 'master_t2.geom.STX'
+ * because t2 is considered as Schema name and not table name because of 3 part call. So, we have to resolve it
+ */
+static char *
+extract_logical_schema_name(char *db_name, char *physical_schema_name, MigrationMode mode)
+{
+
+	if (SINGLE_DB == mode)
+	{
+		if ((strlen(db_name) == 6 && (strncmp(db_name, "master", 6) == 0)) ||
+			(strlen(db_name) == 6 && (strncmp(db_name, "tempdb", 6) == 0)) ||
+			(strlen(db_name) == 4 && (strncmp(db_name, "msdb", 4) == 0)))
+		{
+			return physical_schema_name + strlen(db_name) + 1;
+		}
+		else
+		{
+			/* all schema names are not prepended with db name on single-db */
+			return physical_schema_name;
+		}
+	}
+	else
+	{
+		return physical_schema_name + strlen(db_name) + 1;
+	}
+
+	return physical_schema_name;
+}
+
+/*
+ * This function identifies and modifies geospatial function call for
+ * STX, STY, LAT, LONG functions
+ */
+static Node *
+resolve_geospatial_col_ref(ParseState *pstate, ColumnRef *cref)
+{
+	Node       *last_field;		/* represents the geospatial function name */
+	Node       *col;
+
+	/* if there is only one field then it cannot be a geospatial function call */
+	if (list_length(cref->fields) <= 1){
+		is_geospatial_function = false;	/* resetting the global variable */
+		return NULL;
+	}
+
+	last_field = (Node *) llast(cref->fields);
+	/* 
+	 * checking if last field matches any of the geospatial functions - stx, sty, lat, long
+	 * and also checking if is_geospatial_function variable is false to avoid entering the loop if it is a column rather than function
+	 */
+	if (IsA(last_field, String) && !is_geospatial_function &&
+		(pg_strcasecmp(strVal(last_field), "stx") == 0 ||
+		pg_strcasecmp(strVal(last_field), "sty") == 0 ||
+		pg_strcasecmp(strVal(last_field), "lat") == 0 ||
+		pg_strcasecmp(strVal(last_field), "long") == 0))
+	{
+		FuncExpr   *funcexpr = makeNode(FuncExpr);	/* func expression to modify colref to funcref */
+		char *name = strVal(last_field);
+		Oid funcid;
+		Oid dataTypeId;
+		ColumnRef *colref = makeNode(ColumnRef);	/* new colref to avoid making changes to existing structure */
+
+		colref->fields = list_copy(cref->fields);
+		/* if length is 3 then colref would have been modified by rewrite_column_refs function */
+		if(list_length(colref->fields) == 3)
+		{
+			Node	   *schema = (Node *) linitial(colref->fields);
+			char	   *cur_db = get_cur_db_name();
+			String	   *new_schema;
+
+			if (!is_shared_schema(strVal(schema)))
+			{
+				new_schema = makeString(extract_logical_schema_name(cur_db, strVal(schema), get_migration_mode()));
+				colref->fields = list_delete_first(colref->fields);
+				colref->fields = lcons(new_schema, colref->fields);
+			}
+		}
+		colref->fields = list_delete_last(colref->fields);	/* deletes the function name and check if the remaining colref is a valid column reference */
+		is_geospatial_function = true;
+		PG_TRY();
+		{
+			col = transformExpr(pstate, (Node *) colref, pstate->p_expr_kind);	/* Recursive call to check for valid column reference */
+		}
+		PG_CATCH();
+		{
+			is_geospatial_function = false;	/* resetting the global variable */
+			return NULL;	/* if not a valid column ref */
+		}
+		PG_END_TRY();
+
+		is_geospatial_function = false;	/* resetting the global variable */
+		dataTypeId = ((Var *)col)->vartype;
+		/* Modifying colref to funcref since a valid Geospatial function call is identified */
+		funcid = LookupFuncName(list_make2(makeString("sys"), makeString(name)), 1, &dataTypeId, false);
+		funcexpr = makeFuncExpr(funcid, get_func_rettype(funcid), list_make1(col), 0, 0, COERCE_EXPLICIT_CALL);
+		funcexpr->location = cref->location;
+
+		return (Node *) funcexpr;
+	}
+
+	is_geospatial_function = false;	/* resetting the global variable */
+
+	return NULL;
+}
+
+/*
+ * This function identifies and modifies geospatial function call for
+ * STDistance(), STAsText(), STAsBinary() functions
+ */
+static List *resolve_geospatial_func_ref(ParseState *pstate, FuncCall *fn, List *fargs)
+{
+	Node *fname;	/* represents the geospatial function name */
+
+	if(list_length(fn->funcname) > 1 && list_length(fn->args) <= 1 
+		&& fn->agg_order == NULL && fn->agg_filter == NULL 
+		&& !fn->agg_within_group && !fn->agg_star 
+		&& !fn->agg_distinct && !fn->func_variadic)
+	{
+		Node *col;
+		bool flag = true;
+		fname = (Node *) llast(fn->funcname);
+		/* 
+		 * checking if last field matches any of the geospatial functions - stdistance, stastext, stasbinary
+		 * and also checking if is_geospatial_function variable is false to avoid entering the loop if it is not a function
+		 */
+		if (IsA(fname, String) && !is_geospatial_function &&
+				((pg_strcasecmp(strVal(fname), "stdistance") == 0 && list_length(fn->args) == 1) ||
+				(pg_strcasecmp(strVal(fname), "stastext") == 0 && list_length(fn->args) == 0) || 
+				(pg_strcasecmp(strVal(fname), "stasbinary") == 0 && list_length(fn->args) == 0)))
+		{
+			ColumnRef *cref = makeNode(ColumnRef);		/* new colref to avoid making changes to existing structure */
+			cref->fields = list_copy(fn->funcname);
+			cref->fields = list_delete_last(cref->fields);	/* deletes the function name and check if the remaining colref is a valid column reference */
+
+			/* if length is 1 then colref would have been modified by rewrite_plain_name function */
+			if(list_length(cref->fields) == 1)
+			{
+				Node       *schema = (Node *) linitial(cref->fields);
+				char       *cur_db = get_cur_db_name();
+				String     *new_schema;
+
+				if (!is_shared_schema(strVal(schema)))
+				{
+					new_schema = makeString(extract_logical_schema_name(cur_db, strVal(schema), get_migration_mode()));
+					cref->fields = list_delete_first(cref->fields);
+					cref->fields = lcons(new_schema, cref->fields);
+				}
+			}
+			is_geospatial_function = true;
+			PG_TRY();
+			{
+				col = transformExpr(pstate, (Node *) cref, pstate->p_expr_kind);	/* Recursive call to check for valid column reference */
+			}
+			PG_CATCH();
+			{
+				flag = false;	/* if not a valid column ref. Eg: table.col.stx (where col is not a coulmn of table) */
+			}
+			PG_END_TRY();
+
+			/* Modifying funcref since a valid Geospatial function call is identified */
+			if(flag)
+			{
+				is_geospatial_function = false;	/* resetting the global variable */
+				fn->args = lcons(cref, fn->args);
+				fargs = lcons(col, fargs);
+				fn->funcname = list_delete_first(fn->funcname);
+				return fargs;
+			}
+		}
+	}
+	is_geospatial_function = false;	/* resetting the global variable */
+
+	return NULL;
+}
+
 /*
  * resolve_column_ref		attempt to resolve a ColumnRef as a pltsql var
  *