Golang pq.CopyIn函数代码示例

// TestCopyOne verifies that only one COPY can run at once.
func TestCopyOne(t *testing.T) {
	defer leaktest.AfterTest(t)()

	t.Skip("https://github.com/lib/pq/issues/494")

	params, _ := createTestServerParams()
	s, db, _ := serverutils.StartServer(t, params)
	defer s.Stopper().Stop()

	if _, err := db.Exec(`
		CREATE DATABASE d;
		SET DATABASE = d;
		CREATE TABLE t (
			i INT PRIMARY KEY
		);
	`); err != nil {
		t.Fatal(err)
	}

	txn, err := db.Begin()
	if err != nil {
		t.Fatal(err)
	}

	if _, err := txn.Prepare(pq.CopyIn("t", "i")); err != nil {
		t.Fatal(err)
	}
	if _, err := txn.Prepare(pq.CopyIn("t", "i")); err == nil {
		t.Fatal("expected error")
	}
}

func (qs *QuadStore) copyFrom(tx *sql.Tx, in []graph.Delta) error {
	stmt, err := tx.Prepare(pq.CopyIn("quads", "subject", "predicate", "object", "label", "id", "ts", "subject_hash", "predicate_hash", "object_hash", "label_hash"))
	if err != nil {
		return err
	}
	for _, d := range in {
		_, err := stmt.Exec(
			d.Quad.Subject,
			d.Quad.Predicate,
			d.Quad.Object,
			d.Quad.Label,
			d.ID.Int(),
			d.Timestamp,
			hashOf(d.Quad.Subject),
			hashOf(d.Quad.Predicate),
			hashOf(d.Quad.Object),
			hashOf(d.Quad.Label),
		)
		if err != nil {
			glog.Errorf("couldn't prepare COPY statement: %v", err)
			return err
		}
	}
	_, err = stmt.Exec()
	if err != nil {
		return err
	}
	return stmt.Close()
}

func bulkDirect(db *sql.DB) error {
	tx, err := db.Begin()
	if err != nil {
		return err
	}
	stmt, err := tx.Prepare(pq.CopyIn("foo", "id", "val"))
	if err != nil {
		tx.Rollback()
		return err
	}
	// Build up the COPY FROM by repeated Exec calls
	for i := 0; i < *Count; i++ {
		val := *Start + i
		strVal := fmt.Sprintf("%07d", val)
		if _, err := stmt.Exec(val, strVal); err != nil {
			fmt.Printf("error while building up COPY: %s\n", err)
			tx.Rollback()
			return err
		}
	}
	// Flush the COPY FROM to postgres
	if _, err := stmt.Exec(); err != nil {
		fmt.Printf("error while copying data: %s\n", err)
		return err
	}
	if err := stmt.Close(); err != nil {
		fmt.Printf("error during stmt.Close(): %s\n", err)
		return err
	}
	if err := tx.Commit(); err != nil {
		fmt.Printf("could not commit transaction: %s\n", err)
		return err
	}
	return nil
}

func (db *Database) UpdateDivision(divisionId int, units []Unit) {
	tx, err := db.db.Begin()
	if err != nil {
		panic(err)
	}

	_, err = tx.Query("DELETE FROM unit WHERE unit.division = $1", divisionId)
	if err != nil {
		panic(err)
	}

	stmt, err := tx.Prepare(pq.CopyIn("unit", "division", "unit_type", "amount"))

	for _, unit := range units {
		_, err = stmt.Exec(divisionId, unit.TypeNum, unit.Amount)
		if err != nil {
			panic(err)
		}
	}

	_, err = stmt.Exec()
	if err != nil {
		panic(err)
	}

	err = stmt.Close()
	if err != nil {
		panic(err)
	}

	err = tx.Commit()
	if err != nil {
		panic(err)
	}
}

