Golang csql.Safe函數代碼示例

// Tables returns the names of all tables in the database sorted
// alphabetically in ascending order.
func (db *DB) Tables() (tables []string, err error) {
	defer csql.Safe(&err)

	var q string
	switch db.Driver {
	case "postgres":
		q = `
			SELECT tablename FROM pg_tables
			WHERE schemaname NOT IN ('pg_catalog', 'information_schema')
			ORDER BY tablename ASC
		`
	case "sqlite3":
		q = `
			SELECT tbl_name FROM sqlite_master
			WHERE type = 'table'
			ORDER BY tbl_name ASC
		`
	default:
		return nil, ef("Unrecognized database driver: %s", db.Driver)
	}
	rows := csql.Query(db, q)
	csql.ForRow(rows, func(rs csql.RowScanner) {
		var table string
		csql.Scan(rs, &table)
		if table != "migration_version" {
			tables = append(tables, table)
		}
	})
	return
}

// Set associates the plain text password given with the user that is uniquely
// identified by id. The password is hashed with bcrypt. If there is a problem
// with hashing or with storing the password, an error is returned.
//
// This may be called on a new user.
func (s *Store) Set(id, password string) (cerr error) {
	defer csql.Safe(&cerr)

	hash, err := bcrypt.GenerateFromPassword(
		[]byte(password), bcrypt.DefaultCost)
	if err != nil {
		return err
	}

	// This lock can be avoided if we use some sort of upsert.
	// It's possible with Postgres, but this is just way easier.
	locker.Lock(id)
	defer locker.Unlock(id)

	n := csql.Count(s, `
		SELECT COUNT(*) FROM `+SqlTableName+` WHERE id = $1
		`, id)
	if n == 0 {
		csql.Exec(s, `
			INSERT INTO `+SqlTableName+` (id, hash) VALUES ($1, $2)
			`, id, hash)
	} else {
		csql.Exec(s, `
			UPDATE `+SqlTableName+` SET id = $1, hash = $2 WHERE id = $1
			`, id, hash)
	}
	return nil
}

func doIndices(
	db *DB,
	getSql func(index, *DB) string,
	tables ...string,
) (err error) {
	defer csql.Safe(&err)

	trgmEnabled := db.IsFuzzyEnabled()
	var q string
	var ok bool
	for _, idx := range indices {
		if idx.isFulltext() && !trgmEnabled {
			// Only show the error message if we're on PostgreSQL.
			if db.Driver == "postgres" {
				log.Printf("Skipping fulltext index '%s' since "+
					"the pg_trgm extension is not enabled.", idx.sqlName())
			}
			continue
		}
		if len(tables) == 0 || fun.In(idx.table, tables) {
			q += getSql(idx, db) + "; "
			ok = true
		}
	}
	if ok {
		csql.Exec(db, q)
	}
	return
}

// Get retrieves the current password hash for the user given.
func (s *Store) Get(id string) (hash []byte, err error) {
	defer csql.Safe(&err)

	r := s.QueryRow(`
		SELECT hash FROM `+SqlTableName+` WHERE id = $1
		`, id)
	csql.Scan(r, &hash)
	return
}

func listMovieLinks(db *imdb.DB, atoms *atomizer, r io.ReadCloser) (err error) {
	defer csql.Safe(&err)
	table := startSimpleLoad(db, "link", "atom_id",
		"link_type", "link_atom_id", "entity")
	defer table.done()

	parseMovieLink := func(
		atoms *atomizer,
		text []byte,
		linkType *string,
		linkAtom *imdb.Atom,
		linkEntity *imdb.EntityKind,
	) bool {
		attrName, data, ok := parseNamedAttr(text)
		if !ok {
			logf("Could not parse named attribute '%s'. Skipping.", text)
			return false
		}
		id, ok := atoms.atomOnlyIfExist(data)
		if !ok {
			warnf("Could not find id for '%s'. Skipping.", data)
			return false
		}
		ent, ok := parseMediaEntity(data)
		if !ok {
			logf("Could not find entity type for '%s'. Skipping.", data)
			return false
		}
		*linkType = unicode(attrName)
		*linkAtom = id
		*linkEntity = ent.Type()
		return true
	}

	listAttrRowIds(r, table.atoms, func(id imdb.Atom, line, ent, row []byte) {
		var (
			linkType   string
			linkAtom   imdb.Atom
			linkEntity imdb.EntityKind
		)

		fields := splitListLine(row)
		if len(fields) == 0 {
			return
		}
		if bytes.Contains(fields[0], []byte("(VG)")) {
			return
		}
		ok := parseMovieLink(table.atoms, fields[0],
			&linkType, &linkAtom, &linkEntity)
		if !ok {
			return
		}
		table.add(line, id, linkType, linkAtom, linkEntity.String())
	})
	return
}

