io.prestosql.spi.Plugin Java Examples

@Test
public void testJmxStatsExposure()
        throws Exception
{
    Plugin plugin = new JdbcPlugin("base-jdbc", new TestingH2JdbcModule());
    ConnectorFactory factory = getOnlyElement(plugin.getConnectorFactories());
    factory.create("test", ImmutableMap.of("connection-url", "jdbc"), new TestingConnectorContext());
    MBeanServer mbeanServer = getPlatformMBeanServer();
    Set<ObjectName> objectNames = mbeanServer.queryNames(new ObjectName("io.prestosql.plugin.jdbc:*"), null);

    assertTrue(objectNames.containsAll(
            ImmutableSet.of(
                    new ObjectName("io.prestosql.plugin.jdbc:type=ConnectionFactory,name=test"),
                    new ObjectName("io.prestosql.plugin.jdbc:type=JdbcClient,name=test"))));

    for (ObjectName objectName : objectNames) {
        MBeanInfo mbeanInfo = mbeanServer.getMBeanInfo(objectName);
        assertNotEquals(mbeanInfo.getAttributes().length, 0, format("Object %s doesn't expose JMX stats", objectName.getCanonicalName()));
    }
}
 

public static Set<String> discoverPlugins(Artifact artifact, ClassLoader classLoader)
        throws IOException
{
    if (!artifact.getExtension().equals("presto-plugin")) {
        throw new RuntimeException("Unexpected extension for main artifact: " + artifact);
    }

    File file = artifact.getFile();
    if (!file.getPath().endsWith("/target/classes")) {
        throw new RuntimeException("Unexpected file for main artifact: " + file);
    }
    if (!file.isDirectory()) {
        throw new RuntimeException("Main artifact file is not a directory: " + file);
    }

    if (new File(file, SERVICES_FILE).exists()) {
        throw new RuntimeException("Unexpected services file in artifact directory: " + file);
    }

    return listClasses(file.toPath()).stream()
            .filter(name -> classInterfaces(name, classLoader).contains(Plugin.class.getName()))
            .collect(toImmutableSet());
}
 

@Test
public void testRubixCache()
{
    Plugin plugin = new HiveHadoop2Plugin();
    ConnectorFactory connectorFactory = Iterables.getOnlyElement(plugin.getConnectorFactories());

    connectorFactory.create(
            "test",
            ImmutableMap.<String, String>builder()
                    .put("hive.cache.enabled", "true")
                    .put("hive.metastore.uri", "thrift://foo:1234")
                    .put("hive.cache.location", tempDirectory.toString())
                    .build(),
            new TestingConnectorContext())
            .shutdown();
}
 

@Test
public void testHdfsImpersonationAndHiveCachingMutuallyExclusive()
{
    Plugin plugin = new HiveHadoop2Plugin();
    ConnectorFactory connectorFactory = Iterables.getOnlyElement(plugin.getConnectorFactories());

    assertThatThrownBy(() -> connectorFactory.create(
            "test",
            ImmutableMap.<String, String>builder()
                    .put("hive.hdfs.impersonation.enabled", "true")
                    .put("hive.cache.enabled", "true")
                    .put("hive.metastore.uri", "thrift://foo:1234")
                    .put("hive.cache.location", tempDirectory.toString())
                    .build(),
            new TestingConnectorContext())
            .shutdown())
            .hasMessageContaining("Hdfs impersonation is not compatible with Hive caching");
}
 

@Test
public void testGcsAccessTokenAndHiveCachingMutuallyExclusive()
{
    Plugin plugin = new HiveHadoop2Plugin();
    ConnectorFactory connectorFactory = Iterables.getOnlyElement(plugin.getConnectorFactories());

    assertThatThrownBy(() -> connectorFactory.create(
            "test",
            ImmutableMap.<String, String>builder()
                    .put("hive.gcs.use-access-token", "true")
                    .put("hive.cache.enabled", "true")
                    .put("hive.metastore.uri", "thrift://foo:1234")
                    .put("hive.cache.location", tempDirectory.toString())
                    .build(),
            new TestingConnectorContext())
            .shutdown())
            .hasMessageContaining("Use of GCS access token is not compatible with Hive caching");
}
 

