/*
* Licensed to Crate.io GmbH ("Crate") under one or more contributor
* license agreements. See the NOTICE file distributed with this work for
* additional information regarding copyright ownership. Crate licenses
* this file to you under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License. You may
* obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
* WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
* License for the specific language governing permissions and limitations
* under the License.
*
* However, if you have executed another commercial license agreement
* with Crate these terms will supersede the license and you may use the
* software solely pursuant to the terms of the relevant commercial agreement.
*/
package org.finos.legend.engine.postgres;
import io.netty.buffer.ByteBuf;
import io.netty.channel.Channel;
import io.netty.channel.ChannelFuture;
import io.netty.channel.ChannelFutureListener;
import io.opentelemetry.context.Context;
import io.opentelemetry.context.propagation.TextMapSetter;
import java.nio.charset.StandardCharsets;
import java.sql.ParameterMetaData;
import java.sql.SQLException;
import java.util.HashMap;
import java.util.List;
import java.util.Locale;
import java.util.Map;
import java.util.Objects;
import java.util.SortedSet;
import org.finos.legend.engine.postgres.handler.PostgresResultSet;
import org.finos.legend.engine.postgres.handler.PostgresResultSetMetaData;
import org.finos.legend.engine.postgres.types.PGType;
import org.finos.legend.engine.postgres.types.PGTypes;
import org.finos.legend.engine.postgres.utils.ErrorMessageFormatter;
import org.finos.legend.engine.postgres.utils.OpenTelemetryUtil;
import org.finos.legend.engine.shared.core.operational.Assert;
import org.slf4j.Logger;
/**
* Regular data packet is in the following format:
*
* +----------+-----------+----------+ | char tag | int32 len | payload |
* +----------+-----------+----------+
*
* The tag indicates the message type, the second field is the length of the packet (excluding the
* tag, but including the length itself)
*
*
* See postgresql docs
*/
public class Messages
{
private static final Logger LOGGER = org.slf4j.LoggerFactory.getLogger(Messages.class);
private static final byte[] METHOD_NAME_CLIENT_AUTH = "ClientAuthentication".getBytes(
StandardCharsets.UTF_8);
private final ErrorMessageFormatter errorMessageFormatter;
public Messages(ErrorMessageFormatter errorMessageFormatter)
{
this.errorMessageFormatter = errorMessageFormatter;
}
public ChannelFuture sendAuthenticationOK(Channel channel)
{
ByteBuf buffer = channel.alloc().buffer(9);
buffer.writeByte('R');
buffer.writeInt(8); // size excluding char
buffer.writeInt(0);
ChannelFuture channelFuture = channel.writeAndFlush(buffer);
if (LOGGER.isTraceEnabled())
{
channelFuture.addListener(
(ChannelFutureListener) future -> LOGGER.trace("sentAuthenticationOK"));
}
return channelFuture;
}
/**
* | 'C' | int32 len | str commandTag
*
* @param query :the query
* @param rowCount : number of rows in the result set or number of rows affected by the DML
* statement
*/
ChannelFuture sendCommandComplete(Channel channel, String query, long rowCount)
{
query = query.trim().split(" ", 2)[0].toUpperCase(Locale.ENGLISH);
String commandTag;
/*
* from https://www.postgresql.org/docs/current/static/protocol-message-formats.html:
*
* For an INSERT command, the tag is INSERT oid rows, where rows is the number of rows inserted.
* oid is the object ID of the inserted row if rows is 1 and the target table has OIDs; otherwise oid is 0.
*/
if ("BEGIN".equals(query))
{
commandTag = "BEGIN";
}
else if ("INSERT".equals(query))
{
commandTag = "INSERT 0 " + rowCount;
}
else
{
commandTag = query + " " + rowCount;
}
byte[] commandTagBytes = commandTag.getBytes(StandardCharsets.UTF_8);
int length = 4 + commandTagBytes.length + 1;
ByteBuf buffer = channel.alloc().buffer(length + 1);
buffer.writeByte('C');
buffer.writeInt(length);
writeCString(buffer, commandTagBytes);
ChannelFuture channelFuture = channel.writeAndFlush(buffer);
if (LOGGER.isTraceEnabled())
{
channelFuture.addListener(
(ChannelFutureListener) future -> LOGGER.trace("sentCommandComplete"));
}
return channelFuture;
}
/**
* ReadyForQuery (B)
*
* Byte1('Z') Identifies the message type. ReadyForQuery is sent whenever the backend is ready for
* a new query cycle.
*
* Int32(5) Length of message contents in bytes, including self.
*
* Byte1 Current backend transaction status indicator. Possible values are 'I' if idle (not in a
* transaction block); 'T' if in a transaction block; or 'E' if in a failed transaction block
* (queries will be rejected until block is ended).
*/
ChannelFuture sendReadyForQuery(Channel channel)
{
ByteBuf buffer = channel.alloc().buffer(6);
buffer.writeByte('Z');
buffer.writeInt(5);
buffer.writeByte('I');
ChannelFuture channelFuture = channel.writeAndFlush(buffer);
if (LOGGER.isTraceEnabled())
{
channelFuture.addListener(
(ChannelFutureListener) future -> LOGGER.trace("sentReadyForQuery"));
}
return channelFuture;
}
/**
* | 'S' | int32 len | str name | str value
*
* See https://www.postgresql.org/docs/9.2/static/protocol-flow.html#PROTOCOL-ASYNC
*
* > At present there is a hard-wired set of parameters for which ParameterStatus will be
* generated: they are
*
* - server_version, - server_encoding, - client_encoding, - application_name, - is_superuser, -
* session_authorization, - DateStyle, - IntervalStyle, - TimeZone, - integer_datetimes, -
* standard_conforming_string
*/
void sendParameterStatus(Channel channel, final String name, final String value)
{
byte[] nameBytes = name.getBytes(StandardCharsets.UTF_8);
byte[] valueBytes = value.getBytes(StandardCharsets.UTF_8);
int length = 4 + nameBytes.length + 1 + valueBytes.length + 1;
ByteBuf buffer = channel.alloc().buffer(length + 1);
buffer.writeByte('S');
buffer.writeInt(length);
writeCString(buffer, nameBytes);
writeCString(buffer, valueBytes);
ChannelFuture channelFuture = channel.writeAndFlush(buffer);
if (LOGGER.isTraceEnabled())
{
channelFuture.addListener(
(ChannelFutureListener) future -> LOGGER.trace("sentParameterStatus {}={}", name, value));
}
}
void sendAuthenticationError(Channel channel, String message)
{
LOGGER.warn(message);
byte[] msg = (message != null ? message : "Unknown Auth Error").getBytes(StandardCharsets.UTF_8);
byte[] errorCode = PGErrorStatus.INVALID_AUTHORIZATION_SPECIFICATION.code()
.getBytes(StandardCharsets.UTF_8);
sendErrorResponse(channel, message, msg, PGError.SEVERITY_FATAL, null, null,
METHOD_NAME_CLIENT_AUTH, errorCode);
}
private String buildErrorMessage(Throwable throwable)
{
TextMapSetter