github.com/getsynq/migrate/v4@v4.15.3-0.20220615182648-8e72daaa5ed9/database/clickhouse/clickhouse.go (about) 1 package clickhouse 2 3 import ( 4 "database/sql" 5 "fmt" 6 "io" 7 "io/ioutil" 8 "net/url" 9 "strconv" 10 "strings" 11 "time" 12 13 "go.uber.org/atomic" 14 15 "github.com/getsynq/migrate/v4" 16 "github.com/getsynq/migrate/v4/database" 17 "github.com/getsynq/migrate/v4/database/multistmt" 18 "github.com/hashicorp/go-multierror" 19 ) 20 21 var ( 22 multiStmtDelimiter = []byte(";") 23 24 DefaultMigrationsTable = "schema_migrations" 25 DefaultMigrationsTableEngine = "TinyLog" 26 DefaultMultiStatementMaxSize = 10 * 1 << 20 // 10 MB 27 28 ErrNilConfig = fmt.Errorf("no config") 29 ) 30 31 type Config struct { 32 DatabaseName string 33 ClusterName string 34 MigrationsTable string 35 MigrationsTableEngine string 36 MultiStatementEnabled bool 37 MultiStatementMaxSize int 38 } 39 40 func init() { 41 database.Register("clickhouse", &ClickHouse{}) 42 } 43 44 func WithInstance(conn *sql.DB, config *Config) (database.Driver, error) { 45 if config == nil { 46 return nil, ErrNilConfig 47 } 48 49 if err := conn.Ping(); err != nil { 50 return nil, err 51 } 52 53 ch := &ClickHouse{ 54 conn: conn, 55 config: config, 56 } 57 58 if err := ch.init(); err != nil { 59 return nil, err 60 } 61 62 return ch, nil 63 } 64 65 type ClickHouse struct { 66 conn *sql.DB 67 config *Config 68 isLocked atomic.Bool 69 } 70 71 func (ch *ClickHouse) Open(dsn string) (database.Driver, error) { 72 purl, err := url.Parse(dsn) 73 if err != nil { 74 return nil, err 75 } 76 q := migrate.FilterCustomQuery(purl) 77 q.Scheme = "tcp" 78 conn, err := sql.Open("clickhouse", q.String()) 79 if err != nil { 80 return nil, err 81 } 82 83 multiStatementMaxSize := DefaultMultiStatementMaxSize 84 if s := purl.Query().Get("x-multi-statement-max-size"); len(s) > 0 { 85 multiStatementMaxSize, err = strconv.Atoi(s) 86 if err != nil { 87 return nil, err 88 } 89 } 90 91 migrationsTableEngine := DefaultMigrationsTableEngine 92 if s := purl.Query().Get("x-migrations-table-engine"); len(s) > 0 { 93 migrationsTableEngine = s 94 } 95 96 ch = &ClickHouse{ 97 conn: conn, 98 config: &Config{ 99 MigrationsTable: purl.Query().Get("x-migrations-table"), 100 MigrationsTableEngine: migrationsTableEngine, 101 DatabaseName: purl.Query().Get("database"), 102 ClusterName: purl.Query().Get("x-cluster-name"), 103 MultiStatementEnabled: purl.Query().Get("x-multi-statement") == "true", 104 MultiStatementMaxSize: multiStatementMaxSize, 105 }, 106 } 107 108 if err := ch.init(); err != nil { 109 return nil, err 110 } 111 112 return ch, nil 113 } 114 115 func (ch *ClickHouse) init() error { 116 if len(ch.config.DatabaseName) == 0 { 117 if err := ch.conn.QueryRow("SELECT currentDatabase()").Scan(&ch.config.DatabaseName); err != nil { 118 return err 119 } 120 } 121 122 if len(ch.config.MigrationsTable) == 0 { 123 ch.config.MigrationsTable = DefaultMigrationsTable 124 } 125 126 if ch.config.MultiStatementMaxSize <= 0 { 127 ch.config.MultiStatementMaxSize = DefaultMultiStatementMaxSize 128 } 129 130 if len(ch.config.MigrationsTableEngine) == 0 { 131 ch.config.MigrationsTableEngine = DefaultMigrationsTableEngine 132 } 133 134 return ch.ensureVersionTable() 135 } 136 137 func (ch *ClickHouse) Run(r io.Reader) error { 138 if ch.config.MultiStatementEnabled { 139 var err error 140 if e := multistmt.Parse(r, multiStmtDelimiter, ch.config.MultiStatementMaxSize, func(m []byte) bool { 141 tq := strings.TrimSpace(string(m)) 142 if tq == "" { 143 return true 144 } 145 if _, e := ch.conn.Exec(string(m)); e != nil { 146 err = database.Error{OrigErr: e, Err: "migration failed", Query: m} 147 return false 148 } 149 return true 150 }); e != nil { 151 return e 152 } 153 return err 154 } 155 156 migration, err := ioutil.ReadAll(r) 157 if err != nil { 158 return err 159 } 160 161 if _, err := ch.conn.Exec(string(migration)); err != nil { 162 return database.Error{OrigErr: err, Err: "migration failed", Query: migration} 163 } 164 165 return nil 166 } 167 func (ch *ClickHouse) Version() (int, bool, error) { 168 var ( 169 version int 170 dirty uint8 171 query = "SELECT version, dirty FROM `" + ch.config.MigrationsTable + "` ORDER BY sequence DESC LIMIT 1" 172 ) 173 if err := ch.conn.QueryRow(query).Scan(&version, &dirty); err != nil { 174 if err == sql.ErrNoRows { 175 return database.NilVersion, false, nil 176 } 177 return 0, false, &database.Error{OrigErr: err, Query: []byte(query)} 178 } 179 return version, dirty == 1, nil 180 } 181 182 func (ch *ClickHouse) SetVersion(version int, dirty bool) error { 183 var ( 184 bool = func(v bool) uint8 { 185 if v { 186 return 1 187 } 188 return 0 189 } 190 tx, err = ch.conn.Begin() 191 ) 192 if err != nil { 193 return err 194 } 195 196 query := "INSERT INTO " + ch.config.MigrationsTable + " (version, dirty, sequence) VALUES (?, ?, ?)" 197 stmt, err := tx.Prepare(query) 198 if err != nil { 199 if rollbackErr := tx.Rollback(); rollbackErr != nil { 200 return fmt.Errorf("error during prepare statement %w and rollback %s", err, rollbackErr) 201 } 202 203 return err 204 } 205 206 if _, err := stmt.Exec(int64(version), bool(dirty), uint64(time.Now().UnixNano())); err != nil { 207 return &database.Error{OrigErr: err, Query: []byte(query)} 208 } 209 210 return tx.Commit() 211 } 212 213 // ensureVersionTable checks if versions table exists and, if not, creates it. 214 // Note that this function locks the database, which deviates from the usual 215 // convention of "caller locks" in the ClickHouse type. 216 func (ch *ClickHouse) ensureVersionTable() (err error) { 217 if err = ch.Lock(); err != nil { 218 return err 219 } 220 221 defer func() { 222 if e := ch.Unlock(); e != nil { 223 if err == nil { 224 err = e 225 } else { 226 err = multierror.Append(err, e) 227 } 228 } 229 }() 230 231 var ( 232 table string 233 query = "SHOW TABLES FROM " + ch.config.DatabaseName + " LIKE '" + ch.config.MigrationsTable + "'" 234 ) 235 // check if migration table exists 236 if err := ch.conn.QueryRow(query).Scan(&table); err != nil { 237 if err != sql.ErrNoRows { 238 return &database.Error{OrigErr: err, Query: []byte(query)} 239 } 240 } else { 241 return nil 242 } 243 244 // if not, create the empty migration table 245 if len(ch.config.ClusterName) > 0 { 246 query = fmt.Sprintf(` 247 CREATE TABLE %s ON CLUSTER %s ( 248 version Int64, 249 dirty UInt8, 250 sequence UInt64 251 ) Engine=%s`, ch.config.MigrationsTable, ch.config.ClusterName, ch.config.MigrationsTableEngine) 252 } else { 253 query = fmt.Sprintf(` 254 CREATE TABLE %s ( 255 version Int64, 256 dirty UInt8, 257 sequence UInt64 258 ) Engine=%s`, ch.config.MigrationsTable, ch.config.MigrationsTableEngine) 259 } 260 261 if strings.HasSuffix(ch.config.MigrationsTableEngine, "Tree") { 262 query = fmt.Sprintf(`%s ORDER BY sequence`, query) 263 } 264 265 if _, err := ch.conn.Exec(query); err != nil { 266 return &database.Error{OrigErr: err, Query: []byte(query)} 267 } 268 return nil 269 } 270 271 func (ch *ClickHouse) Drop() (err error) { 272 query := "SHOW TABLES FROM " + ch.config.DatabaseName 273 tables, err := ch.conn.Query(query) 274 275 if err != nil { 276 return &database.Error{OrigErr: err, Query: []byte(query)} 277 } 278 defer func() { 279 if errClose := tables.Close(); errClose != nil { 280 err = multierror.Append(err, errClose) 281 } 282 }() 283 284 for tables.Next() { 285 var table string 286 if err := tables.Scan(&table); err != nil { 287 return err 288 } 289 290 query = "DROP TABLE IF EXISTS " + ch.config.DatabaseName + "." + table 291 292 if _, err := ch.conn.Exec(query); err != nil { 293 return &database.Error{OrigErr: err, Query: []byte(query)} 294 } 295 } 296 if err := tables.Err(); err != nil { 297 return &database.Error{OrigErr: err, Query: []byte(query)} 298 } 299 300 return nil 301 } 302 303 func (ch *ClickHouse) Lock() error { 304 if !ch.isLocked.CAS(false, true) { 305 return database.ErrLocked 306 } 307 308 return nil 309 } 310 func (ch *ClickHouse) Unlock() error { 311 if !ch.isLocked.CAS(true, false) { 312 return database.ErrNotLocked 313 } 314 315 return nil 316 } 317 func (ch *ClickHouse) Close() error { return ch.conn.Close() }