github.com/mssola/todo@v0.0.0-20181029153210-d25348dc3f48/app/db.go (about)

     1  // Copyright (C) 2014-2017 Miquel Sabaté Solà <mikisabate@gmail.com>
     2  //
     3  // This Source Code Form is subject to the terms of the Mozilla Public
     4  // License, v. 2.0. If a copy of the MPL was not distributed with this
     5  // file, You can obtain one at http://mozilla.org/MPL/2.0/.
     6  
     7  package app
     8  
     9  import (
    10  	"database/sql"
    11  	"fmt"
    12  	"log"
    13  	"os"
    14  	"time"
    15  
    16  	"github.com/coopernurse/gorp"
    17  
    18  	// Blank import because we are using postgresql
    19  	_ "github.com/lib/pq"
    20  )
    21  
    22  // maxConnectionTries contains the number of connection attempts that this
    23  // application is going to make before panic'ing.
    24  const maxConnectionTries = 10
    25  
    26  // Db is a global instance that holds a connection to the DB. It gets
    27  // initialized after calling the InitDB function. You have to call CloseDB in
    28  // order to close the connection.
    29  var Db gorp.DbMap
    30  
    31  // EnvOrElse returns the value of the given environment variable. If this
    32  // environment variable is not set, then it returns the provided alternative
    33  // value.
    34  func EnvOrElse(name, value string) string {
    35  	if env := os.Getenv(name); env != "" {
    36  		return env
    37  	}
    38  	return value
    39  }
    40  
    41  // configURL returns the string being used to connect with our PostgreSQL
    42  // database.
    43  func configURL() string {
    44  	user := EnvOrElse("TODO_DB_USER", "postgres")
    45  	dbname := EnvOrElse("TODO_DB_NAME", "todo-dev")
    46  	password := EnvOrElse("TODO_DB_PASSWORD", "")
    47  	host := EnvOrElse("TODO_DB_HOST", "localhost")
    48  	sslmode := EnvOrElse("TODO_DB_SSLMODE", "disable")
    49  
    50  	str := "user=%s host=%s port=5432 dbname=%s sslmode=%s"
    51  	if password != "" {
    52  		str += " password=%s"
    53  		return fmt.Sprintf(str, user, host, dbname, sslmode, password)
    54  	}
    55  	return fmt.Sprintf(str, user, host, dbname, sslmode)
    56  }
    57  
    58  // establishConnection tries to establish a connection to the DB. It tries to
    59  // do so until maxConnectionTries is reached, at which point it panics.
    60  func establishConnection() *sql.DB {
    61  	var err error
    62  
    63  	str := configURL()
    64  	log.Printf("Trying with: '%s'", str)
    65  	d, err := sql.Open("postgres", str)
    66  
    67  	for i := 0; i < maxConnectionTries; i++ {
    68  		if err = d.Ping(); err == nil {
    69  			log.Printf("postgres: connection established.")
    70  			return d
    71  		}
    72  		if i < maxConnectionTries-1 {
    73  			log.Printf("postgres: ping failed: %v", err)
    74  			log.Printf("posgres: retrying in 5 seconds...")
    75  			time.Sleep(5 * time.Second)
    76  		}
    77  	}
    78  	log.Fatalf("postgres: could not establish connection with '%s'.", str)
    79  	return nil
    80  }
    81  
    82  // InitDB initializes the global DB connection.
    83  func InitDB() {
    84  	d := establishConnection()
    85  
    86  	Db = gorp.DbMap{Db: d, Dialect: gorp.PostgresDialect{}}
    87  	Db.AddTableWithName(User{}, "users")
    88  	Db.AddTableWithName(Topic{}, "topics")
    89  }
    90  
    91  // CloseDB close the global DB connection.
    92  func CloseDB() {
    93  	if err := Db.Db.Close(); err != nil {
    94  		log.Printf("Could not close database: %v", err)
    95  	}
    96  }
    97  
    98  // Exists returns true if there is a row in the given table that matches the
    99  // given id. It returns false otherwise.
   100  func Exists(name, id string) bool {
   101  	q := fmt.Sprintf("select count(*) from %v where id=$1", name)
   102  	c, err := Db.SelectInt(q, id)
   103  	return err == nil && c == 1
   104  }
   105  
   106  // Count the number of rows for the given table. Returns a 0 on error. I know
   107  // that this is not idiomatic, but it comes in handy in this case.
   108  func Count(name string) int64 {
   109  	count, err := Db.SelectInt("select count(*) from " + name)
   110  	if err != nil {
   111  		return 0
   112  	}
   113  	return count
   114  }