func (ops *DBOps) CopyDataIntoTable(t []Stock) {
	txn, err := ops.Db.Begin()
	if err != nil {
		log.Fatal(err)
	}
	defer txn.Rollback()

	stmt, err := txn.Prepare(pq.CopyIn("stock", "dealtime", "price", "volume", "amount", "climax"))
	if err != nil {
		log.Fatal(err)
	}

	for _, item := range t {
		_, err = stmt.Exec(item.Dealtime, item.Price, item.Volume, item.Amount, item.Climax)
		if err != nil {
			log.Fatal(err)
		}
	}

	_, err = stmt.Exec()
	if err != nil {
		log.Fatal(err)
	}

	err = stmt.Close()
	if err != nil {
		log.Fatal(err)
	}

	err = txn.Commit()
	if err != nil {
		log.Fatal(err)
	}
}

func insertShooters(db *sql.DB, m *Match) {
	txn, err := db.Begin()
	if err != nil {
		log.Fatal(err)
	}

	stmt, err := txn.Prepare(pq.CopyIn("users", "name", "countryCode", "mos", "flags"))
	if err != nil {
		log.Fatal(err)
	}

	for _, results := range m.Results {
		for _, result := range results {
			_, err = stmt.Exec(result.Name, result.CountryCode, result.MosNumber, result.Flags)
			if err != nil {
				log.Fatal(err)
			}
		}
	}

	_, err = stmt.Exec()
	if err != nil {
		log.Fatal(err)
	}

	err = stmt.Close()
	if err != nil {
		log.Fatal(err)
	}

	err = txn.Commit()
	if err != nil {
		log.Fatal(err)
	}
}

func createSubscriptions(ids []int) error {
	info := getDBSettings()
	db, errCon := sql.Open("postgres", fmt.Sprintf("host=%v user=%v password=%v dbname=%v sslmode=require", info.Host, info.Username, info.Password, info.Database))
	defer db.Close()
	if errCon != nil {
		log.Fatal(errCon)
	}
	txn, errT := db.Begin()
	if errT != nil {
		log.Println(errT)
		return errT
	}
	stmt, errPrep := txn.Prepare(pq.CopyIn("subscription", "topicid", "userid"))
	if errPrep != nil {
		log.Fatal(errPrep)
	}

	log.Println("Start For...")
	var wg sync.WaitGroup
	for _, id := range ids {

		userID := id
		wg.Add(1)
		go func(id int) {
			defer wg.Done()
			mr.Seed(time.Now().UnixNano())
			numSubs := mr.Intn(5) + 1
			for i := 0; i < numSubs; i++ {
				topicID := i + 1
				_, errA := stmt.Exec(topicID, userID)
				if errA != nil {
					log.Fatal(errA)
				}
			}
		}(userID)

	}
	wg.Wait()
	log.Println("End For")
	log.Println("Start Exec")
	_, errEX := stmt.Exec()
	if errEX != nil {
		log.Fatal(errEX)
	}
	log.Println("End Exec")

	errClose := stmt.Close()
	if errClose != nil {
		log.Fatal(errClose)
	}
	log.Println("Start Commit")
	errCommit := txn.Commit()
	if errCommit != nil {
		log.Fatal(errCommit)
	}
	log.Println("End Commit")
	return nil

}

func (dfr *DataFileReader) executeCopyLine(line string, c *spiffy.DbConnection, tx *sql.Tx) (*sql.Stmt, error) {
	pieces := dfr.extractCopyLine(line)
	if len(pieces) < 3 {
		return nil, exception.New("Invalid `COPY ...` line, cannot continue.")
	}
	tableName := pieces[1]
	columnCSV := pieces[2]
	columns := strings.Split(columnCSV, ", ")
	return tx.Prepare(pq.CopyIn(tableName, columns...))
}