func listActors(db *imdb.DB, ractor, ractress io.ReadCloser) (err error) {
	defer csql.Safe(&err)

	logf("Reading actors list...")

	// PostgreSQL wants different transactions for each inserter.
	// SQLite can't handle them. The wrapper type here ensures that
	// PostgreSQL gets multiple transactions while SQLite only gets one.
	tx, err := db.Begin()
	csql.Panic(err)

	txactor := wrapTx(db, tx)
	txcredit := txactor.another()
	txname := txactor.another()
	txatom := txactor.another()

	// Drop data from the actor and credit tables. They will be rebuilt below.
	// The key here is to leave the atom and name tables alone. Invariably,
	// they will contain stale data. But the only side effect, I think, is
	// taking up space.
	// (Stale data can be removed with 'goim clean'.)
	csql.Truncate(txactor, db.Driver, "actor")
	csql.Truncate(txcredit.Tx, db.Driver, "credit")

	actIns, err := csql.NewInserter(txactor.Tx, db.Driver, "actor",
		"atom_id", "sequence")
	csql.Panic(err)
	credIns, err := csql.NewInserter(txcredit.Tx, db.Driver, "credit",
		"actor_atom_id", "media_atom_id", "character", "position", "attrs")
	csql.Panic(err)
	nameIns, err := csql.NewInserter(txname.Tx, db.Driver, "name",
		"atom_id", "name")
	csql.Panic(err)
	atoms, err := newAtomizer(db, txatom.Tx)
	csql.Panic(err)

	// Unfortunately, it looks like credits for an actor can appear in
	// multiple locations. (Or there are different actors that erroneously
	// have the same name.)
	added := make(map[imdb.Atom]struct{}, 3000000)
	n1, nc1 := listActs(db, ractress, atoms, added, actIns, credIns, nameIns)
	n2, nc2 := listActs(db, ractor, atoms, added, actIns, credIns, nameIns)

	csql.Panic(actIns.Exec())
	csql.Panic(credIns.Exec())
	csql.Panic(nameIns.Exec())
	csql.Panic(atoms.Close())

	csql.Panic(txactor.Commit())
	csql.Panic(txcredit.Commit())
	csql.Panic(txname.Commit())
	csql.Panic(txatom.Commit())

	logf("Done. Added %d actors/actresses and %d credits.", n1+n2, nc1+nc2)
	return
}

func listTaglines(db *imdb.DB, atoms *atomizer, r io.ReadCloser) (err error) {
	defer csql.Safe(&err)
	table := startSimpleLoad(db, "tagline", "atom_id", "tag")
	defer table.done()

	do := func(id imdb.Atom, item []byte) {
		table.add(item, id, unicode(item))
	}
	listPrefixItems(r, table.atoms, []byte{'#'}, []byte{'\t'}, do)
	return
}