@Test
public void testS3SecurityMappingAndHiveCachingMutuallyExclusive()
        throws IOException
{
    Path mappingConfig = Files.createTempFile(null, null);
    Plugin plugin = new HiveHadoop2Plugin();
    ConnectorFactory connectorFactory = Iterables.getOnlyElement(plugin.getConnectorFactories());

    assertThatThrownBy(() -> connectorFactory.create(
            "test",
            ImmutableMap.<String, String>builder()
                    .put("hive.s3.security-mapping.config-file", mappingConfig.toString())
                    .put("hive.cache.enabled", "true")
                    .put("hive.metastore.uri", "thrift://foo:1234")
                    .put("hive.cache.location", tempDirectory.toString())
                    .build(),
            new TestingConnectorContext())
            .shutdown()).hasMessageContaining("S3 security mapping is not compatible with Hive caching");
}
 

public static RuleTester defaultRuleTester(List<Plugin> plugins, Map<String, String> sessionProperties, Optional<Integer> nodeCountForStats)
{
    Session.SessionBuilder sessionBuilder = testSessionBuilder()
            .setCatalog(CATALOG_ID)
            .setSchema("tiny")
            .setSystemProperty("task_concurrency", "1"); // these tests don't handle exchanges from local parallel

    for (Map.Entry<String, String> entry : sessionProperties.entrySet()) {
        sessionBuilder.setSystemProperty(entry.getKey(), entry.getValue());
    }

    Session session = sessionBuilder.build();

    LocalQueryRunner queryRunner = nodeCountForStats
            .map(nodeCount -> LocalQueryRunner.builder(session)
                    .withNodeCountForStats(nodeCount)
                    .build())
            .orElseGet(() -> LocalQueryRunner.create(session));

    queryRunner.createCatalog(session.getCatalog().get(),
            new TpchConnectorFactory(1),
            ImmutableMap.of());
    plugins.stream().forEach(queryRunner::installPlugin);

    return new RuleTester(queryRunner);
}
 

@Test
public void testPlugin()
        throws Exception
{
    Plugin plugin = new RaptorPlugin();
    ConnectorFactory factory = getOnlyElement(plugin.getConnectorFactories());
    assertInstanceOf(factory, RaptorConnectorFactory.class);

    File tmpDir = Files.createTempDir();
    try {
        Map<String, String> config = ImmutableMap.<String, String>builder()
                .put("metadata.db.type", "h2")
                .put("metadata.db.filename", tmpDir.getAbsolutePath())
                .put("storage.data-directory", tmpDir.getAbsolutePath())
                .build();

        factory.create("test", config, new TestingConnectorContext());
    }
    finally {
        deleteRecursively(tmpDir.toPath(), ALLOW_INSECURE);
    }
}
 

@BeforeClass
public void setUp()
        throws Exception
{
    queryRunner = TpchQueryRunnerBuilder.builder().build();
    queryRunner.installPlugin(new Plugin()
    {
        @Override
        public Iterable<ConnectorFactory> getConnectorFactories()
        {
            MockConnectorFactory connectorFactory = MockConnectorFactory.builder()
                    .withListSchemaNames(session -> ImmutableList.of("UPPER_CASE_SCHEMA"))
                    .withListTables((session, schemaNameOrNull) -> {
                        throw new UnsupportedOperationException();
                    })
                    .withGetViews((session, prefix) -> ImmutableMap.of())
                    .build();
            return ImmutableList.of(connectorFactory);
        }
    });
    queryRunner.createCatalog("upper_case_schema_catalog", "mock");
    metadataManager = (MetadataManager) queryRunner.getMetadata();
}
 

