github.com/iron-io/functions@v0.0.0-20180820112432-d59d7d1c40b2/examples/postgres/func.go (about)

     1  package main
     2  
     3  import (
     4  	"bytes"
     5  	"database/sql"
     6  	"encoding/json"
     7  	"io/ioutil"
     8  	"log"
     9  	"os"
    10  	"strconv"
    11  
    12  	"github.com/pkg/errors"
    13  
    14  	_ "github.com/lib/pq"
    15  )
    16  
    17  var (
    18  	// command to execute, 'SELECT' or 'INSERT'
    19  	command = os.Getenv("HEADER_COMMAND")
    20  	// postgres host:port, e.g. 'postgres:5432'
    21  	server = os.Getenv("HEADER_SERVER")
    22  	// postgres table name
    23  	table = os.Getenv("HEADER_TABLE")
    24  )
    25  
    26  func main() {
    27  	req, err := ioutil.ReadAll(os.Stdin)
    28  	if err != nil {
    29  		log.Fatal(errors.Wrap(err, "failed to read stdin"))
    30  	}
    31  
    32  	db, err := sql.Open("postgres", "postgres://postgres@"+server+"?sslmode=disable")
    33  	if err != nil {
    34  		log.Println("Failed to connect to postgres server")
    35  		log.Fatal(err)
    36  		return
    37  	}
    38  
    39  	switch command {
    40  	case "SELECT":
    41  		if resp, err := selectCommand(req, db); err != nil {
    42  			log.Fatal(errors.Wrap(err, "select command failed"))
    43  		} else {
    44  			log.Println(resp)
    45  		}
    46  	case "INSERT":
    47  		if err := insertCommand(req, db); err != nil {
    48  			log.Fatal(errors.Wrap(err, "insert command failed"))
    49  		}
    50  	default:
    51  		log.Fatalf("invalid command: %q", command)
    52  	}
    53  }
    54  
    55  func selectCommand(req []byte, db *sql.DB) (string, error) {
    56  	// Parse request JSON
    57  	var params map[string]interface{}
    58  	if err := json.Unmarshal(req, &params); err != nil {
    59  		return "", errors.Wrap(err, "failed to parse json")
    60  	}
    61  
    62  	// Build query and gather arguments
    63  	var query bytes.Buffer
    64  	var args []interface{}
    65  
    66  	query.WriteString("SELECT json_agg(t) FROM (SELECT * FROM ")
    67  	query.WriteString(table)
    68  	query.WriteString(" WHERE")
    69  	first := true
    70  	arg := 1
    71  	for k, v := range params {
    72  		args = append(args, v)
    73  
    74  		if !first {
    75  			query.WriteString(" AND")
    76  		}
    77  		query.WriteString(" ")
    78  		query.WriteString(k)
    79  		query.WriteString("=$")
    80  		query.WriteString(strconv.Itoa(arg))
    81  		arg += 1
    82  		first = false
    83  	}
    84  	query.WriteString(") AS t")
    85  
    86  	// Execute query
    87  	r := db.QueryRow(query.String(), args...)
    88  	var resp string
    89  	if err := r.Scan(&resp); err != nil {
    90  		return "", errors.Wrap(err, "failed to execute select query")
    91  	}
    92  
    93  	return resp, nil
    94  }
    95  
    96  func insertCommand(req []byte, db *sql.DB) error {
    97  	q := "INSERT INTO " + table + " SELECT * FROM json_populate_record(null::" + table + ", $1)"
    98  	_, err := db.Exec(q, req)
    99  	if err != nil {
   100  		return errors.Wrap(err, "Failed to execute insert query")
   101  	}
   102  	return nil
   103  }