github.com/seeker-insurance/kit@v0.0.13/db/psql/psql.go (about)

     1  package psql
     2  
     3  import (
     4  	"database/sql/driver"
     5  	"encoding/json"
     6  	"errors"
     7  	"fmt"
     8  
     9  	"github.com/jinzhu/gorm"
    10  	//register postgres dialect
    11  	_ "github.com/jinzhu/gorm/dialects/postgres"
    12  	"github.com/spf13/cobra"
    13  	"github.com/spf13/viper"
    14  )
    15  
    16  type Scope func(db *gorm.DB) *gorm.DB
    17  type JsonB map[string]interface{}
    18  
    19  const nullDataValue = "null"
    20  
    21  var (
    22  	DB      *gorm.DB
    23  	DBError error
    24  )
    25  
    26  func init() {
    27  	cobra.OnInitialize(ConnectDB)
    28  }
    29  
    30  func ConnectDB() {
    31  	viper.SetDefault("database_scheme", "postgres")
    32  	scheme := viper.GetString("database_scheme")
    33  	url := viper.GetString("database_url")
    34  
    35  	if len(url) == 0 {
    36  		DBError = errors.New("Missing database_url")
    37  	} else {
    38  		DB, DBError = gorm.Open(scheme, url)
    39  		// DB.LogMode(true)
    40  	}
    41  }
    42  
    43  func (j JsonB) Value() (driver.Value, error) {
    44  	return json.Marshal(j)
    45  }
    46  func (j *JsonB) Scan(src interface{}) error {
    47  	source, ok := src.([]byte)
    48  	if !ok {
    49  		return errors.New("Type assertion .([]byte) failed.")
    50  	}
    51  
    52  	if sourceIsNull(source) {
    53  		return nil
    54  	}
    55  
    56  	var i interface{}
    57  	err := json.Unmarshal(source, &i)
    58  	if err != nil {
    59  		return err
    60  	}
    61  
    62  	*j, ok = i.(map[string]interface{})
    63  	if !ok {
    64  		return fmt.Errorf("type assertion .(map[string]interface{}) failed. got %s", source)
    65  	}
    66  
    67  	return nil
    68  }
    69  
    70  func UndeleteOrCreate(model interface{}, query string, args ...interface{}) error {
    71  	tx := DB.Begin()
    72  	tx = tx.Unscoped().Model(model).Where(query, args...).UpdateColumn("deleted_at", nil)
    73  
    74  	if err := tx.Error; err != nil {
    75  		tx.Rollback()
    76  		return err
    77  	}
    78  
    79  	if tx.RowsAffected > 1 {
    80  		tx.Rollback()
    81  		return fmt.Errorf("UndeleteOrCreate can only undelete one record, you're trying to undelete %v", tx.RowsAffected)
    82  	}
    83  
    84  	if tx.RowsAffected < 1 {
    85  		if err := tx.Create(model).Error; err != nil {
    86  			tx.Rollback()
    87  			return err
    88  		}
    89  	}
    90  	if err := tx.Commit().Error; err != nil {
    91  		tx.Rollback()
    92  		return err
    93  	}
    94  	return nil
    95  }
    96  
    97  func DeleteModelWithAssociations(value interface{}, associations ...string) error {
    98  	tx := DB.Begin()
    99  
   100  	if tx.Error != nil {
   101  		return tx.Error
   102  	}
   103  
   104  	for _, a := range associations {
   105  		if err := tx.Model(value).Association(a).Clear().Error; err != nil {
   106  			return err
   107  		}
   108  	}
   109  
   110  	tx = tx.Delete(value)
   111  
   112  	if tx.Error != nil {
   113  		tx.Rollback()
   114  		return tx.Error
   115  	}
   116  
   117  	if tx.RowsAffected < 1 {
   118  		tx.Rollback()
   119  		return gorm.ErrRecordNotFound
   120  	}
   121  
   122  	if err := tx.Commit().Error; err != nil {
   123  		tx.Rollback()
   124  		return err
   125  	}
   126  	return nil
   127  }
   128  
   129  func Transact(db *gorm.DB, txFunc func(*gorm.DB) error) (err error) {
   130  	tx := db.Begin()
   131  	defer func() {
   132  		if p := recover(); p != nil {
   133  			tx.Rollback()
   134  			panic(p)
   135  		} else if err != nil {
   136  			tx.Rollback()
   137  		} else {
   138  			tx.Commit()
   139  		}
   140  	}()
   141  	err = txFunc(tx)
   142  	return err
   143  }
   144  
   145  func DbWithError(e error) *gorm.DB {
   146  	db := DB.New()
   147  	db.Error = e
   148  	return db
   149  }
   150  
   151  func sourceIsNull(b []byte) bool {
   152  	return string(b) == nullDataValue
   153  }