InterceptImportMessages sits between PSQL and Doltgres, returning all error messages that are encountered. As we rely on PSQL to handle the import process, we normally wouldn't be able to associate error messages with queries, as this information is not returned by PSQL itself. Therefore, we create
(t *testing.T, args InterceptArgs)
| 58 | // information is not returned by PSQL itself. Therefore, we create our own connection to Doltgres, and a server that |
| 59 | // PSQL listens to. We then forward everything from PSQL to Doltgres, while inspecting the messages as they come and go. |
| 60 | func InterceptImportMessages(t *testing.T, args InterceptArgs) (int, chan ImportQueryError) { |
| 61 | psqlPort, err := sql.GetEmptyPort() |
| 62 | require.NoError(t, err) |
| 63 | qeChan := make(chan ImportQueryError) |
| 64 | listener, err := server.NewListener("tcp", fmt.Sprintf("127.0.0.1:%d", psqlPort), "") |
| 65 | if err != nil { |
| 66 | t.Fatal(err) |
| 67 | return psqlPort, qeChan |
| 68 | } |
| 69 | timer := time.NewTimer(5 * time.Second) |
| 70 | timer.Stop() |
| 71 | go func() { |
| 72 | <-timer.C |
| 73 | _ = listener.Close() |
| 74 | }() |
| 75 | |
| 76 | go func() { |
| 77 | for { |
| 78 | psqlConn, err := listener.Accept() |
| 79 | if err != nil { |
| 80 | return |
| 81 | } |
| 82 | timer.Stop() |
| 83 | terminate := &sync.WaitGroup{} |
| 84 | terminate.Add(1) |
| 85 | psqlConnBackend := pgproto3.NewBackend(psqlConn, psqlConn) |
| 86 | doltgresConn, err := (&net.Dialer{}).Dial("tcp", fmt.Sprintf("127.0.0.1:%d", args.DoltgresPort)) |
| 87 | if err != nil { |
| 88 | fmt.Println(err) |
| 89 | return |
| 90 | } |
| 91 | doltgresConnFrontend := pgproto3.NewFrontend(doltgresConn, doltgresConn) |
| 92 | |
| 93 | if err = handleStartup(t, psqlConnBackend, doltgresConnFrontend, psqlConn); err != nil { |
| 94 | fmt.Println(err) |
| 95 | return |
| 96 | } |
| 97 | createPassthrough(passthroughArguments{ |
| 98 | qeChan: qeChan, |
| 99 | terminate: terminate, |
| 100 | psqlConnBackend: psqlConnBackend, |
| 101 | doltgresConnFrontend: doltgresConnFrontend, |
| 102 | triggerBreakpoint: args.TriggerBreakpoint, |
| 103 | skippedQueries: args.SkippedQueries, |
| 104 | breakpointQueries: args.BreakpointQueries, |
| 105 | }) |
| 106 | terminate.Wait() |
| 107 | _ = psqlConn.Close() |
| 108 | _ = doltgresConn.Close() |
| 109 | timer.Reset(5 * time.Second) |
| 110 | } |
| 111 | }() |
| 112 | return psqlPort, qeChan |
| 113 | } |
| 114 | |
| 115 | // handleStartup handles the startup messages. |
| 116 | func handleStartup(t *testing.T, psqlConnBackend *pgproto3.Backend, doltgresConnFrontend *pgproto3.Frontend, clientConn net.Conn) error { |