@BeforeClass
private void setUp()
        throws Exception
{
    session = testSessionBuilder()
            .setSystemProperty("task_concurrency", "1")
            .setCatalog("tpch")
            .setSchema("tiny")
            .setClientInfo("{\"clientVersion\":\"testVersion\"}")
            .build();
    queryRunner = DistributedQueryRunner.builder(session).setNodeCount(1).build();
    queryRunner.installPlugin(new Plugin()
    {
        @Override
        public Iterable<ConnectorFactory> getConnectorFactories()
        {
            return ImmutableList.of(new MockConnectorFactory.Builder()
                    .withEventListener(new TestingEventListener(generatedEvents))
                    .build());
        }
    });
    queryRunner.createCatalog("mock-catalog", "mock");
    queries = new EventsAwaitingQueries(generatedEvents, queryRunner);
}
 

public static DistributedQueryRunner createQueryRunner(String dbConfigUrl, H2ResourceGroupsDao dao, String environment)
        throws Exception
{
    DistributedQueryRunner queryRunner = DistributedQueryRunner
            .builder(testSessionBuilder().setCatalog("tpch").setSchema("tiny").build())
            .setNodeCount(2)
            .setEnvironment(environment)
            .build();
    try {
        Plugin h2ResourceGroupManagerPlugin = new H2ResourceGroupManagerPlugin();
        queryRunner.installPlugin(h2ResourceGroupManagerPlugin);
        queryRunner.getCoordinator().getResourceGroupManager().get()
                .setConfigurationManager(CONFIGURATION_MANAGER_TYPE, ImmutableMap.of("resource-groups.config-db-url", dbConfigUrl, "node.environment", environment));
        queryRunner.installPlugin(new TpchPlugin());
        queryRunner.createCatalog("tpch", "tpch");
        setup(queryRunner, dao, environment);
        return queryRunner;
    }
    catch (Exception e) {
        queryRunner.close();
        throw e;
    }
}
 

@Override
protected QueryRunner createQueryRunner()
        throws Exception
{
    Session defaultSession = testSessionBuilder()
            .setCatalog("mock")
            .setSchema("default")
            .build();

    DistributedQueryRunner queryRunner = DistributedQueryRunner.builder(defaultSession).build();
    queryRunner.installPlugin(new Plugin()
    {
        @Override
        public Iterable<ConnectorFactory> getConnectorFactories()
        {
            MockConnectorFactory connectorFactory = MockConnectorFactory.builder()
                    .withGetViews((session, schemaTablePrefix) -> ImmutableMap.of())
                    .withListTables((session, s) -> ImmutableList.of(SCHEMA_TABLE_NAME))
                    .withGetColumns(tableName -> getColumns.apply(tableName))
                    .build();
            return ImmutableList.of(connectorFactory);
        }
    });
    queryRunner.createCatalog("mock", "mock", ImmutableMap.of());
    return queryRunner;
}
 

public void run()
        throws Exception
{
    try (TestingPrestoServer server = TestingPrestoServer.create()) {
        for (String pluginClass : options.getPluginClassNames()) {
            Plugin plugin = (Plugin) Class.forName(pluginClass).getConstructor().newInstance();
            server.installPlugin(plugin);
        }

        for (Catalog catalog : options.getCatalogs()) {
            server.createCatalog(catalog.getCatalogName(), catalog.getConnectorName());
        }

        System.out.println(server.getAddress());
        waitForInterruption();
    }
}
 

@Test
public void testCreateConnector()
{
    Plugin plugin = new KuduPlugin();
    ConnectorFactory factory = getOnlyElement(plugin.getConnectorFactories());
    factory.create("test", ImmutableMap.of("kudu.client.master-addresses", "localhost:7051"), new TestingConnectorContext());
}
 

@Test
public void testCreateConnector()
{
    Plugin plugin = new MemSqlPlugin();
    ConnectorFactory factory = getOnlyElement(plugin.getConnectorFactories());
    factory.create("test", ImmutableMap.of("connection-url", "jdbc:mariadb://test"), new TestingConnectorContext());
}
 

@Test
public void testCreateConnector()
        throws Exception
{
    Plugin plugin = new SheetsPlugin();
    ConnectorFactory factory = getOnlyElement(plugin.getConnectorFactories());
    Builder<String, String> propertiesMap = new Builder<String, String>().put("credentials-path", getTestCredentialsPath()).put("metadata-sheet-id", TEST_METADATA_SHEET_ID);
    Connector c = factory.create(GOOGLE_SHEETS, propertiesMap.build(), new TestingConnectorContext());
    assertNotNull(c);
}
 