func listRunningTimes(
	db *imdb.DB,
	atoms *atomizer,
	r io.ReadCloser,
) (err error) {
	defer csql.Safe(&err)
	table := startSimpleLoad(db, "running_time",
		"atom_id", "country", "minutes", "attrs")
	defer table.done()

	parseRunningTime := func(text []byte, country *string, minutes *int) bool {
		sep := bytes.IndexByte(text, ':')
		var runtime []byte
		if sep > -1 {
			*country = unicode(bytes.TrimSpace(text[:sep]))
			runtime = bytes.TrimSpace(text[sep+1:])
		} else {
			*country = ""
			runtime = bytes.TrimSpace(text)
		}

		var err error
		*minutes, err = strconv.Atoi(unicode(runtime))
		if err != nil {
			// There are a lot of these.
			// From the looks of it, IMDb's web site just ignores them.
			// It's almost like it's freeform text... Yikes.
			return false
		}
		return true
	}

	listAttrRowIds(r, table.atoms, func(id imdb.Atom, line, ent, row []byte) {
		var (
			country string
			minutes int
			attrs   []byte
		)

		rowFields := splitListLine(row)
		if len(rowFields) == 0 {
			return // herp derp...
		}
		if !parseRunningTime(rowFields[0], &country, &minutes) {
			return
		}
		if len(rowFields) > 1 {
			attrs = rowFields[1]
		}
		table.add(line, id, country, minutes, unicode(attrs))
	})
	return
}

func listGenres(db *imdb.DB, atoms *atomizer, r io.ReadCloser) (err error) {
	defer csql.Safe(&err)
	table := startSimpleLoad(db, "genre", "atom_id", "name")
	defer table.done()

	listAttrRowIds(r, table.atoms, func(id imdb.Atom, line, ent, row []byte) {
		fields := splitListLine(row)
		if len(fields) == 0 {
			return
		}
		table.add(line, id, strings.ToLower(unicode(fields[0])))
	})
	return
}

func listAlternateVersions(
	db *imdb.DB,
	atoms *atomizer,
	r io.ReadCloser,
) (err error) {
	defer csql.Safe(&err)
	table := startSimpleLoad(db, "alternate_version", "atom_id", "about")
	defer table.done()

	do := func(id imdb.Atom, item []byte) {
		table.add(item, id, unicode(item))
	}
	listPrefixItems(r, table.atoms, []byte{'#'}, []byte{'-'}, do)
	return
}

// newAtomizer returns an atomizer that can be used to access or create new
// atom identifiers. Note that if tx is nil, then the atomizer returned is
// read-only (attempting to write will cause a panic).
//
// A read-only atomizer may be accessed from multiple goroutines
// simultaneously, but a read/write atomizer may NOT.
//
// If a read/write atomizer is created, then the caller is responsible for
// closing the transaction (which should be done immediately after a call to
// atomizer.Close).
//
// Note that this function loads the entire set of atoms from the database
// into memory, so it is costly.
func newAtomizer(db *imdb.DB, tx *sql.Tx) (az *atomizer, err error) {
	defer csql.Safe(&err)

	az = &atomizer{db, make(atomMap, 1000000), 0, nil}
	if tx != nil {
		var err error
		az.ins, err = csql.NewInserter(
			tx, db.Driver, "atom", "id", "hash")
		csql.Panic(err)
	}

	rs := csql.Query(db, "SELECT id, hash FROM atom ORDER BY id ASC")
	csql.ForRow(rs, az.readRow)
	az.nextId++
	return
}

func listPlots(db *imdb.DB, atoms *atomizer, r io.ReadCloser) (err error) {
	defer csql.Safe(&err)
	table := startSimpleLoad(db, "plot", "atom_id", "entry", "by")
	defer table.done()

	var curAtom imdb.Atom
	var curPlot []byte
	var curBy []byte
	var ok bool
	add := func(line []byte) {
		if curAtom > 0 && len(curPlot) > 0 {
			plot := unicode(bytes.TrimSpace(curPlot))
			by := unicode(bytes.TrimSpace(curBy))
			table.add(line, curAtom, plot, by)
		}
		curPlot, curBy = nil, nil
	}
	listLines(r, func(line []byte) {
		if bytes.HasPrefix(line, []byte("MV:")) {
			if len(curPlot) > 0 {
				add(line)
			}
			entity := bytes.TrimSpace(line[3:])
			if curAtom, ok = table.atoms.atomOnlyIfExist(entity); !ok {
				warnf("Could not find id for '%s'. Skipping.", entity)
				curAtom, curPlot, curBy = 0, nil, nil
			}
			return
		}
		if len(line) == 0 {
			return
		}
		if bytes.HasPrefix(line, []byte("PL:")) {
			curPlot = append(curPlot, bytes.TrimSpace(line[3:])...)
			curPlot = append(curPlot, ' ')
			return
		}
		if bytes.HasPrefix(line, []byte("BY:")) {
			curBy = line[3:]
			add(line)
			return
		}
	})
	add([]byte("UNKNOWN (last line?)"))
	return
}