func createLang(ids []int) error {
	info := getDBSettings()
	db, errCon := sql.Open("postgres", fmt.Sprintf("host=%v user=%v password=%v dbname=%v sslmode=require", info.Host, info.Username, info.Password, info.Database))
	defer db.Close()
	if errCon != nil {
		log.Fatal(errCon)
	}
	txn, errT := db.Begin()
	if errT != nil {
		log.Println(errT)
		return errT
	}
	stmt, errPrep := txn.Prepare(pq.CopyIn("userinfo", "userid", "langcode"))
	if errPrep != nil {
		log.Fatal(errPrep)
	}

	log.Println("Start For...")
	var wg sync.WaitGroup
	for _, id := range ids {

		userID := id
		wg.Add(1)
		go func(id int) {
			defer wg.Done()
			_, errA := stmt.Exec(userID, "en")
			if errA != nil {
				log.Fatal(errA)
			}
		}(userID)

	}
	wg.Wait()
	log.Println("End For")
	log.Println("Start Exec")
	_, errEX := stmt.Exec()
	if errEX != nil {
		log.Fatal(errEX)
	}
	log.Println("End Exec")

	errClose := stmt.Close()
	if errClose != nil {
		log.Fatal(errClose)
	}
	log.Println("Start Commit")
	errCommit := txn.Commit()
	if errCommit != nil {
		log.Fatal(errCommit)
	}
	log.Println("End Commit")
	return nil
}

func commitData(db *sql.DB, data map[string]string, columns ...string) {

	var dataitems = []string{}
	columns = append(columns, "symbol")
	for k, v := range columns {
		if k == 0 && len(data[v]) == 0 {
			break
		}
		dataitems = append(dataitems, data[v])
	}

	if len(dataitems) > 0 {

		fmt.Printf("%v{%d}\n", dataitems, len(dataitems))

		txn, err := db.Begin()
		if err != nil {
			log.Fatal(err)
		}

		stmt, err := txn.Prepare(pq.CopyIn("balancesheet", columns...))
		if err != nil {
			log.Fatal(err)
		}

		new := make([]interface{}, len(dataitems))
		for i, v := range dataitems {
			new[i] = v
		}

		_, err = stmt.Exec(new...)
		if err != nil {
			log.Fatal(err)
		}

		_, err = stmt.Exec()
		if err != nil {
			log.Fatal(err)
		}

		err = stmt.Close()
		if err != nil {
			log.Fatal(err)
		}

		err = txn.Commit()
		if err != nil {
			log.Fatal(err)
		}
	}

}

// insert operates on a list of dumpFormat and inserts them to the provided Postgres
// database.
func (data dataset) insert(db *sql.DB) error {
	transaction, err := db.Begin()
	if err != nil {
		return fmt.Errorf("Error starting PG txn => %s", err.Error())
	}

	// PG's COPY FROM used for fast mass insertions. Syntax is table followed by columns.
	// http://godoc.org/github.com/lib/pq#hdr-Bulk_imports
	stmt, err := transaction.Prepare(pq.CopyIn(
		"density_data", // table
		"dump_time",    // columns.....
		"group_id",
		"group_name",
		"parent_id",
		"parent_name",
		"client_count",
	))
	if err != nil {
		return fmt.Errorf("Error prepping PG txn => %s", err.Error())
	}
	defer stmt.Close()

	// Add all data from the set
	for _, d := range data {
		_, err = stmt.Exec(
			d.DumpTime,
			d.GroupID,
			d.GroupName,
			d.ParentID,
			d.ParentName,
			d.ClientCount,
		)
		if err != nil {
			return fmt.Errorf("Failed to add to bulk insert => %s", err.Error())
		}
	}

	// execute the transaction
	if _, err = stmt.Exec(); err != nil {
		return fmt.Errorf("Failed to execute bulk insert => %s", err.Error())
	}

	// commit the transaction if there's been no errors
	if err = transaction.Commit(); err != nil {
		log.Printf("ERROR: Failed to commit txn => %s", err.Error())
		if err = transaction.Rollback(); err != nil {
			log.Printf("ERROR: Failed to rollback txn => %s", err.Error())
		}
	}
	return nil
}

