Fix JMX connector

This fixes two problems: * Calling toString() the Optional ObjectName * Using toMap() which does not support null values
prestodb · Jul 24, 2015 · 5d2d290 · 5d2d290
1 parent 5821c74
commit 5d2d290
Show file tree

Hide file tree

Showing 4 changed files with 60 additions and 5 deletions.
diff --git a/presto-docs/src/main/sphinx/release/release-0.113.rst b/presto-docs/src/main/sphinx/release/release-0.113.rst
@@ -17,3 +17,4 @@ General Changes
 ---------------
 
 * Add :func:`element_at` function.
+* Fix JMX connector. In the previous release it always returned zero rows.
diff --git a/presto-jmx/src/main/java/com/facebook/presto/connector/jmx/JmxMetadata.java b/presto-jmx/src/main/java/com/facebook/presto/connector/jmx/JmxMetadata.java
@@ -92,7 +92,7 @@ public JmxTableHandle getTableHandle(ConnectorSession session, SchemaTableName t
                 }
                 columns.add(new JmxColumnHandle(connectorId, attribute.getName(), getColumnType(attribute)));
             }
-            return new JmxTableHandle(connectorId, objectName.toString(), columns.build());
+            return new JmxTableHandle(connectorId, objectName.get().toString(), columns.build());
         }
         catch (JMException e) {
             return null;

diff --git a/presto-jmx/src/main/java/com/facebook/presto/connector/jmx/JmxRecordSetProvider.java b/presto-jmx/src/main/java/com/facebook/presto/connector/jmx/JmxRecordSetProvider.java
@@ -30,6 +30,7 @@
 
 import java.util.ArrayList;
 import java.util.Arrays;
+import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
 import java.util.Map.Entry;
@@ -38,7 +39,6 @@
 import static com.facebook.presto.connector.jmx.Types.checkType;
 import static com.google.common.base.Preconditions.checkArgument;
 import static java.util.Objects.requireNonNull;
-import static java.util.stream.Collectors.toMap;
 
 public class JmxRecordSetProvider
         implements ConnectorRecordSetProvider
@@ -166,8 +166,10 @@ private Map<String, Object> getAttributes(Set<String> uniqueColumnNames, JmxTabl
 
         String[] columnNamesArray = uniqueColumnNames.toArray(new String[uniqueColumnNames.size()]);
 
-        return mbeanServer.getAttributes(objectName, columnNamesArray)
-                .asList().stream()
-                .collect(toMap(Attribute::getName, Attribute::getValue));
+        Map<String, Object> map = new HashMap<>();
+        for (Attribute attribute : mbeanServer.getAttributes(objectName, columnNamesArray).asList()) {
+            map.put(attribute.getName(), attribute.getValue());
+        }
+        return map;
     }
 }
diff --git a/presto-jmx/src/test/java/com/facebook/presto/connector/jmx/TestJmxMetadata.java b/presto-jmx/src/test/java/com/facebook/presto/connector/jmx/TestJmxMetadata.java
@@ -0,0 +1,52 @@
+package com.facebook.presto.connector.jmx;
+
+import com.facebook.presto.spi.ConnectorSession;
+import com.facebook.presto.spi.SchemaTableName;
+import com.facebook.presto.spi.type.BigintType;
+import com.google.common.collect.ImmutableList;
+import org.testng.annotations.Test;
+
+import java.util.List;
+
+import static com.facebook.presto.spi.type.TimeZoneKey.UTC_KEY;
+import static com.facebook.presto.spi.type.VarcharType.VARCHAR;
+import static java.lang.management.ManagementFactory.getPlatformMBeanServer;
+import static java.util.Locale.ENGLISH;
+import static org.testng.Assert.assertEquals;
+import static org.testng.Assert.assertTrue;
+
+public class TestJmxMetadata
+{
+    private static final ConnectorSession SESSION = new ConnectorSession("user", UTC_KEY, ENGLISH, System.currentTimeMillis(), null);
+    private static final String RUNTIME_OBJECT = "java.lang:type=Runtime";
+    private static final SchemaTableName RUNTIME_TABLE = new SchemaTableName("jmx", RUNTIME_OBJECT.toLowerCase(ENGLISH));
+
+    private final JmxMetadata metadata = new JmxMetadata("test", getPlatformMBeanServer());
+
+    @Test
+    public void testListSchemas()
+            throws Exception
+    {
+        assertEquals(metadata.listSchemaNames(SESSION), ImmutableList.of("jmx"));
+    }
+
+    @Test
+    public void testListTables()
+    {
+        assertTrue(metadata.listTables(SESSION, "jmx").contains(RUNTIME_TABLE));
+    }
+
+    @Test
+    public void testGetTableHandle()
+            throws Exception
+    {
+        JmxTableHandle handle = metadata.getTableHandle(SESSION, RUNTIME_TABLE);
+        assertEquals(handle.getConnectorId(), "test");
+        assertEquals(handle.getObjectName(), RUNTIME_OBJECT);
+
+        List<JmxColumnHandle> columns = handle.getColumns();
+        assertTrue(columns.contains(new JmxColumnHandle("test", "node", VARCHAR)));
+        assertTrue(columns.contains(new JmxColumnHandle("test", "Name", VARCHAR)));
+        assertTrue(columns.contains(new JmxColumnHandle("test", "StartTime", BigintType.BIGINT)));
+    }
+}