@Test
public void testPlugin()
{
    Plugin plugin = new ThriftPlugin();
    ConnectorFactory factory = getOnlyElement(plugin.getConnectorFactories());
    assertInstanceOf(factory, ThriftConnectorFactory.class);

    Map<String, String> config = ImmutableMap.of("presto.thrift.client.addresses", "localhost:7779");

    Connector connector = factory.create("test", config, new TestingConnectorContext());
    assertNotNull(connector);
    assertInstanceOf(connector, ThriftConnector.class);
}
 

@Test
public void testCreateConnector()
{
    Plugin plugin = new RedshiftPlugin();
    ConnectorFactory factory = getOnlyElement(plugin.getConnectorFactories());
    factory.create("test", ImmutableMap.of("connection-url", "test"), new TestingConnectorContext());
}
 

@Test
public void testCreateConnector()
{
    Plugin plugin = new OraclePlugin();
    ConnectorFactory factory = getOnlyElement(plugin.getConnectorFactories());
    factory.create("test", ImmutableMap.of("connection-url", "jdbc:oracle:thin//test"), new TestingConnectorContext());
}
 

private void loadPlugin(PluginClassLoader pluginClassLoader)
{
    ServiceLoader<Plugin> serviceLoader = ServiceLoader.load(Plugin.class, pluginClassLoader);
    List<Plugin> plugins = ImmutableList.copyOf(serviceLoader);
    checkState(!plugins.isEmpty(), "No service providers of type %s", Plugin.class.getName());
    for (Plugin plugin : plugins) {
        log.info("Installing %s", plugin.getClass().getName());
        installPlugin(plugin, pluginClassLoader::duplicate);
    }
}
 

@Test
public void testCreateConnector()
{
    Plugin plugin = new DruidJdbcPlugin();
    ConnectorFactory factory = getOnlyElement(plugin.getConnectorFactories());
    factory.create("test", ImmutableMap.of("connection-url", "test"), new TestingConnectorContext());
}
 

@BeforeClass
private void setUp()
        throws Exception
{
    session = testSessionBuilder()
            .setSystemProperty("task_concurrency", "1")
            .setCatalog("tpch")
            .setSchema("tiny")
            .setClientInfo("{\"clientVersion\":\"testVersion\"}")
            .build();
    queryRunner = DistributedQueryRunner.builder(session).setNodeCount(1).build();
    queryRunner.installPlugin(new TpchPlugin());
    queryRunner.installPlugin(new TestingEventListenerPlugin(generatedEvents));
    queryRunner.installPlugin(new ResourceGroupManagerPlugin());
    queryRunner.createCatalog("tpch", "tpch", ImmutableMap.of("tpch.splits-per-node", Integer.toString(SPLITS_PER_NODE)));
    queryRunner.installPlugin(new Plugin()
    {
        @Override
        public Iterable<ConnectorFactory> getConnectorFactories()
        {
            MockConnectorFactory connectorFactory = MockConnectorFactory.builder()
                    .withListTables((session, s) -> ImmutableList.of(new SchemaTableName("default", "test_table")))
                    .withApplyProjection((session, handle, projections, assignments) -> {
                        throw new RuntimeException("Throw from apply projection");
                    })
                    .build();
            return ImmutableList.of(connectorFactory);
        }
    });
    queryRunner.createCatalog("mock", "mock", ImmutableMap.of());
    queryRunner.getCoordinator().getResourceGroupManager().get()
            .setConfigurationManager("file", ImmutableMap.of("resource-groups.config-file", getResourceFilePath("resource_groups_config_simple.json")));
    queries = new EventsAwaitingQueries(generatedEvents, queryRunner);
}
 

@Test
public void testCreateConnector()
{
    Plugin plugin = new SqlServerPlugin();
    ConnectorFactory factory = getOnlyElement(plugin.getConnectorFactories());
    factory.create("test", ImmutableMap.of("connection-url", "test"), new TestingConnectorContext());
}
 

