Golang pgx.Connect函数代码示例

func TestConnectWithTLSFallback(t *testing.T) {
	t.Parallel()

	if tlsConnConfig == nil {
		return
	}

	connConfig := *tlsConnConfig
	connConfig.TLSConfig = &tls.Config{ServerName: "bogus.local"} // bogus ServerName should ensure certificate validation failure

	conn, err := pgx.Connect(connConfig)
	if err == nil {
		t.Fatal("Expected failed connection, but succeeded")
	}

	connConfig.UseFallbackTLS = true
	connConfig.FallbackTLSConfig = &tls.Config{InsecureSkipVerify: true}

	conn, err = pgx.Connect(connConfig)
	if err != nil {
		t.Fatal("Unable to establish connection: " + err.Error())
	}

	err = conn.Close()
	if err != nil {
		t.Fatal("Unable to close connection")
	}
}

func TestConnectCustomDialer(t *testing.T) {
	t.Parallel()

	if customDialerConnConfig == nil {
		return
	}

	dialled := false
	conf := *customDialerConnConfig
	conf.Dial = func(network, address string) (net.Conn, error) {
		dialled = true
		return net.Dial(network, address)
	}

	conn, err := pgx.Connect(conf)
	if err != nil {
		t.Fatalf("Unable to establish connection: %s", err)
	}
	if !dialled {
		t.Fatal("Connect did not use custom dialer")
	}

	err = conn.Close()
	if err != nil {
		t.Fatal("Unable to close connection")
	}
}

func TestFatalRxError(t *testing.T) {
	t.Parallel()

	conn := mustConnect(t, *defaultConnConfig)
	defer closeConn(t, conn)

	var wg sync.WaitGroup
	wg.Add(1)
	go func() {
		defer wg.Done()
		var n int32
		var s string
		err := conn.QueryRow("select 1::int4, pg_sleep(10)::varchar").Scan(&n, &s)
		if err, ok := err.(pgx.PgError); !ok || err.Severity != "FATAL" {
			t.Fatalf("Expected QueryRow Scan to return fatal PgError, but instead received %v", err)
		}
	}()

	otherConn, err := pgx.Connect(*defaultConnConfig)
	if err != nil {
		t.Fatalf("Unable to establish connection: %v", err)
	}
	defer otherConn.Close()

	if _, err := otherConn.Exec("select pg_terminate_backend($1)", conn.Pid); err != nil {
		t.Fatalf("Unable to kill backend PostgreSQL process: %v", err)
	}

	wg.Wait()

	if conn.IsAlive() {
		t.Fatal("Connection should not be live but was")
	}
}

func mustConnect(t testing.TB, config pgx.ConnConfig) *pgx.Conn {
	conn, err := pgx.Connect(config)
	if err != nil {
		t.Fatalf("Unable to establish connection: %v", err)
	}
	return conn
}

func TestFatalTxError(t *testing.T) {
	t.Parallel()

	conn := mustConnect(t, *defaultConnConfig)
	defer closeConn(t, conn)

	otherConn, err := pgx.Connect(*defaultConnConfig)
	if err != nil {
		t.Fatalf("Unable to establish connection: %v", err)
	}
	defer otherConn.Close()

	_, err = otherConn.Exec("select pg_terminate_backend($1)", conn.Pid)
	if err != nil {
		t.Fatalf("Unable to kill backend PostgreSQL process: %v", err)
	}

	_, err = conn.Query("select 1")
	if err == nil {
		t.Fatal("Expected error but none occurred")
	}

	if conn.IsAlive() {
		t.Fatal("Connection should not be live but was")
	}
}

func SetupPostgres(dbname string) error {
	if os.Getenv("PGDATABASE") != "" {
		dbname = os.Getenv("PGDATABASE")
	} else {
		os.Setenv("PGDATABASE", dbname)
	}
	if os.Getenv("PGSSLMODE") == "" {
		os.Setenv("PGSSLMODE", "disable")
	}

	connConfig := pgx.ConnConfig{
		Host:     "/var/run/postgresql",
		Database: "postgres",
	}

	db, err := pgx.Connect(connConfig)
	if err != nil {
		return err
	}

	defer db.Close()
	if _, err := db.Exec(fmt.Sprintf("DROP DATABASE IF EXISTS %s", dbname)); err != nil {
		return err
	}
	if _, err := db.Exec(fmt.Sprintf("CREATE DATABASE %s", dbname)); err != nil {
		return err
	}
	return nil
}

func Example_CustomType() {
	conn, err := pgx.Connect(*defaultConnConfig)
	if err != nil {
		fmt.Printf("Unable to establish connection: %v", err)
		return
	}

	var p NullPoint
	err = conn.QueryRow("select null::point").Scan(&p)
	if err != nil {
		fmt.Println(err)
		return
	}
	fmt.Println(p)

	err = conn.QueryRow("select point(1.5,2.5)").Scan(&p)
	if err != nil {
		fmt.Println(err)
		return
	}
	fmt.Println(p)
	// Output:
	// null point
	// 1.5, 2.5
}

