github.com/orofarne/hammy@v0.0.0-20130409105742-374fadfd6ecb/src/hammy/mysql_saver.go (about)

     1  package hammy
     2  
     3  import (
     4  	"fmt"
     5  	"log"
     6  	"strings"
     7  	"database/sql"
     8  	_ "github.com/ziutek/mymysql/godrv"
     9  )
    10  
    11  // Driver for saving historical data in MySQL database
    12  // It's assumes the tables structure like this:
    13  //
    14  //  CREATE TABLE `history_host` (
    15  //    `id` int(11) NOT NULL AUTO_INCREMENT,
    16  //    `name` varchar(255) NOT NULL,
    17  //    PRIMARY KEY (`id`),
    18  //    UNIQUE KEY `by_name` (`name`)
    19  //  ) ENGINE=InnoDB DEFAULT CHARSET=utf8
    20  //
    21  //  CREATE TABLE `history_item` (
    22  //    `id` int(11) NOT NULL AUTO_INCREMENT,
    23  //    `host_id` int(11) NOT NULL,
    24  //    `name` varchar(255) NOT NULL,
    25  //    PRIMARY KEY (`id`),
    26  //    UNIQUE KEY `by_name` (`host_id`, `name`)
    27  //  ) ENGINE=InnoDB DEFAULT CHARSET=utf8
    28  //
    29  //  CREATE TABLE `history` (
    30  //    `item_id` int(11) NOT NULL,
    31  //    `timestamp` DATETIME NOT NULL,
    32  //    `value` DOUBLE NOT NULL,
    33  //    PRIMARY KEY (`item_id`, `timestamp`)
    34  //  ) ENGINE=InnoDB DEFAULT CHARSET=utf8
    35  //
    36  //  CREATE TABLE `history_log` (
    37  //    `item_id` int(11) NOT NULL,
    38  //    `timestamp` DATETIME NOT NULL,
    39  //    `value` TEXT NOT NULL,
    40  //    PRIMARY KEY (`item_id`, `timestamp`)
    41  //  ) ENGINE=InnoDB DEFAULT CHARSET=utf8
    42  //
    43  type MySQLSaver struct {
    44  	db *sql.DB
    45  	tableName string
    46  	logTableName string
    47  	hostTableName string
    48  	itemTableName string
    49  	pool chan int
    50  }
    51  
    52  func NewMySQLSaver(cfg Config) (s *MySQLSaver, err error) {
    53  	s = new(MySQLSaver)
    54  	s.db, err = sql.Open("mymysql", cfg.MySQLSaver.Database + "/" + cfg.MySQLSaver.User + "/" + cfg.MySQLSaver.Password)
    55  	if err != nil {
    56  		return
    57  	}
    58  
    59  	s.tableName = cfg.MySQLSaver.Table
    60  	s.logTableName = cfg.MySQLSaver.LogTable
    61  	s.hostTableName = cfg.MySQLSaver.HostTable
    62  	s.itemTableName = cfg.MySQLSaver.ItemTable
    63  
    64  	s.pool = make(chan int, cfg.MySQLSaver.MaxConn)
    65  	for i := 0; i < cfg.MySQLSaver.MaxConn; i++ {
    66  		s.pool <- 1
    67  	}
    68  
    69  	return
    70  }
    71  
    72  func (s *MySQLSaver) getOrCreateHost(host string) (int, error) {
    73  	sqlp := fmt.Sprintf("SELECT `id` FROM `%s` WHERE `name` = ?", s.hostTableName)
    74  	row := s.db.QueryRow(sqlp, host)
    75  
    76  	var hostId int
    77  	err := row.Scan(&hostId)
    78  
    79  	switch err {
    80  		case nil:
    81  			// Do nothing
    82  		case sql.ErrNoRows:
    83  			sqlp = fmt.Sprintf("INSERT INTO `%s` SET `name` = ?", s.hostTableName)
    84  			_, err := s.db.Exec(sqlp, host)
    85  			if err != nil {
    86  				return 0, err
    87  			}
    88  
    89  			row = s.db.QueryRow("SELECT LAST_INSERT_ID()")
    90  			err = row.Scan(&hostId)
    91  			if err != nil {
    92  				return 0, err
    93  			}
    94  		default:
    95  			return 0, err
    96  	}
    97  
    98  	return hostId, nil
    99  }
   100  
   101  func (s *MySQLSaver) getOrCreateItem(host_id int, item string) (int, error) {
   102  	sqlp := fmt.Sprintf("SELECT `id` FROM `%s` WHERE `host_id` = ? AND `name` = ?", s.itemTableName)
   103  	row := s.db.QueryRow(sqlp, host_id, item)
   104  
   105  	var itemId int
   106  	err := row.Scan(&itemId)
   107  
   108  	switch err {
   109  		case nil:
   110  			// Do nothing
   111  		case sql.ErrNoRows:
   112  			sqlp = fmt.Sprintf("INSERT INTO `%s` SET `host_id` = ?, `name` = ?", s.itemTableName)
   113  			_, err := s.db.Exec(sqlp, host_id, item)
   114  			if err != nil {
   115  				return 0, err
   116  			}
   117  
   118  			row = s.db.QueryRow("SELECT LAST_INSERT_ID()")
   119  			err = row.Scan(&itemId)
   120  			if err != nil {
   121  				return 0, err
   122  			}
   123  		default:
   124  			return 0, err
   125  	}
   126  
   127  	return itemId, nil
   128  }
   129  
   130  func (s *MySQLSaver) Push(data *IncomingData) {
   131  	// Pool limits
   132  	<- s.pool
   133  	defer func() {
   134  		s.pool <- 1
   135  	}()
   136  
   137  	for hostK, hostV := range *data {
   138  		for itemK, itemV := range hostV {
   139  			for _, v := range itemV {
   140  				hostId, err := s.getOrCreateHost(hostK)
   141  				if err != nil {
   142  					log.Printf("MySQLSaver can't get host_id: %v", err)
   143  					continue
   144  				}
   145  
   146  				itemId, err := s.getOrCreateItem(hostId, itemK)
   147  				if err != nil {
   148  					log.Printf("MySQLSaver can't get item_id: %v", err)
   149  					continue
   150  				}
   151  
   152  				var tName string
   153  				if strings.HasSuffix(itemK, "#log") {
   154  					tName = s.logTableName
   155  				} else {
   156  					tName = s.tableName
   157  				}
   158  				sqlp := fmt.Sprintf("INSERT INTO `%s` SET `item_id` = ?, `timestamp` = FROM_UNIXTIME(?), `value` = ?", tName)
   159  				_, err = s.db.Exec(sqlp, itemId, v.Timestamp, v.Value)
   160  				if err != nil {
   161  					log.Printf("MySQLSaver error: %v", err)
   162  				}
   163  			}
   164  		}
   165  	}
   166  }