Support TSQL syntax for geospatial functions

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
@@ -74,6 +77,12 @@ MemoryContext pltsql_compile_tmp_cxt;
  */
 static HTAB *pltsql_HashTable = NULL;
 
+/* ----------
+ * Hash table for geospatial functions
+ * ----------
+ */
+static HTAB *ht_spatial_func_info = NULL;
+
 typedef struct pltsql_hashent
 {
 	PLtsql_func_hashkey key;
@@ -82,6 +91,13 @@ typedef struct pltsql_hashent
 
 #define FUNCS_PER_USER		128 /* initial table size */
 
+/* Macro to check if field is a geospatial function */
+#define IS_GEOSPATIAL_FIELD(field) \
+    (pg_strcasecmp(strVal(field), "stx") == 0 || \
+     pg_strcasecmp(strVal(field), "sty") == 0 || \
+     pg_strcasecmp(strVal(field), "lat") == 0 || \
+     pg_strcasecmp(strVal(field), "long") == 0)
+
 /* ----------
  * Lookup table for EXCEPTION condition names
  * ----------
@@ -97,13 +113,38 @@ static const ExceptionLabelMap exception_label_map[] = {
 	{NULL, 0}
 };
 
+typedef struct geospatial_hash_entry_t
+{
+	int key;
+	Oid funcid;
+	Oid	returntype;
+} geospatial_hash_entry_t;
+
+enum geospatial_type
+{
+	GEOM_STX,
+	GEOM_STY,
+	GEOG_LAT,
+	GEOG_LONG,
+	INVALID_TYPE
+};
+
+static Oid	geomId = InvalidOid;
+static Oid	geogId = InvalidOid;
+
 /* ----------
  * Current session's handler
  * ----------
  */
 
 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
  * ----------
@@ -142,6 +183,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 +1523,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_pre_funcref_hook = resolve_geospatial_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 */
@@ -1539,6 +1584,20 @@ 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)
+	{
+		Node       *geovar;
+		geovar = resolve_geospatial_col_ref(pstate, cref);
+		if(geovar != NULL){
+			printf("lol");
+		}
+			//return geovar;
+	}
 
 	return myvar;
 }
@@ -1633,15 +1692,255 @@ pltsql_param_ref(ParseState *pstate, ParamRef *pref)
 	snprintf(pname, sizeof(pname), "$%d", pref->number);
 
 	nse = pltsql_ns_lookup(expr->ns, false,
-						   pname, NULL, NULL,
-						   NULL);
+						pname, NULL, NULL,
+						NULL);
 
 	if (nse == NULL)
 		return NULL;			/* name not known to pltsql */
 
 	return make_datum_param(expr, nse->itemno, pref->location);
 }
 