func TestCopyError(t *testing.T) {
	defer leaktest.AfterTest(t)()

	params, _ := createTestServerParams()
	s, db, _ := serverutils.StartServer(t, params)
	defer s.Stopper().Stop()

	if _, err := db.Exec(`
		CREATE DATABASE d;
		SET DATABASE = d;
		CREATE TABLE t (
			i INT PRIMARY KEY
		);
	`); err != nil {
		t.Fatal(err)
	}

	txn, err := db.Begin()
	if err != nil {
		t.Fatal(err)
	}

	stmt, err := txn.Prepare(pq.CopyIn("t", "i"))
	if err != nil {
		t.Fatal(err)
	}

	// Insert conflicting primary keys.
	for i := 0; i < 2; i++ {
		_, err = stmt.Exec(1)
		if err != nil {
			t.Fatal(err)
		}
	}

	err = stmt.Close()
	if err == nil {
		t.Fatal("expected error")
	}

	// Make sure we can query after an error.
	var i int
	if err := db.QueryRow("SELECT 1").Scan(&i); err != nil {
		t.Fatal(err)
	} else if i != 1 {
		t.Fatalf("expected 1, got %d", i)
	}
	if err := txn.Rollback(); err != nil {
		t.Fatal(err)
	}
}

func InsertMultiDataValues(db *sql.DB, table_name string, table_columns []string, data [][]interface{}) error {
	// Transaction Begins and must end with a commit or rollback
	transaction, err := db.Begin()
	if err != nil {
		transaction.Rollback()
		return err
	}

	// Preparing statement with the table name and columns passed
	statement, err := transaction.Prepare(pq.CopyIn(table_name, table_columns...))
	if err != nil {
		statement.Close()
		transaction.Rollback()
		return err
	}

	// Looping though all the data rows passed
	for _, data_row := range data {
		// Inserting Single Data row into the statement
		_, err = statement.Exec(data_row...)
		if err != nil {
			statement.Close()
			transaction.Rollback()
			return err
		}
	}

	/*
		_, err = stmt.Exec()
		if err != nil {
			return err
		}
	*/

	// Closing the connection of the statement
	err = statement.Close()
	if err != nil {
		statement.Close()
		transaction.Rollback()
		return err
	}

	// Commiting and closing the transaction saving changes we have made in the database
	err = transaction.Commit()
	if err != nil {
		transaction.Rollback()
		return err
	}
	return nil
}

// TestCopyTransaction verifies that COPY data can be used after it is done
// within a transaction.
func TestCopyTransaction(t *testing.T) {
	defer leaktest.AfterTest(t)()

	params, _ := createTestServerParams()
	s, db, _ := serverutils.StartServer(t, params)
	defer s.Stopper().Stop()

	if _, err := db.Exec(`
		CREATE DATABASE d;
		SET DATABASE = d;
		CREATE TABLE t (
			i INT PRIMARY KEY
		);
	`); err != nil {
		t.Fatal(err)
	}

	txn, err := db.Begin()
	if err != nil {
		t.Fatal(err)
	}

	stmt, err := txn.Prepare(pq.CopyIn("t", "i"))
	if err != nil {
		t.Fatal(err)
	}

	const val = 2

	_, err = stmt.Exec(val)
	if err != nil {
		t.Fatal(err)
	}

	if err = stmt.Close(); err != nil {
		t.Fatal(err)
	}

	var i int
	if err := txn.QueryRow("SELECT i FROM d.t").Scan(&i); err != nil {
		t.Fatal(err)
	} else if i != val {
		t.Fatalf("expected 1, got %d", i)
	}
	if err := txn.Commit(); err != nil {
		t.Fatal(err)
	}
}

func NewInserter(
	tx *sql.Tx,
	driver string,
	table string,
	columns ...string,
) (ins *Inserter, err error) {
	defer Safe(&err)

	ins = &Inserter{tx, table, columns, driver, nil}
	if driver == "postgres" {
		ins.stmt = Prepare(tx, pq.CopyIn(table, columns...))
	} else {
		ins.stmt = ins.preparedInsert()
	}
	return
}