/client/store.go

https://github.com/canonical/go-dqlite · Go · 239 lines · 167 code · 43 blank · 29 comment · 45 complexity · 3875eebdcd3ecf2f5aaaa41ef1a622d4 MD5 · raw file

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