[Legend SQL] Fix error handling (finos#2887)

* Use proper assert library (so that it is active at runtime) * Set error status code on failure * Minor updates in code style * Update config format
levlaz · Jun 6, 2024 · ff47fdc · ff47fdc
1 parent af6f1f9
commit ff47fdc
Show file tree

Hide file tree

Showing 15 changed files with 57 additions and 30 deletions.
diff --git a/...ngine-xt-sql-postgres-server/src/main/java/org/finos/legend/engine/postgres/Messages.java b/...ngine-xt-sql-postgres-server/src/main/java/org/finos/legend/engine/postgres/Messages.java
@@ -42,6 +42,7 @@
 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;
 
 /**
@@ -358,9 +359,9 @@ void sendDataRow(Channel channel, PostgresResultSet rs, List<PGType<?>> columnTy
                      FormatCodes.FormatCode[] formatCodes) throws Exception
     {
         int length = 4 + 2;
-        assert columnTypes.size() == rs.getMetaData().getColumnCount()
-                : "Number of columns in the row must match number of columnTypes. Row: " + rs + " types: "
-                + columnTypes;
+        Assert.assertTrue(columnTypes.size() == rs.getMetaData().getColumnCount(),
+                () -> "Number of columns in the row must match number of columnTypes. Row: " + rs + " types: "
+                        + columnTypes);
 
         ByteBuf buffer = channel.alloc().buffer();
         buffer.writeByte('D');

diff --git a/...gine-xt-sql-postgres-server/src/main/java/org/finos/legend/engine/postgres/PgDecoder.java b/...gine-xt-sql-postgres-server/src/main/java/org/finos/legend/engine/postgres/PgDecoder.java
@@ -26,6 +26,8 @@
 import io.netty.channel.ChannelHandlerContext;
 import io.netty.handler.codec.ByteToMessageDecoder;
 import io.netty.handler.ssl.SslContext;
+import org.finos.legend.engine.shared.core.operational.Assert;
+
 import java.util.List;
 import java.util.function.Supplier;
 
@@ -169,8 +171,8 @@ private ByteBuf decode(ChannelHandlerContext ctx, ByteBuf in)
 
     public void startupDone()
     {
-        assert state == State.STARTUP_PARAMETERS
-                : "Must only call startupDone if state == STARTUP_PARAMETERS";
+        Assert.assertTrue(state == State.STARTUP_PARAMETERS,
+                () -> "Must only call startupDone if state == STARTUP_PARAMETERS");
         state = State.MSG;
     }
 

diff --git a/...-postgres-server/src/main/java/org/finos/legend/engine/postgres/PostgresWireProtocol.java b/...-postgres-server/src/main/java/org/finos/legend/engine/postgres/PostgresWireProtocol.java
@@ -30,6 +30,7 @@
 import io.netty.handler.ssl.SslContext;
 import io.netty.handler.ssl.SslHandler;
 import io.opentelemetry.api.trace.Span;
+import io.opentelemetry.api.trace.StatusCode;
 import io.opentelemetry.api.trace.Tracer;
 import io.opentelemetry.context.Scope;
 import org.finos.legend.engine.postgres.auth.AuthenticationMethod;
@@ -43,6 +44,7 @@
 import org.finos.legend.engine.postgres.utils.OpenTelemetryUtil;
 import org.finos.legend.engine.shared.core.identity.Identity;
 import org.finos.legend.engine.shared.core.kerberos.SubjectTools;
+import org.finos.legend.engine.shared.core.operational.Assert;
 import org.ietf.jgss.GSSContext;
 import org.ietf.jgss.GSSCredential;
 import org.ietf.jgss.GSSException;
@@ -325,7 +327,7 @@ public boolean acceptInboundMessage(Object msg) throws Exception
         @Override
         public void channelRead0(ChannelHandlerContext ctx, ByteBuf buffer) throws Exception
         {
-            assert channel != null : "Channel must be initialized";
+            Assert.assertTrue(channel != null, () -> "Channel must be initialized");
             try
             {
                 dispatchState(buffer, channel);
@@ -541,7 +543,7 @@ private void initAuthentication(Channel channel)
 
     private void finishAuthentication(Channel channel)
     {
-        assert authContext != null : "finishAuthentication() requires an authContext instance";
+        Assert.assertTrue(authContext != null, () -> "finishAuthentication() requires an authContext instance");
         try
         {
             Identity authenticatedUser = authContext.authenticate();
@@ -561,7 +563,7 @@ private void finishAuthentication(Channel channel)
 
     private void finishAuthentication(Channel channel, Subject delegSubject)
     {
-        assert authContext != null : "finishAuthentication() requires an authContext instance";
+        Assert.assertTrue(authContext != null, () -> "finishAuthentication() requires an authContext instance");
         try
         {
             Identity authenticatedUser = KerberosIdentityProvider.getIdentityForSubject(delegSubject);
@@ -846,6 +848,7 @@ private void handleDescribeMessage(ByteBuf buffer, Channel channel) throws Excep
         }
         catch (Exception e)
         {
+            span.setStatus(StatusCode.ERROR, "Failed to handle describe message");
             span.recordException(e);
             OpenTelemetryUtil.TOTAL_FAILURE_METADATA.add(1);
             throw e;
@@ -932,6 +935,7 @@ private void handleExecute(ByteBuf buffer, DelayableWriteChannel channel)
         }
         catch (Exception e)
         {
+            span.setStatus(StatusCode.ERROR, "Failed to execute");
             span.recordException(e);
             throw PostgresServerException.wrapException(e);
         }
@@ -1017,7 +1021,7 @@ void handleSimpleQuery(ByteBuf buffer, final DelayableWriteChannel channel)
         try (Scope scope = span.makeCurrent())
         {
             String queryString = readCString(buffer);
-            assert queryString != null : "query must not be nulL";
+            Assert.assertTrue(queryString != null, () -> "query must not be nulL");
             span.setAttribute("query", queryString);
 
             if (queryString.isEmpty() || ";".equals(queryString))
@@ -1037,6 +1041,7 @@ void handleSimpleQuery(ByteBuf buffer, final DelayableWriteChannel channel)
         }
         catch (Exception e)
         {
+            span.setStatus(StatusCode.ERROR, "Failed to handle simple query");
             span.recordException(e);
             throw e;
         }

diff --git a/...engine-xt-sql-postgres-server/src/main/java/org/finos/legend/engine/postgres/Session.java b/...engine-xt-sql-postgres-server/src/main/java/org/finos/legend/engine/postgres/Session.java
@@ -22,6 +22,7 @@
 package org.finos.legend.engine.postgres;
 
 import io.opentelemetry.api.trace.Span;
+import io.opentelemetry.api.trace.StatusCode;
 import io.opentelemetry.api.trace.Tracer;
 import io.opentelemetry.context.Context;
 import io.opentelemetry.context.Scope;
@@ -381,6 +382,7 @@ public CompletableFuture<?> execute(String portalName, int maxRows, ResultSetRec
         }
         catch (Exception e)
         {
+            span.setStatus(StatusCode.ERROR, "Failed to execute");
             span.recordException(e);
             throw PostgresServerException.wrapException(e);
         }
@@ -418,6 +420,7 @@ public CompletableFuture<?> executeSimple(String query, ResultSetReceiver result
         }
         catch (Exception e)
         {
+            span.setStatus(StatusCode.ERROR, "Failed to execute simple query");
             span.recordException(e);
             throw PostgresServerException.wrapException(e);
         }
@@ -564,6 +567,7 @@ public Boolean call() throws Exception
             }
             catch (Exception e)
             {
+                span.setStatus(StatusCode.ERROR, "Failed to execute");
                 span.recordException(e);
                 resultSetReceiver.fail(e);
                 OpenTelemetryUtil.TOTAL_FAILURE_EXECUTE.add(1);
@@ -615,6 +619,7 @@ public Boolean call() throws Exception
             }
             catch (Exception e)
             {
+                span.setStatus(StatusCode.ERROR, "Failed to execute");
                 span.recordException(e);
                 resultSetReceiver.fail(e);
                 OpenTelemetryUtil.TOTAL_FAILURE_EXECUTE.add(1);

diff --git a/...stgres-server/src/main/java/org/finos/legend/engine/postgres/types/BaseTimestampType.java b/...stgres-server/src/main/java/org/finos/legend/engine/postgres/types/BaseTimestampType.java
@@ -22,6 +22,7 @@
 package org.finos.legend.engine.postgres.types;
 
 import io.netty.buffer.ByteBuf;
+import org.finos.legend.engine.shared.core.operational.Assert;
 
 abstract class BaseTimestampType extends PGType
 {
@@ -73,8 +74,8 @@ private static long toPgTimestamp(long unixTsInMs)
     @Override
     public Object readBinaryValue(ByteBuf buffer, int valueLength)
     {
-        assert valueLength == TYPE_LEN : "valueLength must be " + TYPE_LEN +
-                " because timestamp is a 64 bit long. Actual length: " + valueLength;
+        Assert.assertTrue(valueLength == TYPE_LEN, () -> "valueLength must be " + TYPE_LEN +
+                " because timestamp is a 64 bit long. Actual length: " + valueLength);
         long microSecondsSince2K = buffer.readLong();
         return toCrateTimestamp(microSecondsSince2K);
     }

diff --git a/...t-sql-postgres-server/src/main/java/org/finos/legend/engine/postgres/types/BitString.java b/...t-sql-postgres-server/src/main/java/org/finos/legend/engine/postgres/types/BitString.java
@@ -21,15 +21,17 @@
 
 package org.finos.legend.engine.postgres.types;
 
+import org.finos.legend.engine.shared.core.operational.Assert;
+
 import java.util.BitSet;
 
 class BitString
 {
 
     public static BitString ofBitString(String bitString)
     {
-        assert bitString.startsWith("B'") : "Bitstring must start with B'";
-        assert bitString.endsWith("'") : "Bitstrign must end with '";
+        Assert.assertTrue(bitString.startsWith("B'"), () -> "BitString must start with B'");
+        Assert.assertTrue(bitString.endsWith("'"), () -> "BitString must end with '");
         return ofRawBits(bitString.substring(2, bitString.length() - 1));
     }
 

diff --git a/...-xt-sql-postgres-server/src/main/java/org/finos/legend/engine/postgres/types/BitType.java b/...-xt-sql-postgres-server/src/main/java/org/finos/legend/engine/postgres/types/BitType.java
@@ -22,6 +22,8 @@
 package org.finos.legend.engine.postgres.types;
 
 import io.netty.buffer.ByteBuf;
+import org.finos.legend.engine.shared.core.operational.Assert;
+
 import java.nio.charset.StandardCharsets;
 import java.util.BitSet;
 
@@ -102,7 +104,7 @@ public String readBinaryValue(ByteBuf buffer, int valueLength)
     @Override
     byte[] encodeAsUTF8Text(String value)
     {
-        assert length >= 0 : "BitType length must be set";
+        Assert.assertTrue(length >= 0, () -> "BitType length must be set");
         // See `varbit_out` in src/backend/utils/adt/varbit.c of PostgreSQL
         return value.getBytes(StandardCharsets.UTF_8);
     }

diff --git a/...sql-postgres-server/src/main/java/org/finos/legend/engine/postgres/types/BooleanType.java b/...sql-postgres-server/src/main/java/org/finos/legend/engine/postgres/types/BooleanType.java
@@ -23,6 +23,8 @@
 
 import com.google.common.collect.ImmutableList;
 import io.netty.buffer.ByteBuf;
+import org.finos.legend.engine.shared.core.operational.Assert;
+
 import java.nio.ByteBuffer;
 import java.util.Collection;
 
@@ -87,8 +89,8 @@ byte[] encodeAsUTF8Text(Boolean value)
     @Override
     public Boolean readBinaryValue(ByteBuf buffer, int valueLength)
     {
-        assert valueLength == TYPE_LEN : "length should be " + TYPE_LEN +
-                " because boolean is just a byte. Actual length: " + valueLength;
+        Assert.assertTrue(valueLength == TYPE_LEN, () -> "length should be " + TYPE_LEN +
+                " because boolean is just a byte. Actual length: " + valueLength);
         byte value = buffer.readByte();
         switch (value)
         {

diff --git a/...xt-sql-postgres-server/src/main/java/org/finos/legend/engine/postgres/types/CharType.java b/...xt-sql-postgres-server/src/main/java/org/finos/legend/engine/postgres/types/CharType.java
@@ -22,6 +22,8 @@
 package org.finos.legend.engine.postgres.types;
 
 import io.netty.buffer.ByteBuf;
+import org.finos.legend.engine.shared.core.operational.Assert;
+
 import java.nio.charset.StandardCharsets;
 
 class CharType extends PGType<Byte>
@@ -65,7 +67,7 @@ public int writeAsBinary(ByteBuf buffer, Byte value)
     @Override
     public Byte readBinaryValue(ByteBuf buffer, int valueLength)
     {
-        assert valueLength == 1 : "char must have 1 byte";
+        Assert.assertTrue(valueLength == 1, () -> "char must have 1 byte");
         return buffer.readByte();
     }
 

diff --git a/...sql-postgres-server/src/main/java/org/finos/legend/engine/postgres/types/IntegerType.java b/...sql-postgres-server/src/main/java/org/finos/legend/engine/postgres/types/IntegerType.java
@@ -22,6 +22,8 @@
 package org.finos.legend.engine.postgres.types;
 
 import io.netty.buffer.ByteBuf;
+import org.finos.legend.engine.shared.core.operational.Assert;
+
 import java.nio.charset.StandardCharsets;
 
 class IntegerType extends PGType<Integer>
@@ -74,8 +76,8 @@ protected byte[] encodeAsUTF8Text(Integer value)
     @Override
     public Integer readBinaryValue(ByteBuf buffer, int valueLength)
     {
-        assert valueLength == TYPE_LEN
-                : "length should be " + TYPE_LEN + " because int is int32. Actual length: " + valueLength;
+        Assert.assertTrue(valueLength == TYPE_LEN,
+                () -> "length should be " + TYPE_LEN + " because int is int32. Actual length: " + valueLength);
         return buffer.readInt();
     }
 

diff --git a/...-xt-sql-postgres-server/src/main/java/org/finos/legend/engine/postgres/types/OidType.java b/...-xt-sql-postgres-server/src/main/java/org/finos/legend/engine/postgres/types/OidType.java
@@ -22,6 +22,8 @@
 package org.finos.legend.engine.postgres.types;
 
 import io.netty.buffer.ByteBuf;
+import org.finos.legend.engine.shared.core.operational.Assert;
+
 import java.nio.charset.StandardCharsets;
 
 public class OidType extends PGType<Integer>
@@ -68,8 +70,8 @@ public int writeAsBinary(ByteBuf buffer, Integer value)
     @Override
     public Integer readBinaryValue(ByteBuf buffer, int valueLength)
     {
-        assert valueLength == TYPE_LEN
-                : "length should be " + TYPE_LEN + " because oid is int32. Actual length: " + valueLength;
+        Assert.assertTrue(valueLength == TYPE_LEN,
+                () -> "length should be " + TYPE_LEN + " because oid is int32. Actual length: " + valueLength);
         return buffer.readInt();
     }
 

diff --git a/...sql-postgres-server/src/main/java/org/finos/legend/engine/postgres/types/VarCharType.java b/...sql-postgres-server/src/main/java/org/finos/legend/engine/postgres/types/VarCharType.java
@@ -22,6 +22,8 @@
 package org.finos.legend.engine.postgres.types;
 
 import io.netty.buffer.ByteBuf;
+import org.finos.legend.engine.shared.core.operational.Assert;
+
 import java.nio.charset.StandardCharsets;
 
 class VarCharType extends PGType<Object>
@@ -69,7 +71,7 @@ public String typeCategory()
     @Override
     public int writeAsBinary(ByteBuf buffer, Object value)
     {
-        assert value instanceof String : "value must be string, got: " + value;
+        Assert.assertTrue(value instanceof String, () -> "value must be string, got: " + value);
         String string = (String) value;
         int writerIndex = buffer.writerIndex();
         buffer.writeInt(0);
@@ -87,7 +89,7 @@ public int writeAsText(ByteBuf buffer, Object value)
     @Override
     protected byte[] encodeAsUTF8Text(Object value)
     {
-        assert value instanceof String : "value must be string, got: " + value;
+        Assert.assertTrue(value instanceof String, () -> "value must be string, got: " + value);
         String string = (String) value;
         return string.getBytes(StandardCharsets.UTF_8);
     }

diff --git a/...stgres-server/src/main/java/org/finos/legend/engine/postgres/utils/OpenTelemetryUtil.java b/...stgres-server/src/main/java/org/finos/legend/engine/postgres/utils/OpenTelemetryUtil.java
@@ -25,7 +25,7 @@
 
 public class OpenTelemetryUtil
 {
-    private static final String LEGEND_ENGINE_XTS_SQL = "alloy-sql-server";
+    private static final String LEGEND_ENGINE_XTS_SQL = "legend-sql-server";
     private static final OpenTelemetry OPEN_TELEMETRY = GlobalOpenTelemetry.get();
 
     public static final LongUpDownCounter

diff --git a/...ql-postgres-server/src/test/java/org/finos/legend/engine/postgres/PostgresServerTest.java b/...ql-postgres-server/src/test/java/org/finos/legend/engine/postgres/PostgresServerTest.java
@@ -65,7 +65,7 @@ public static void setUp()
 
         testPostgresServer = new TestPostgresServer(serverConfig, legendSessionFactory,
                 (user, connectionProperties) -> new NoPasswordAuthenticationMethod(new AnonymousIdentityProvider()),
-                new Messages((exception) -> exception.getMessage()));
+                new Messages(Throwable::getMessage));
         testPostgresServer.startUp();
     }
 
@@ -107,10 +107,11 @@ public void testParameterMetadata() throws SQLException
         }
     }
 
-    @Test
     /**
      * Verify that schema created as part of the metadata H2 DB creation exits
+     * @throws SQLException on errors
      */
+    @Test
     public void testInitSchemaCreation() throws SQLException
     {
         try (

diff --git a/...sql/legend-engine-xt-sql-postgres-server/src/test/resources/legendConfigWithKerberos.json b/...sql/legend-engine-xt-sql-postgres-server/src/test/resources/legendConfigWithKerberos.json
@@ -3,12 +3,10 @@
   "authenticationMethod": "GSS",
   "identityType": "KERBEROS",
   "handler": {
+    "className": "org.finos.legend.engine.postgres.config.LegendHandlerConfig",
     "protocol": "http",
     "host": "localhost",
-    "port": "123",
-    "projectId": "project",
-    "sessionCookie": "cookie",
-    "type": "LEGEND"
+    "port": "123"
   },
   "gss": {
     "kerberosConfigFile": "/etc/krb5.conf",