/vendor/github.com/canonical/go-dqlite/client/store.go

https://github.com/rancher/k3s · Go · 231 lines · 162 code · 42 blank · 27 comment · 44 complexity · c236ad706b8980ec57a25b7fe41951ec MD5 · raw file

  1. package client
  2. import (
  3. "context"
  4. "database/sql"
  5. "fmt"
  6. "io/ioutil"
  7. "os"
  8. "sync"
  9. "github.com/ghodss/yaml"
  10. "github.com/pkg/errors"
  11. "github.com/canonical/go-dqlite/internal/protocol"
  12. _ "github.com/mattn/go-sqlite3" // Go SQLite bindings
  13. )
  14. // NodeStore is used by a dqlite client to get an initial list of candidate
  15. // dqlite nodes that it can dial in order to find a leader dqlite node to use.
  16. type NodeStore = protocol.NodeStore
  17. // NodeRole identifies the role of a node.
  18. type NodeRole = protocol.NodeRole
  19. // NodeInfo holds information about a single server.
  20. type NodeInfo = protocol.NodeInfo
  21. // InmemNodeStore keeps the list of target dqlite nodes in memory.
  22. type InmemNodeStore = protocol.InmemNodeStore
  23. // NewInmemNodeStore creates NodeStore which stores its data in-memory.
  24. var NewInmemNodeStore = protocol.NewInmemNodeStore
  25. // DatabaseNodeStore persists a list addresses of dqlite nodes in a SQL table.
  26. type DatabaseNodeStore struct {
  27. db *sql.DB // Database handle to use.
  28. schema string // Name of the schema holding the servers table.
  29. table string // Name of the servers table.
  30. column string // Column name in the servers table holding the server address.
  31. where string // Optional WHERE filter
  32. }
  33. // DefaultNodeStore creates a new NodeStore using the given filename to
  34. // open a SQLite database, with default names for the schema, table and column
  35. // parameters.
  36. //
  37. // It also creates the table if it doesn't exist yet.
  38. func DefaultNodeStore(filename string) (*DatabaseNodeStore, error) {
  39. // Open the database.
  40. db, err := sql.Open("sqlite3", filename)
  41. if err != nil {
  42. return nil, errors.Wrap(err, "failed to open database")
  43. }
  44. // Since we're setting SQLite single-thread mode, we need to have one
  45. // connection at most.
  46. db.SetMaxOpenConns(1)
  47. // Create the servers table if it does not exist yet.
  48. _, err = db.Exec("CREATE TABLE IF NOT EXISTS servers (address TEXT, UNIQUE(address))")
  49. if err != nil {
  50. return nil, errors.Wrap(err, "failed to create servers table")
  51. }
  52. store := NewNodeStore(db, "main", "servers", "address")
  53. return store, nil
  54. }
  55. // Option that can be used to tweak node store parameters.
  56. type NodeStoreOption func(*nodeStoreOptions)
  57. type nodeStoreOptions struct {
  58. Where string
  59. }
  60. // WithNodeStoreWhereClause configures the node store to append the given
  61. // hard-coded where clause to the SELECT query used to fetch nodes. Only the
  62. // clause itself must be given, without the "WHERE" prefix.
  63. func WithNodeStoreWhereClause(where string) NodeStoreOption {
  64. return func(options *nodeStoreOptions) {
  65. options.Where = where
  66. }
  67. }
  68. // NewNodeStore creates a new NodeStore.
  69. func NewNodeStore(db *sql.DB, schema, table, column string, options ...NodeStoreOption) *DatabaseNodeStore {
  70. o := &nodeStoreOptions{}
  71. for _, option := range options {
  72. option(o)
  73. }
  74. return &DatabaseNodeStore{
  75. db: db,
  76. schema: schema,
  77. table: table,
  78. column: column,
  79. where: o.Where,
  80. }
  81. }
  82. // Get the current servers.
  83. func (d *DatabaseNodeStore) Get(ctx context.Context) ([]NodeInfo, error) {
  84. tx, err := d.db.Begin()
  85. if err != nil {
  86. return nil, errors.Wrap(err, "failed to begin transaction")
  87. }
  88. defer tx.Rollback()
  89. query := fmt.Sprintf("SELECT %s FROM %s.%s", d.column, d.schema, d.table)
  90. if d.where != "" {
  91. query += " WHERE " + d.where
  92. }
  93. rows, err := tx.QueryContext(ctx, query)
  94. if err != nil {
  95. return nil, errors.Wrap(err, "failed to query servers table")
  96. }
  97. defer rows.Close()
  98. servers := make([]NodeInfo, 0)
  99. for rows.Next() {
  100. var address string
  101. err := rows.Scan(&address)
  102. if err != nil {
  103. return nil, errors.Wrap(err, "failed to fetch server address")
  104. }
  105. servers = append(servers, NodeInfo{ID: 1, Address: address})
  106. }
  107. if err := rows.Err(); err != nil {
  108. return nil, errors.Wrap(err, "result set failure")
  109. }
  110. return servers, nil
  111. }
  112. // Set the servers addresses.
  113. func (d *DatabaseNodeStore) Set(ctx context.Context, servers []NodeInfo) error {
  114. tx, err := d.db.Begin()
  115. if err != nil {
  116. return errors.Wrap(err, "failed to begin transaction")
  117. }
  118. query := fmt.Sprintf("DELETE FROM %s.%s", d.schema, d.table)
  119. if _, err := tx.ExecContext(ctx, query); err != nil {
  120. tx.Rollback()
  121. return errors.Wrap(err, "failed to delete existing servers rows")
  122. }
  123. query = fmt.Sprintf("INSERT INTO %s.%s(%s) VALUES (?)", d.schema, d.table, d.column)
  124. stmt, err := tx.PrepareContext(ctx, query)
  125. if err != nil {
  126. tx.Rollback()
  127. return errors.Wrap(err, "failed to prepare insert statement")
  128. }
  129. defer stmt.Close()
  130. for _, server := range servers {
  131. if _, err := stmt.ExecContext(ctx, server.Address); err != nil {
  132. tx.Rollback()
  133. return errors.Wrapf(err, "failed to insert server %s", server.Address)
  134. }
  135. }
  136. if err := tx.Commit(); err != nil {
  137. return errors.Wrap(err, "failed to commit transaction")
  138. }
  139. return nil
  140. }
  141. // Persists a list addresses of dqlite nodes in a YAML file.
  142. type YamlNodeStore struct {
  143. path string
  144. servers []NodeInfo
  145. mu sync.RWMutex
  146. }
  147. // NewYamlNodeStore creates a new YamlNodeStore backed by the given YAML file.
  148. func NewYamlNodeStore(path string) (*YamlNodeStore, error) {
  149. servers := []NodeInfo{}
  150. _, err := os.Stat(path)
  151. if err != nil {
  152. if !os.IsNotExist(err) {
  153. return nil, err
  154. }
  155. } else {
  156. data, err := ioutil.ReadFile(path)
  157. if err != nil {
  158. return nil, err
  159. }
  160. if err := yaml.Unmarshal(data, &servers); err != nil {
  161. return nil, err
  162. }
  163. }
  164. store := &YamlNodeStore{
  165. path: path,
  166. servers: servers,
  167. }
  168. return store, nil
  169. }
  170. // Get the current servers.
  171. func (s *YamlNodeStore) Get(ctx context.Context) ([]NodeInfo, error) {
  172. s.mu.RLock()
  173. defer s.mu.RUnlock()
  174. return s.servers, nil
  175. }
  176. // Set the servers addresses.
  177. func (s *YamlNodeStore) Set(ctx context.Context, servers []NodeInfo) error {
  178. s.mu.Lock()
  179. defer s.mu.Unlock()
  180. data, err := yaml.Marshal(servers)
  181. if err != nil {
  182. return err
  183. }
  184. if err := ioutil.WriteFile(s.path, data, 0600); err != nil {
  185. return err
  186. }
  187. s.servers = servers
  188. return nil
  189. }