func listRatings(db *imdb.DB, atoms *atomizer, r io.ReadCloser) (err error) {
	defer csql.Safe(&err)
	table := startSimpleLoad(db, "rating", "atom_id", "votes", "rank")
	defer table.done()

	done := false
	listLines(r, func(line []byte) {
		var (
			id    imdb.Atom
			ok    bool
			votes int
			rank  float64
		)
		if done {
			return
		}

		fields := bytes.Fields(line)
		if bytes.HasPrefix(line, []byte("REPORT FORMAT")) {
			done = true
			return
		}
		if len(fields) < 4 {
			return
		}
		if bytes.Equal(fields[0], []byte("New")) {
			return
		}

		entity := bytes.Join(fields[3:], []byte{' '})
		if id, ok = table.atoms.atomOnlyIfExist(entity); !ok {
			warnf("Could not find id for '%s'. Skipping.", entity)
			return
		}
		if err := parseInt(fields[1], &votes); err != nil {
			logf("Could not parse integer '%s' in: '%s'", fields[1], line)
			return
		}
		if err := parseFloat(fields[2], &rank); err != nil {
			logf("Could not parse float '%s' in: '%s'", fields[2], line)
			return
		}
		table.add(line, id, votes, int(10*rank))
	})
	return
}

// attrs uses reflection to automatically construct a list of simple attribute
// rows from the database based on information in the attribute's struct.
// This includes building the SELECT query and the slice itself.
//
// zero MUST be a pointer to a simple struct. A simple struct MUST ONLY contain
// fields that can be encoded/decoded as declared by the 'database/sql'
// package. Column names are the lowercase version of their struct field name
// unless the 'imdb_name' struct tag is set, in which case, that name is used.
//
// extra is passed to the end of the query executed. Useful for specifying
// ORDER BY or LIMIT clauses.
func attrs(
	zero interface{},
	db csql.Queryer,
	e Entity,
	tableName string,
	idColumn string,
	extra string,
) (v interface{}, err error) {
	defer csql.Safe(&err)

	rz := reflect.ValueOf(zero).Elem()
	tz := rz.Type()
	nfields := tz.NumField()
	columns := make([]string, nfields)
	for i := 0; i < nfields; i++ {
		f := tz.Field(i)
		column := f.Tag.Get("imdb_name")
		if len(column) == 0 {
			column = strings.ToLower(f.Name)
		}
		columns[i] = column
	}
	tattrs := reflect.SliceOf(tz)
	vattrs := reflect.MakeSlice(tattrs, 0, 10)
	v = vattrs.Interface()

	q := sf("SELECT %s FROM %s WHERE %s = $1 %s",
		strings.Join(columns, ", "), tableName, idColumn, extra)
	rs := csql.Query(db, q, e.Ident())
	csql.ForRow(rs, func(s csql.RowScanner) {
		loadCols := make([]interface{}, nfields)
		for i := 0; i < nfields; i++ {
			loadCols[i] = reflect.New(tz.Field(i).Type).Interface()
		}
		csql.Scan(s, loadCols...)

		row := reflect.New(tz).Elem()
		for i := 0; i < nfields; i++ {
			row.Field(i).Set(reflect.ValueOf(loadCols[i]).Elem())
		}
		vattrs = reflect.Append(vattrs, row)
	})
	v = vattrs.Interface() // not sure if this is necessary.
	return
}

func listLanguages(db *imdb.DB, atoms *atomizer, r io.ReadCloser) (err error) {
	defer csql.Safe(&err)
	table := startSimpleLoad(db, "language", "atom_id", "name", "attrs")
	defer table.done()

	listAttrRowIds(r, table.atoms, func(id imdb.Atom, line, ent, row []byte) {
		var attrs []byte
		fields := splitListLine(row)
		if len(fields) == 0 {
			return
		}
		if len(fields) > 1 {
			attrs = fields[1]
		}
		table.add(line, id, unicode(fields[0]), unicode(attrs))
	})
	return
}