Open opens a Postgres driver.
(ctx context.Context, _ storepb.Engine, config db.ConnectionConfig)
| 62 | |
| 63 | // Open opens a Postgres driver. |
| 64 | func (d *Driver) Open(ctx context.Context, _ storepb.Engine, config db.ConnectionConfig) (db.Driver, error) { |
| 65 | var pgxConnConfig *pgx.ConnConfig |
| 66 | var err error |
| 67 | |
| 68 | switch config.DataSource.GetAuthenticationType() { |
| 69 | case storepb.DataSource_GOOGLE_CLOUD_SQL_IAM: |
| 70 | pgxConnConfig, err = getCloudSQLConnectionConfig(ctx, config) |
| 71 | case storepb.DataSource_AWS_RDS_IAM: |
| 72 | pgxConnConfig, err = getRDSConnectionConfig(ctx, config) |
| 73 | default: |
| 74 | pgxConnConfig, err = getPGConnectionConfig(config) |
| 75 | } |
| 76 | if err != nil { |
| 77 | return nil, err |
| 78 | } |
| 79 | appName := "bytebase" |
| 80 | if config.ConnectionContext.TaskRunUID != nil { |
| 81 | appName = fmt.Sprintf("bytebase-taskrun-%d", *config.ConnectionContext.TaskRunUID) |
| 82 | } |
| 83 | pgxConnConfig.RuntimeParams["application_name"] = appName |
| 84 | if config.ConnectionContext.ReadOnly { |
| 85 | pgxConnConfig.RuntimeParams["default_transaction_read_only"] = "true" |
| 86 | } |
| 87 | |
| 88 | if config.DataSource.GetSshHost() != "" { |
| 89 | sshClient, err := util.GetSSHClient(config.DataSource) |
| 90 | if err != nil { |
| 91 | return nil, err |
| 92 | } |
| 93 | d.sshClient = sshClient |
| 94 | |
| 95 | pgxConnConfig.DialFunc = func(_ context.Context, network, addr string) (net.Conn, error) { |
| 96 | conn, err := sshClient.Dial(network, addr) |
| 97 | if err != nil { |
| 98 | return nil, err |
| 99 | } |
| 100 | return &util.NoDeadlineConn{Conn: conn}, nil |
| 101 | } |
| 102 | } |
| 103 | |
| 104 | d.databaseName = config.ConnectionContext.DatabaseName |
| 105 | if config.ConnectionContext.DatabaseName != "" { |
| 106 | pgxConnConfig.Database = config.ConnectionContext.DatabaseName |
| 107 | } else if config.DataSource.GetDatabase() != "" { |
| 108 | pgxConnConfig.Database = config.DataSource.GetDatabase() |
| 109 | } else { |
| 110 | pgxConnConfig.Database = "postgres" |
| 111 | } |
| 112 | d.config = config |
| 113 | |
| 114 | pgxConnConfig.OnNotice = d.onNotice |
| 115 | |
| 116 | d.connectionString = stdlib.RegisterConnConfig(pgxConnConfig) |
| 117 | db, err := sql.Open(driverName, d.connectionString) |
| 118 | if err != nil { |
| 119 | return nil, err |
| 120 | } |
| 121 | d.db = db |