+/*
+ * extract_logical_schema_name
+ * converts physical schema name to logical schema name using given db_name and migration mode.
+ * This can be used in places where we want to revert the schema rewrites which have happened earlier, such
+ * as 'master_t2.geom.STX' -> 't2.geom.STX'
+ */
+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 handles geospatial syntax such as STX, STY, LAT, LONG notations and converts them to function calls.
+ */
+static Node *
+resolve_geospatial_col_ref(ParseState *pstate, ColumnRef *cref)
+{
+	Node       *last_field;		/* represents the geospatial function name */
+	Node       *col;
+
+	/* Initializing Hash Table to store func and return id's for geospatial functions */
+	if (ht_spatial_func_info == NULL)
+	{
+		geospatial_hash_entry_t *t1 = (geospatial_hash_entry_t*)palloc(sizeof(geospatial_hash_entry_t));
+		int key = GEOM_STX;
+		HASHCTL		hashCtl;
+		geomId = GetSysCacheOid2(TYPENAMENSP, Anum_pg_type_oid, CStringGetDatum("geometry"), ObjectIdGetDatum(get_namespace_oid("sys", false)));
+		geogId = GetSysCacheOid2(TYPENAMENSP, Anum_pg_type_oid, CStringGetDatum("geography"), ObjectIdGetDatum(get_namespace_oid("sys", false)));
+		MemSet(&hashCtl, 0, sizeof(hashCtl));
+		hashCtl.keysize = sizeof(int);
+		hashCtl.entrysize = sizeof(geospatial_hash_entry_t);
+		ht_spatial_func_info = hash_create("T-SQL geospatial function info hash",
+										SPI_processed,
+										&hashCtl,
+										HASH_ELEM | HASH_BLOBS);
+		/* stx */
+		t1 = hash_search(ht_spatial_func_info, &key, HASH_ENTER, NULL);
+		t1->funcid = LookupFuncName(list_make2(makeString("sys"), makeString("stx")), 1, &geomId, false);
+		t1->returntype = get_func_rettype(t1->funcid);
+		/* sty */
+		key = GEOM_STY;
+		t1 = hash_search(ht_spatial_func_info, &key, HASH_ENTER, NULL);
+		t1->funcid = LookupFuncName(list_make2(makeString("sys"), makeString("sty")), 1, &geomId, false);
+		t1->returntype = get_func_rettype(t1->funcid);
+		/* lat */
+		key = GEOG_LAT;
+		t1 = hash_search(ht_spatial_func_info, &key, HASH_ENTER, NULL);
+		t1->funcid = LookupFuncName(list_make2(makeString("sys"), makeString("lat")), 1, &geogId, false);
+		t1->returntype = get_func_rettype(t1->funcid);
+		/* long */
+		key = GEOG_LONG;
+		t1 = hash_search(ht_spatial_func_info, &key, HASH_ENTER, NULL);
+		t1->funcid = LookupFuncName(list_make2(makeString("sys"), makeString("long")), 1, &geogId, false);
+		t1->returntype = get_func_rettype(t1->funcid);
+	}
+
+	/* 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);
+	/* 
+	 * check if last field matches any of the geospatial functions - stx, sty, lat, long
+	 * is_geospatial_function is used to make sure that we don't get stuck in infinite loop.
+	 */
+	if (IsA(last_field, String) && !is_geospatial_function &&
+		IS_GEOSPATIAL_FIELD(last_field))
+	{
+		FuncExpr   *funcexpr = makeNode(FuncExpr);	/* func expression to modify colref to funcref */
+		char *name = strVal(last_field);
+		Oid dataTypeId;
+		int key;
+		geospatial_hash_entry_t *t1;
+		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);
+			}
+			pfree(cur_db);
+		}
+		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 */
+			pfree(colref);
+			pfree(last_field);
+			return NULL;	/* if not a valid column ref */
+		}
+		PG_END_TRY();
+
+		is_geospatial_function = false;	/* resetting the global variable */
+		dataTypeId = ((Var *)col)->vartype;
+
+		/* check key for hash table search */
+		if (pg_strcasecmp(name, "stx") == 0 && (int)geomId == (int)dataTypeId)
+			key = GEOM_STX;
+		else if (pg_strcasecmp(name, "sty") == 0 && (int)geomId == (int)dataTypeId)
+			key = GEOM_STY;
+		else if (pg_strcasecmp(name, "lat") == 0 && (int)geogId == (int)dataTypeId)
+			key = GEOG_LAT;
+		else if (pg_strcasecmp(name, "long") == 0 && (int)geogId == (int)dataTypeId)
+			key = GEOG_LONG;
+		else
+			key = INVALID_TYPE;
+
+		t1 = hash_search(ht_spatial_func_info, &key, HASH_FIND, NULL);
+		/* Prepare a Funcexpr node for the corresponding geospatial function. */
+		/* if key is Invalid then we need to throw an error */
+		if (key == INVALID_TYPE)
+		{
+			ereport(ERROR,
+				(errcode(ERRCODE_UNDEFINED_FUNCTION),
+				 errmsg("function sys.%s(%s) does not exist", name,  format_type_be(dataTypeId))));
+		}
+
+		funcexpr = makeFuncExpr(t1->funcid, t1->returntype, list_make1(col), 0, 0, COERCE_EXPLICIT_CALL);
+
+		funcexpr->location = cref->location;
+
+		pfree(colref);
+		pfree(last_field);
+		return (Node *) funcexpr;
+	}
+
+	is_geospatial_function = false;	/* resetting the global variable */
+	pfree(last_field);
+	return NULL;
+}
+
+/*
+ * This function handles geospatial syntax such as STDistance, STAsText, STAsBinary notations and converts them to function calls.
+ */
+static List *
+resolve_geospatial_func_ref(ParseState *pstate, FuncCall *fn, List *fargs)
+{
+	return fargs;
+
+	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 *fname;	/* represents the geospatial function name */
+		Node *col;
+		fname = (Node *) llast(fn->funcname);
+		/* 
+		 * check if last field matches any of the geospatial functions - stdistance, stastext, stasbinary
+		 * is_geospatial_function is used to make sure that we don't get stuck in infinite loop.
+		 */
+		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);
+				}
+				pfree(cur_db);
+			}
+			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();
+			{
+				is_geospatial_function = false;	/* resetting the global variable */
+				pfree(cref);
+				return fargs;
+			}
+			PG_END_TRY();
+
+			/* Modify Funcexpr node for the corresponding geospatial function. */
+			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);
+			pfree(cref);
+			return fargs;
+		}
+	}
+	is_geospatial_function = false;	/* resetting the global variable */
+
+	return fargs;
+}
+
 /*
  * resolve_column_ref		attempt to resolve a ColumnRef as a pltsql var
  *

contrib/babelfishpg_tsql/src/pl_handler.c

@@ -487,6 +487,11 @@ pltsql_pre_parse_analyze(ParseState *pstate, RawStmt *parseTree)
 	if (prev_pre_parse_analyze_hook)
 		prev_pre_parse_analyze_hook(pstate, parseTree);
 
+	// if (pstate->p_pre_columnref_hook == NULL && pstate->p_post_columnref_hook == NULL 
+	// 	&& pstate->p_pre_funcref_hook == NULL && pstate->p_post_expand_star_hook == NULL 
+	// 	&& pstate->p_paramref_hook == NULL && parseTree->stmt->type == T_SelectStmt && sql_dialect == SQL_DIALECT_TSQL)
+	// 	pltsql_parser_setup(pstate, NULL); 
+
 	switch (parseTree->stmt->type)
 	{
 		case T_InsertStmt: