github.com/rjgonzale/pop/v5@v5.1.3-dev/dialect_cockroach.go (about) 1 package pop 2 3 import ( 4 "database/sql" 5 "fmt" 6 "io" 7 "os" 8 "os/exec" 9 "path/filepath" 10 "strings" 11 "sync" 12 13 "github.com/gobuffalo/fizz" 14 "github.com/gobuffalo/fizz/translators" 15 "github.com/gobuffalo/pop/v5/columns" 16 "github.com/gobuffalo/pop/v5/internal/defaults" 17 "github.com/gobuffalo/pop/v5/logging" 18 "github.com/jmoiron/sqlx" 19 "github.com/pkg/errors" 20 ) 21 22 const nameCockroach = "cockroach" 23 const portCockroach = "26257" 24 25 const selectTablesQueryCockroach = "select table_name from information_schema.tables where table_schema = 'public' and table_type = 'BASE TABLE' and table_name <> ? and table_catalog = ?" 26 const selectTablesQueryCockroachV1 = "select table_name from information_schema.tables where table_name <> ? and table_schema = ?" 27 28 func init() { 29 AvailableDialects = append(AvailableDialects, nameCockroach) 30 dialectSynonyms["cockroachdb"] = nameCockroach 31 dialectSynonyms["crdb"] = nameCockroach 32 finalizer[nameCockroach] = finalizerCockroach 33 newConnection[nameCockroach] = newCockroach 34 } 35 36 var _ dialect = &cockroach{} 37 38 // ServerInfo holds informational data about connected database server. 39 type cockroachInfo struct { 40 VersionString string `db:"version"` 41 product string `db:"-"` 42 license string `db:"-"` 43 version string `db:"-"` 44 buildInfo string `db:"-"` 45 client string `db:"-"` 46 } 47 48 type cockroach struct { 49 commonDialect 50 translateCache map[string]string 51 mu sync.Mutex 52 info cockroachInfo 53 } 54 55 func (p *cockroach) Name() string { 56 return nameCockroach 57 } 58 59 func (p *cockroach) DefaultDriver() string { 60 return namePostgreSQL 61 } 62 63 func (p *cockroach) Details() *ConnectionDetails { 64 return p.ConnectionDetails 65 } 66 67 func (p *cockroach) Create(s store, model *Model, cols columns.Columns) error { 68 keyType := model.PrimaryKeyType() 69 switch keyType { 70 case "int", "int64": 71 cols.Remove("id") 72 id := struct { 73 ID int `db:"id"` 74 }{} 75 w := cols.Writeable() 76 var query string 77 if len(w.Cols) > 0 { 78 query = fmt.Sprintf("INSERT INTO %s (%s) VALUES (%s) returning id", p.Quote(model.TableName()), w.QuotedString(p), w.SymbolizedString()) 79 } else { 80 query = fmt.Sprintf("INSERT INTO %s DEFAULT VALUES returning id", p.Quote(model.TableName())) 81 } 82 log(logging.SQL, query) 83 stmt, err := s.PrepareNamed(query) 84 if err != nil { 85 return err 86 } 87 err = stmt.Get(&id, model.Value) 88 if err != nil { 89 if err := stmt.Close(); err != nil { 90 return errors.WithMessage(err, "failed to close statement") 91 } 92 return err 93 } 94 model.setID(id.ID) 95 return errors.WithMessage(stmt.Close(), "failed to close statement") 96 } 97 return genericCreate(s, model, cols, p) 98 } 99 100 func (p *cockroach) Update(s store, model *Model, cols columns.Columns) error { 101 return genericUpdate(s, model, cols, p) 102 } 103 104 func (p *cockroach) Destroy(s store, model *Model) error { 105 stmt := p.TranslateSQL(fmt.Sprintf("DELETE FROM %s WHERE %s", p.Quote(model.TableName()), model.whereID())) 106 _, err := genericExec(s, stmt, model.ID()) 107 return err 108 } 109 110 func (p *cockroach) SelectOne(s store, model *Model, query Query) error { 111 return genericSelectOne(s, model, query) 112 } 113 114 func (p *cockroach) SelectMany(s store, models *Model, query Query) error { 115 return genericSelectMany(s, models, query) 116 } 117 118 func (p *cockroach) CreateDB() error { 119 // createdb -h db -p 5432 -U cockroach enterprise_development 120 deets := p.ConnectionDetails 121 db, err := sql.Open(deets.Dialect, p.urlWithoutDb()) 122 if err != nil { 123 return errors.Wrapf(err, "error creating Cockroach database %s", deets.Database) 124 } 125 defer db.Close() 126 query := fmt.Sprintf("CREATE DATABASE %s", p.Quote(deets.Database)) 127 log(logging.SQL, query) 128 129 _, err = db.Exec(query) 130 if err != nil { 131 return errors.Wrapf(err, "error creating Cockroach database %s", deets.Database) 132 } 133 134 log(logging.Info, "created database %s", deets.Database) 135 return nil 136 } 137 138 func (p *cockroach) DropDB() error { 139 deets := p.ConnectionDetails 140 db, err := sql.Open(deets.Dialect, p.urlWithoutDb()) 141 if err != nil { 142 return errors.Wrapf(err, "error dropping Cockroach database %s", deets.Database) 143 } 144 defer db.Close() 145 query := fmt.Sprintf("DROP DATABASE %s CASCADE;", p.Quote(deets.Database)) 146 log(logging.SQL, query) 147 148 _, err = db.Exec(query) 149 if err != nil { 150 return errors.Wrapf(err, "error dropping Cockroach database %s", deets.Database) 151 } 152 153 log(logging.Info, "dropped database %s", deets.Database) 154 return nil 155 } 156 157 func (p *cockroach) URL() string { 158 c := p.ConnectionDetails 159 if c.URL != "" { 160 return c.URL 161 } 162 s := "postgres://%s:%s@%s:%s/%s?%s" 163 return fmt.Sprintf(s, c.User, c.Password, c.Host, c.Port, c.Database, c.OptionsString("")) 164 } 165 166 func (p *cockroach) urlWithoutDb() string { 167 c := p.ConnectionDetails 168 s := "postgres://%s:%s@%s:%s/?%s" 169 return fmt.Sprintf(s, c.User, c.Password, c.Host, c.Port, c.OptionsString("")) 170 } 171 172 func (p *cockroach) MigrationURL() string { 173 return p.URL() 174 } 175 176 func (p *cockroach) TranslateSQL(sql string) string { 177 defer p.mu.Unlock() 178 p.mu.Lock() 179 180 if csql, ok := p.translateCache[sql]; ok { 181 return csql 182 } 183 csql := sqlx.Rebind(sqlx.DOLLAR, sql) 184 185 p.translateCache[sql] = csql 186 return csql 187 } 188 189 func (p *cockroach) FizzTranslator() fizz.Translator { 190 return translators.NewCockroach(p.URL(), p.Details().Database) 191 } 192 193 func (p *cockroach) DumpSchema(w io.Writer) error { 194 cmd := exec.Command("cockroach", "dump", p.Details().Database, "--dump-mode=schema") 195 196 c := p.ConnectionDetails 197 if defaults.String(c.Options["sslmode"], "disable") == "disable" || strings.Contains(c.RawOptions, "sslmode=disable") { 198 cmd.Args = append(cmd.Args, "--insecure") 199 } 200 return genericDumpSchema(p.Details(), cmd, w) 201 } 202 203 func (p *cockroach) LoadSchema(r io.Reader) error { 204 return genericLoadSchema(p.ConnectionDetails, p.MigrationURL(), r) 205 } 206 207 func (p *cockroach) TruncateAll(tx *Connection) error { 208 type table struct { 209 TableName string `db:"table_name"` 210 } 211 212 tableQuery := p.tablesQuery() 213 214 var tables []table 215 if err := tx.RawQuery(tableQuery, tx.MigrationTableName(), tx.Dialect.Details().Database).All(&tables); err != nil { 216 return err 217 } 218 219 if len(tables) == 0 { 220 return nil 221 } 222 223 tableNames := make([]string, len(tables)) 224 for i, t := range tables { 225 tableNames[i] = t.TableName 226 //! work around for current limitation of DDL and DML at the same transaction. 227 // it should be fixed when cockroach support it or with other approach. 228 // https://www.cockroachlabs.com/docs/stable/known-limitations.html#schema-changes-within-transactions 229 if err := tx.RawQuery(fmt.Sprintf("delete from %s", p.Quote(t.TableName))).Exec(); err != nil { 230 return err 231 } 232 } 233 return nil 234 // TODO! 235 // return tx3.RawQuery(fmt.Sprintf("truncate %s cascade;", strings.Join(tableNames, ", "))).Exec() 236 } 237 238 func (p *cockroach) AfterOpen(c *Connection) error { 239 if err := c.RawQuery(`select version() AS "version"`).First(&p.info); err != nil { 240 return err 241 } 242 if s := strings.Split(p.info.VersionString, " "); len(s) > 3 { 243 p.info.product = s[0] 244 p.info.license = s[1] 245 p.info.version = s[2] 246 p.info.buildInfo = s[3] 247 } 248 log(logging.Debug, "server: %v %v %v", p.info.product, p.info.license, p.info.version) 249 250 return nil 251 } 252 253 func newCockroach(deets *ConnectionDetails) (dialect, error) { 254 deets.Dialect = "postgres" 255 d := &cockroach{ 256 commonDialect: commonDialect{ConnectionDetails: deets}, 257 translateCache: map[string]string{}, 258 mu: sync.Mutex{}, 259 } 260 d.info.client = deets.Options["application_name"] 261 return d, nil 262 } 263 264 func finalizerCockroach(cd *ConnectionDetails) { 265 appName := filepath.Base(os.Args[0]) 266 cd.Options["application_name"] = defaults.String(cd.Options["application_name"], appName) 267 cd.Port = defaults.String(cd.Port, portCockroach) 268 if cd.URL != "" { 269 cd.URL = "postgres://" + trimCockroachPrefix(cd.URL) 270 } 271 } 272 273 func trimCockroachPrefix(u string) string { 274 parts := strings.Split(u, "://") 275 if len(parts) != 2 { 276 return u 277 } 278 return parts[1] 279 } 280 281 func (p *cockroach) tablesQuery() string { 282 // See https://www.cockroachlabs.com/docs/stable/information-schema.html for more info about information schema changes 283 tableQuery := selectTablesQueryCockroach 284 if strings.HasPrefix(p.info.version, "v1.") { 285 tableQuery = selectTablesQueryCockroachV1 286 } 287 return tableQuery 288 }