feat: add support for single argument Postgres functions

core/internal/psql/columns.go

@@ -167,7 +167,7 @@ func (c *compilerContext) renderColumnTypename(sel *qcode.Select, ti *DBTableInf
 }
 
 func (c *compilerContext) renderColumnFunction(sel *qcode.Select, ti *DBTableInfo, col qcode.Column, columnsRendered int) error {
-	pl := funcPrefixLen(col.Name)
+	pl := funcPrefixLen(c.schema.fm, col.Name)
 	// if pl == 0 {
 	// 	//fmt.Fprintf(w, `'%s not defined' AS %s`, cn, col.Name)
 	// 	io.WriteString(c.w, `'`)

core/internal/psql/query.go

@@ -543,7 +543,7 @@ func (c *compilerContext) renderColumns(sel *qcode.Select, ti *DBTableInfo, skip
 	var cn string
 
 	for _, col := range sel.Cols {
-		if n := funcPrefixLen(col.Name); n != 0 {
+		if n := funcPrefixLen(c.schema.fm, col.Name); n != 0 {
 			if !sel.Functions {
 				continue
 			}
@@ -1193,7 +1193,7 @@ func (c *compilerContext) renderVal(ex *qcode.Exp, vars map[string]string, col *
 	io.WriteString(c.w, col.Type)
 }
 
-func funcPrefixLen(fn string) int {
+func funcPrefixLen(fm map[string]*DBFunction, fn string) int {
 	switch {
 	case strings.HasPrefix(fn, "avg_"):
 		return 4
@@ -1218,6 +1218,14 @@ func funcPrefixLen(fn string) int {
 	case strings.HasPrefix(fn, "var_samp_"):
 		return 9
 	}
+	fnLen := len(fn)
+
+	for k := range fm {
+		kLen := len(k)
+		if kLen < fnLen && k[0] == fn[0] && strings.HasPrefix(fn, k) && fn[kLen] == '_' {
+			return kLen + 1
+		}
+	}
 	return 0
 }
 

core/internal/psql/schema.go

@@ -11,6 +11,7 @@ type DBSchema struct {
 	ver int
 	t   map[string]*DBTableInfo
 	rm  map[string]map[string]*DBRel
+	fm  map[string]*DBFunction
 }
 
 type DBTableInfo struct {
@@ -58,6 +59,7 @@ func NewDBSchema(info *DBInfo, aliases map[string][]string) (*DBSchema, error) {
 	schema := &DBSchema{
 		t:  make(map[string]*DBTableInfo),
 		rm: make(map[string]map[string]*DBRel),
+		fm: make(map[string]*DBFunction, len(info.Functions)),
 	}
 
 	for i, t := range info.Tables {
@@ -81,6 +83,12 @@ func NewDBSchema(info *DBInfo, aliases map[string][]string) (*DBSchema, error) {
 		}
 	}
 
+	for k, f := range info.Functions {
+		if len(f.Params) == 1 {
+			schema.fm[strings.ToLower(f.Name)] = &info.Functions[k]
+		}
+	}
+
 	return schema, nil
 }
 
@@ -439,3 +447,11 @@ func (s *DBSchema) GetRel(child, parent string) (*DBRel, error) {
 	}
 	return rel, nil
 }
+
+func (s *DBSchema) GetFunctions() []*DBFunction {
+	var funcs []*DBFunction
+	for _, f := range s.fm {
+		funcs = append(funcs, f)
+	}
+	return funcs
+}

core/internal/psql/tables.go

@@ -10,10 +10,11 @@ import (
 )
 
 type DBInfo struct {
-	Version int
-	Tables  []DBTable
-	Columns [][]DBColumn
-	colmap  map[string]map[string]*DBColumn
+	Version   int
+	Tables    []DBTable
+	Columns   [][]DBColumn
+	Functions []DBFunction
+	colmap    map[string]map[string]*DBColumn
 }
 
 func GetDBInfo(db *sql.DB) (*DBInfo, error) {
@@ -51,6 +52,11 @@ func GetDBInfo(db *sql.DB) (*DBInfo, error) {
 		}
 	}
 
+	di.Functions, err = GetFunctions(db)
+	if err != nil {
+		return nil, err
+	}
+
 	return di, nil
 }
 
@@ -237,6 +243,64 @@ ORDER BY id;`
 	return cols, nil
 }
 
+type DBFunction struct {
+	Name   string
+	Params []DBFuncParam
+}
+
+type DBFuncParam struct {
+	ID   int
+	Name string
+	Type string
+}
+
+func GetFunctions(db *sql.DB) ([]DBFunction, error) {
+	sqlStmt := `
+SELECT 
+	routines.routine_name, 
+	parameters.specific_name,
+	parameters.data_type, 
+	parameters.parameter_name,
+	parameters.ordinal_position	
+FROM 
+	information_schema.routines
+RIGHT JOIN 
+	information_schema.parameters 
+	ON (routines.specific_name = parameters.specific_name and parameters.ordinal_position IS NOT NULL)	
+WHERE 
+	routines.specific_schema = 'public'
+ORDER BY 
+	routines.routine_name, parameters.ordinal_position;`
+
+	rows, err := db.Query(sqlStmt)
+	if err != nil {
+		return nil, fmt.Errorf("Error fetching functions: %s", err)
+	}
+	defer rows.Close()
+
+	var funcs []DBFunction
+	fm := make(map[string]int)
+
+	for rows.Next() {
+		var fn, fid string
+		fp := DBFuncParam{}
+
+		err = rows.Scan(&fn, &fid, &fp.Type, &fp.Name, &fp.ID)
+		if err != nil {
+			return nil, err
+		}
+
+		if i, ok := fm[fid]; ok {
+			funcs[i].Params = append(funcs[i].Params, fp)
+		} else {
+			funcs = append(funcs, DBFunction{Name: fn, Params: []DBFuncParam{fp}})
+			fm[fid] = len(funcs) - 1
+		}
+	}
+
+	return funcs, nil
+}
+
 // func GetValType(type string) qcode.ValType {
 // 	switch {
 // 		case "bigint", "integer", "smallint", "numeric", "bigserial":