Examples with QueryExecutorImpl org.postgresql.core.v3.QueryExecutorImpl used on opensource projects

Example 1 with QueryExecutorImpl

use of org.postgresql.core.v3.QueryExecutorImpl in project pgadapter by GoogleCloudPlatform.

the class CopyInMockServerTest method testCopyIn_QueryDuringCopy.

@Test
public void testCopyIn_QueryDuringCopy() throws SQLException, NoSuchFieldException, IllegalAccessException, IOException {
    setupCopyInformationSchemaResults();
    byte[] payload = "5\t5\t5\n".getBytes(StandardCharsets.UTF_8);
    try (Connection connection = DriverManager.getConnection(createUrl())) {
        CopyManager copyManager = new CopyManager(connection.unwrap(BaseConnection.class));
        // Start a copy operation and then try to execute a query during the copy.
        CopyIn copyOperation = copyManager.copyIn("COPY users FROM STDIN;");
        copyOperation.writeToCopy(payload, 0, payload.length);
        copyOperation.flushCopy();
        // Use reflection to get hold of the underlying stream, so we can send an invalid message.
        java.lang.reflect.Field queryExecutorField = CopyOperationImpl.class.getDeclaredField("queryExecutor");
        queryExecutorField.setAccessible(true);
        QueryExecutorImpl queryExecutor = (QueryExecutorImpl) queryExecutorField.get(copyOperation);
        java.lang.reflect.Field pgStreamField = QueryExecutorBase.class.getDeclaredField("pgStream");
        pgStreamField.setAccessible(true);
        PGStream stream = (PGStream) pgStreamField.get(queryExecutor);
        stream.sendChar('Q');
        // Length = 4 + 8 + 1 = 13
        // (msg length 4 bytes, 8 bytes for SELECT 1, 1 byte for \0)
        stream.sendInteger4(13);
        stream.send("SELECT 1".getBytes(StandardCharsets.UTF_8));
        stream.sendChar(0);
        stream.flush();
        // PGAdapter drops the connection if an invalid message is received during a COPY. This is a
        // safety measure as there is no other error handling in the COPY protocol, and the server
        // could otherwise have been completely flushed with garbage if it continued to receive
        // messages after receiving an invalid message.
        SQLException exception = assertThrows(SQLException.class, copyOperation::endCopy);
        assertEquals("Database connection failed when ending copy", exception.getMessage());
    }
    List<CommitRequest> commitRequests = mockSpanner.getRequestsOfType(CommitRequest.class);
    assertTrue(commitRequests.isEmpty());
}