| 19 | ) |
| 20 | |
| 21 | func CreateTestServer(t *testing.T, port int) (ctx context.Context, pgServer *pgserver.Server, conn *pgx.Conn, close func() error, err error) { |
| 22 | provider := catalog.NewInMemoryDBProvider() |
| 23 | |
| 24 | // Postgres tables are created in the `public` schema by default. |
| 25 | // Create the `public` schema if it doesn't exist. |
| 26 | _, err = provider.Pool().ExecContext(context.Background(), "CREATE SCHEMA IF NOT EXISTS public") |
| 27 | if err != nil { |
| 28 | return nil, nil, nil, nil, err |
| 29 | } |
| 30 | |
| 31 | engine := sqle.NewDefault(provider) |
| 32 | |
| 33 | builder := backend.NewDuckBuilder(engine.Analyzer.ExecBuilder, provider) |
| 34 | engine.Analyzer.ExecBuilder = builder |
| 35 | |
| 36 | config := server.Config{ |
| 37 | Address: fmt.Sprintf("127.0.0.1:%d", port-1), // Unused |
| 38 | } |
| 39 | |
| 40 | sb := backend.NewSessionBuilder(provider) |
| 41 | tracer := sql.NoopTracer |
| 42 | |
| 43 | sm := server.NewSessionManager( |
| 44 | sb, tracer, |
| 45 | engine.Analyzer.Catalog.Database, |
| 46 | engine.MemoryManager, |
| 47 | engine.ProcessList, |
| 48 | config.Address, |
| 49 | ) |
| 50 | |
| 51 | var connID atomic.Uint32 |
| 52 | |
| 53 | pgServer, err = pgserver.NewServer( |
| 54 | provider, |
| 55 | "127.0.0.1", port, |
| 56 | "", |
| 57 | func() *sql.Context { |
| 58 | session := backend.NewSession(memory.NewSession(sql.NewBaseSession(), provider), provider) |
| 59 | return sql.NewContext(context.Background(), sql.WithSession(session)) |
| 60 | }, |
| 61 | pgserver.WithEngine(engine), |
| 62 | pgserver.WithSessionManager(sm), |
| 63 | pgserver.WithConnID(&connID), |
| 64 | ) |
| 65 | if err != nil { |
| 66 | panic(err) |
| 67 | } |
| 68 | pgConfig.Init() |
| 69 | go pgServer.Start() |
| 70 | |
| 71 | ctx = context.Background() |
| 72 | |
| 73 | close = func() error { |
| 74 | pgServer.Listener.Close() |
| 75 | return errors.Join( |
| 76 | provider.Pool().Close(), |
| 77 | provider.Close(), |
| 78 | ) |