@Test
public void testCreateConnector()
{
    Plugin plugin = new PostgreSqlPlugin();
    ConnectorFactory factory = getOnlyElement(plugin.getConnectorFactories());
    factory.create("test", ImmutableMap.of("connection-url", "test"), new TestingConnectorContext());
}
 

@Override
public void installPlugin(Plugin plugin)
{
    long start = System.nanoTime();
    for (TestingPrestoServer server : servers) {
        server.installPlugin(plugin);
    }
    log.info("Installed plugin %s in %s", plugin.getClass().getSimpleName(), nanosSince(start).convertToMostSuccinctTimeUnit());
}
 

public Plugin getPlugin()
{
    return new Plugin()
    {
        @Override
        public Iterable<ConnectorFactory> getConnectorFactories()
        {
            return ImmutableList.of(getConnectorFactory());
        }
    };
}
 

@Test
public void testRubixCacheWithNonExistingCacheDirectory()
{
    Plugin plugin = new HiveHadoop2Plugin();
    ConnectorFactory connectorFactory = Iterables.getOnlyElement(plugin.getConnectorFactories());

    assertThatThrownBy(() -> connectorFactory.create(
            "test",
            ImmutableMap.<String, String>builder()
                    .put("hive.cache.enabled", "true")
                    .put("hive.cache.start-server-on-coordinator", "true")
                    .put("hive.metastore.uri", "thrift://foo:1234")
                    .put("hive.cache.location", "/tmp/non/existing/directory")
                    .build(),
            new TestingConnectorContext())
            .shutdown())
            .hasRootCauseMessage("None of the cache parent directories exists");

    // cache directories should not be required when cache is not explicitly started on coordinator
    connectorFactory.create(
            "test",
            ImmutableMap.<String, String>builder()
                    .put("hive.cache.enabled", "true")
                    .put("hive.metastore.uri", "thrift://foo:1234")
                    .put("hive.cache.location", "/tmp/non/existing/directory")
                    .build(),
            new TestingConnectorContext())
            .shutdown();
}
 

@Override
public void installPlugin(Plugin plugin)
{
    server.installPlugin(plugin);
}
 

@Setup
public void setup()
        throws Exception
{
    queryRunner = DistributedQueryRunner.builder(session).build();
    queryRunner.installPlugin(new Plugin()
    {
        @Override
        public Iterable<ConnectorFactory> getConnectorFactories()
        {
            Function<ConnectorSession, List<String>> listSchemaNames = session -> IntStream.range(0, Integer.parseInt(schemasCount))
                    .boxed()
                    .map(i -> "stream_" + i)
                    .collect(toImmutableList());

            BiFunction<ConnectorSession, String, List<SchemaTableName>> listTables = (session, schemaNameOrNull) -> {
                List<String> tables = IntStream.range(0, Integer.parseInt(tablesCount))
                        .boxed()
                        .map(i -> "table_" + i)
                        .collect(toImmutableList());
                List<String> schemas;
                if (schemaNameOrNull == null) {
                    schemas = listSchemaNames.apply(session);
                }
                else {
                    schemas = ImmutableList.of(schemaNameOrNull);
                }
                return schemas.stream()
                        .flatMap(schema -> tables.stream().map(table -> new SchemaTableName(schema, table)))
                        .collect(toImmutableList());
            };

            MockConnectorFactory connectorFactory = MockConnectorFactory.builder()
                    .withListSchemaNames(listSchemaNames)
                    .withListTables(listTables)
                    .withGetViews((session, prefix) -> ImmutableMap.of())
                    .build();
            return ImmutableList.of(connectorFactory);
        }
    });
    queryRunner.createCatalog("test_catalog", "mock", ImmutableMap.of());

    query = queries.get(queryId);
}
 

private static void createMySqlPlugin(String url)
{
    Plugin plugin = new MySqlPlugin();
    ConnectorFactory factory = getOnlyElement(plugin.getConnectorFactories());
    factory.create("test", ImmutableMap.of("connection-url", url), new TestingConnectorContext());
}