func TestConnectWithRuntimeParams(t *testing.T) {
	t.Parallel()

	connConfig := *defaultConnConfig
	connConfig.RuntimeParams = map[string]string{
		"application_name": "pgxtest",
		"search_path":      "myschema",
	}

	conn, err := pgx.Connect(connConfig)
	if err != nil {
		t.Fatalf("Unable to establish connection: %v", err)
	}
	defer conn.Close()

	var s string
	err = conn.QueryRow("show application_name").Scan(&s)
	if err != nil {
		t.Fatalf("QueryRow Scan unexpectedly failed: %v", err)
	}
	if s != "pgxtest" {
		t.Errorf("Expected application_name to be %s, but it was %s", "pgxtest", s)
	}

	err = conn.QueryRow("show search_path").Scan(&s)
	if err != nil {
		t.Fatalf("QueryRow Scan unexpectedly failed: %v", err)
	}
	if s != "myschema" {
		t.Errorf("Expected search_path to be %s, but it was %s", "myschema", s)
	}
}

func TestStressTLSConnection(t *testing.T) {
	t.Parallel()

	if tlsConnConfig == nil {
		t.Skip("Skipping due to undefined tlsConnConfig")
	}

	if testing.Short() {
		t.Skip("Skipping due to testing -short")
	}

	conn, err := pgx.Connect(*tlsConnConfig)
	if err != nil {
		t.Fatalf("Unable to establish connection: %v", err)
	}
	defer conn.Close()

	for i := 0; i < 50; i++ {
		sql := `select * from generate_series(1, $1)`

		rows, err := conn.Query(sql, 2000000)
		if err != nil {
			t.Fatal(err)
		}

		var n int32
		for rows.Next() {
			rows.Scan(&n)
		}

		if rows.Err() != nil {
			t.Fatalf("queryCount: %d, Row number: %d. %v", i, n, rows.Err())
		}
	}
}

func main() {
	var err error
	conn, err = pgx.Connect(extractConfig())
	if err != nil {
		fmt.Fprintf(os.Stderr, "Unable to connection to database: %v\n", err)
		os.Exit(1)
	}

	if len(os.Args) == 1 {
		printHelp()
		os.Exit(0)
	}

	switch os.Args[1] {
	case "list":
		err = listTasks()
		if err != nil {
			fmt.Fprintf(os.Stderr, "Unable to list tasks: %v\n", err)
			os.Exit(1)
		}

	case "add":
		err = addTask(os.Args[2])
		if err != nil {
			fmt.Fprintf(os.Stderr, "Unable to add task: %v\n", err)
			os.Exit(1)
		}

	case "update":
		n, err := strconv.ParseInt(os.Args[2], 10, 32)
		if err != nil {
			fmt.Fprintf(os.Stderr, "Unable convert task_num into int32: %v\n", err)
			os.Exit(1)
		}
		err = updateTask(int32(n), os.Args[3])
		if err != nil {
			fmt.Fprintf(os.Stderr, "Unable to update task: %v\n", err)
			os.Exit(1)
		}

	case "remove":
		n, err := strconv.ParseInt(os.Args[2], 10, 32)
		if err != nil {
			fmt.Fprintf(os.Stderr, "Unable convert task_num into int32: %v\n", err)
			os.Exit(1)
		}
		err = removeTask(int32(n))
		if err != nil {
			fmt.Fprintf(os.Stderr, "Unable to remove task: %v\n", err)
			os.Exit(1)
		}

	default:
		fmt.Fprintln(os.Stderr, "Invalid command")
		printHelp()
		os.Exit(1)
	}
}

func connect(c *C, n int, db string) *pgx.Conn {
	conn, err := pgx.Connect(pgx.ConnConfig{
		Host:     "127.0.0.1",
		Port:     54320 + uint16(n),
		User:     "flynn",
		Password: "password",
		Database: db,
	})
	c.Assert(err, IsNil)
	return conn
}

func connect(c *C, s state.Database, db string) *pgx.Conn {
	port, _ := strconv.Atoi(s.(*Postgres).port)
	conn, err := pgx.Connect(pgx.ConnConfig{
		Host:     "127.0.0.1",
		Port:     uint16(port),
		User:     "flynn",
		Password: "password",
		Database: db,
	})
	c.Assert(err, IsNil)
	return conn
}

func TestConnectWithConnectionRefused(t *testing.T) {
	t.Parallel()

	// Presumably nothing is listening on 127.0.0.1:1
	bad := *defaultConnConfig
	bad.Host = "127.0.0.1"
	bad.Port = 1

	_, err := pgx.Connect(bad)
	if err == nil {
		t.Fatal("Expected error establishing connection to bad port")
	}
}

func TestConnectWithInvalidUser(t *testing.T) {
	t.Parallel()

	if invalidUserConnConfig == nil {
		return
	}

	_, err := pgx.Connect(*invalidUserConnConfig)
	pgErr, ok := err.(pgx.PgError)
	if !ok {
		t.Fatalf("Expected to receive a PgError with code 28000, instead received: %v", err)
	}
	if pgErr.Code != "28000" && pgErr.Code != "28P01" {
		t.Fatalf("Expected to receive a PgError with code 28000 or 28P01, instead received: %v", pgErr)
	}
}

func TestConnectWithMD5Password(t *testing.T) {
	t.Parallel()

	if md5ConnConfig == nil {
		t.Skip("Skipping due to undefined md5ConnConfig")
	}

	conn, err := pgx.Connect(*md5ConnConfig)
	if err != nil {
		t.Fatal("Unable to establish connection: " + err.Error())
	}

	err = conn.Close()
	if err != nil {
		t.Fatal("Unable to close connection")
	}
}