Getting list of tables on servers and counts records
You can not select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
 
 
 

187 lines
3.9 KiB

  1. package main
  2. import (
  3. "database/sql"
  4. "flag"
  5. "fmt"
  6. "io/ioutil"
  7. "os"
  8. _ "github.com/lib/pq"
  9. "github.com/withmandala/go-log"
  10. "gopkg.in/yaml.v2"
  11. )
  12. var (
  13. configFile = flag.String("config", "", "Config file")
  14. singleTable = flag.String("table", "", "Table name")
  15. logger *log.Logger
  16. pg pgConnections
  17. )
  18. type config struct {
  19. Credentials map[string]string `yaml:"credentials"`
  20. }
  21. type pgConnections struct {
  22. Handlers map[string]*sql.DB
  23. }
  24. func loadConfig(configFile string) config {
  25. var conf config
  26. yamldata, err := ioutil.ReadFile(configFile)
  27. if err != nil {
  28. logger.Errorf("Error opening file %v", err)
  29. os.Exit(1)
  30. }
  31. marshErr := yaml.Unmarshal(yamldata, &conf)
  32. if marshErr != nil {
  33. logger.Fatalf("Config file malformed: %v", marshErr)
  34. os.Exit(1)
  35. }
  36. logger.Info("Loaded config:", conf)
  37. return conf
  38. }
  39. func getCount(db *sql.DB, tableName string) int {
  40. sqlStatement := fmt.Sprintf("SELECT count(*) FROM \"%v\";", tableName)
  41. var cnt int
  42. row := db.QueryRow(sqlStatement)
  43. switch err := row.Scan(&cnt); err {
  44. case sql.ErrNoRows:
  45. logger.Warn("No rows were returned!")
  46. case nil:
  47. return cnt
  48. default:
  49. logger.Error("Error", err, "DB", db)
  50. }
  51. return -1
  52. }
  53. func getSize(db *sql.DB, tableName string) string {
  54. sqlStatement := fmt.Sprintf("SELECT pg_size_pretty( pg_total_relation_size('%v') );", tableName)
  55. var size string
  56. row := db.QueryRow(sqlStatement)
  57. switch err := row.Scan(&size); err {
  58. case sql.ErrNoRows:
  59. logger.Warn("No rows were returned!")
  60. case nil:
  61. return size
  62. default:
  63. logger.Error("Error", err, "DB", db)
  64. }
  65. return "NAN"
  66. }
  67. func getTableNames(db *sql.DB) []string {
  68. rows, err := db.Query("SELECT tablename FROM pg_catalog.pg_tables WHERE schemaname != 'pg_catalog' AND schemaname != 'information_schema';")
  69. resultSet := []string{}
  70. if err != nil {
  71. logger.Error(err)
  72. os.Exit(1)
  73. }
  74. defer rows.Close()
  75. for rows.Next() {
  76. var tableName string
  77. err = rows.Scan(&tableName)
  78. if err != nil {
  79. logger.Error(err)
  80. os.Exit(1)
  81. }
  82. resultSet = append(resultSet, tableName)
  83. }
  84. err = rows.Err()
  85. if err != nil {
  86. logger.Error(err)
  87. os.Exit(1)
  88. }
  89. return resultSet
  90. }
  91. func inList(list []string, value string) bool {
  92. for _, val := range list {
  93. if val == value {
  94. return true
  95. }
  96. }
  97. return false
  98. }
  99. func init() {
  100. logger = log.New(os.Stderr).WithColor()
  101. // Config
  102. flag.Parse()
  103. if *configFile == "" {
  104. logger.Errorf("provide -config")
  105. os.Exit(1)
  106. }
  107. }
  108. func main() {
  109. conf := loadConfig(*configFile)
  110. // Connect to DBs
  111. pgConfigs := make(map[string]*sql.DB)
  112. for pgName, pgConnectString := range conf.Credentials {
  113. db, err := sql.Open("postgres", pgConnectString)
  114. if err != nil {
  115. logger.Error(err)
  116. os.Exit(1)
  117. }
  118. defer db.Close()
  119. err = db.Ping()
  120. if err != nil {
  121. logger.Error(err)
  122. os.Exit(1)
  123. }
  124. logger.Info(pgName, "successfully connected!")
  125. pgConfigs[pgName] = db
  126. }
  127. pg.Handlers = pgConfigs
  128. // Now get tables list
  129. tablesList := []string{}
  130. for _, handlerDB := range pg.Handlers {
  131. dbList := getTableNames(handlerDB)
  132. for _, tableName := range dbList {
  133. if !inList(tablesList, tableName) {
  134. tablesList = append(tablesList, tableName)
  135. }
  136. }
  137. }
  138. if *singleTable != "" {
  139. if !inList(tablesList, *singleTable) {
  140. // table not in database
  141. logger.Error("Table not found")
  142. os.Exit(1)
  143. } else {
  144. tablesList = []string{*singleTable}
  145. }
  146. }
  147. // Show records count
  148. for _, tableName := range tablesList {
  149. logger.Info("Table name:", tableName)
  150. counts := []int{}
  151. differenceFlag := false
  152. for pgName, handlerDB := range pg.Handlers {
  153. size := getSize(handlerDB, tableName)
  154. count := getCount(handlerDB, tableName)
  155. logger.Info("Postgres name:", pgName, "\tCount:", count, "\tSize:", size)
  156. if len(counts) > 0 && counts[len(counts)-1] != count {
  157. differenceFlag = true
  158. }
  159. counts = append(counts, count)
  160. }
  161. if differenceFlag {
  162. logger.Warn("------------ TABLES DIFFERENT ------------")
  163. } else {
  164. logger.Info("------------")
  165. }
  166. }
  167. }