Examples with LocalKeyFactory com.amazonaws.athena.connector.lambda.security.LocalKeyFactory used on opensource projects

Example 11 with LocalKeyFactory

use of com.amazonaws.athena.connector.lambda.security.LocalKeyFactory in project aws-athena-query-federation by awslabs.

the class ElasticsearchMetadataHandlerTest method doListTables.

/**
 * Used to test the doListTables() functionality in the ElasticsearchMetadataHandler class.
 * @throws IOException
 */
@Test
public void doListTables() throws Exception {
    logger.info("doListTables - enter");
    // Hardcoded response with 2 indices.
    Collection<TableName> mockIndices = ImmutableList.of(new TableName("movies", "customer"), new TableName("movies", "movies"));
    // Get real indices.
    when(domainMapProvider.getDomainMap(null)).thenReturn(ImmutableMap.of("movies", "https://search-movies-ne3fcqzfipy6jcrew2wca6kyqu.us-east-1.es.amazonaws.com"));
    handler = new ElasticsearchMetadataHandler(awsGlue, new LocalKeyFactory(), awsSecretsManager, amazonAthena, "spill-bucket", "spill-prefix", domainMapProvider, clientFactory, 10);
    when(mockClient.getAliases()).thenReturn(ImmutableSet.of("movies", ".kibana_1", "customer"));
    ListTablesRequest req = new ListTablesRequest(fakeIdentity(), "queryId", "elasticsearch", "movies", null, UNLIMITED_PAGE_SIZE_VALUE);
    Collection<TableName> realIndices = handler.doListTables(allocator, req).getTables();
    logger.info("doListTables - {}", realIndices);
    // Test 1 - Indices list should NOT be empty.
    assertFalse("Real indices list is empty!", realIndices.isEmpty());
    // Test 2 - Real list and mocked list should have the same indices.
    assertTrue("Real and mocked indices list are different!", indicesEqual(realIndices, mockIndices));
    logger.info("doListTables - exit");
}

Example 12 with LocalKeyFactory

use of com.amazonaws.athena.connector.lambda.security.LocalKeyFactory in project aws-athena-query-federation by awslabs.

the class ElasticsearchMetadataHandlerTest method doGetTable.

/**
 * Used to test the doGetTable() functionality in the ElasticsearchMetadataHandler class.
 * @throws IOException
 */
@Test
public void doGetTable() throws Exception {
    logger.info("doGetTable - enter");
    // Mock mapping.
    Schema mockMapping = SchemaBuilder.newBuilder().addField("mytext", Types.MinorType.VARCHAR.getType()).addField("mykeyword", Types.MinorType.VARCHAR.getType()).addField(new Field("mylong", FieldType.nullable(Types.MinorType.LIST.getType()), Collections.singletonList(new Field("mylong", FieldType.nullable(Types.MinorType.BIGINT.getType()), null)))).addField("myinteger", Types.MinorType.INT.getType()).addField("myshort", Types.MinorType.SMALLINT.getType()).addField("mybyte", Types.MinorType.TINYINT.getType()).addField("mydouble", Types.MinorType.FLOAT8.getType()).addField(new Field("myscaled", new FieldType(true, Types.MinorType.BIGINT.getType(), null, ImmutableMap.of("scaling_factor", "10.0")), null)).addField("myfloat", Types.MinorType.FLOAT4.getType()).addField("myhalf", Types.MinorType.FLOAT4.getType()).addField("mydatemilli", Types.MinorType.DATEMILLI.getType()).addField("mydatenano", Types.MinorType.DATEMILLI.getType()).addField("myboolean", Types.MinorType.BIT.getType()).addField("mybinary", Types.MinorType.VARCHAR.getType()).addField("mynested", Types.MinorType.STRUCT.getType(), ImmutableList.of(new Field("l1long", FieldType.nullable(Types.MinorType.BIGINT.getType()), null), new Field("l1date", FieldType.nullable(Types.MinorType.DATEMILLI.getType()), null), new Field("l1nested", FieldType.nullable(Types.MinorType.STRUCT.getType()), ImmutableList.of(new Field("l2short", FieldType.nullable(Types.MinorType.LIST.getType()), Collections.singletonList(new Field("l2short", FieldType.nullable(Types.MinorType.SMALLINT.getType()), null))), new Field("l2binary", FieldType.nullable(Types.MinorType.VARCHAR.getType()), null))))).build();
    // Real mapping.
    LinkedHashMap<String, Object> mapping = new ObjectMapper().readValue("{\n" + // Index: mishmash
    "  \"mishmash\" : {\n" + "    \"mappings\" : {\n" + // _meta:
    "      \"_meta\" : {\n" + // mynested.l1nested.l2short: LIST<SMALLINT>
    "        \"mynested.l1nested.l2short\" : \"list\",\n" + // mylong: LIST<BIGINT>
    "        \"mylong\" : \"list\"\n" + "      },\n" + "      \"properties\" : {\n" + // mybinary:
    "        \"mybinary\" : {\n" + // type: binary (VARCHAR)
    "          \"type\" : \"binary\"\n" + "        },\n" + // myboolean:
    "        \"myboolean\" : {\n" + // type: boolean (BIT)
    "          \"type\" : \"boolean\"\n" + "        },\n" + // mybyte:
    "        \"mybyte\" : {\n" + // type: byte (TINYINT)
    "          \"type\" : \"byte\"\n" + "        },\n" + // mydatemilli:
    "        \"mydatemilli\" : {\n" + // type: date (DATEMILLI)
    "          \"type\" : \"date\"\n" + "        },\n" + // mydatenano:
    "        \"mydatenano\" : {\n" + // type: date_nanos (DATEMILLI)
    "          \"type\" : \"date_nanos\"\n" + "        },\n" + // mydouble:
    "        \"mydouble\" : {\n" + // type: double (FLOAT8)
    "          \"type\" : \"double\"\n" + "        },\n" + // myfloat:
    "        \"myfloat\" : {\n" + // type: float (FLOAT4)
    "          \"type\" : \"float\"\n" + "        },\n" + // myhalf:
    "        \"myhalf\" : {\n" + // type: half_float (FLOAT4)
    "          \"type\" : \"half_float\"\n" + "        },\n" + // myinteger:
    "        \"myinteger\" : {\n" + // type: integer (INT)
    "          \"type\" : \"integer\"\n" + "        },\n" + // mykeyword:
    "        \"mykeyword\" : {\n" + // type: keyword (VARCHAR)
    "          \"type\" : \"keyword\"\n" + "        },\n" + // mylong: LIST
    "        \"mylong\" : {\n" + // type: long (BIGINT)
    "          \"type\" : \"long\"\n" + "        },\n" + // mynested: STRUCT
    "        \"mynested\" : {\n" + "          \"properties\" : {\n" + // mynested.l1date:
    "            \"l1date\" : {\n" + // type: date_nanos (DATEMILLI)
    "              \"type\" : \"date_nanos\"\n" + "            },\n" + // mynested.l1long:
    "            \"l1long\" : {\n" + // type: long (BIGINT)
    "              \"type\" : \"long\"\n" + "            },\n" + // mynested.l1nested: STRUCT
    "            \"l1nested\" : {\n" + "              \"properties\" : {\n" + // mynested.l1nested.l2binary:
    "                \"l2binary\" : {\n" + // type: binary (VARCHAR)
    "                  \"type\" : \"binary\"\n" + "                },\n" + // mynested.l1nested.l2short: LIST
    "                \"l2short\" : {\n" + // type: short (SMALLINT)
    "                  \"type\" : \"short\"\n" + "                }\n" + "              }\n" + "            }\n" + "          }\n" + "        },\n" + // myscaled:
    "        \"myscaled\" : {\n" + // type: scaled_float (BIGINT)
    "          \"type\" : \"scaled_float\",\n" + // factor: 10
    "          \"scaling_factor\" : 10.0\n" + "        },\n" + // myshort:
    "        \"myshort\" : {\n" + // type: short (SMALLINT)
    "          \"type\" : \"short\"\n" + "        },\n" + // mytext:
    "        \"mytext\" : {\n" + // type: text (VARCHAR)
    "          \"type\" : \"text\"\n" + "        }\n" + "      }\n" + "    }\n" + "  }\n" + "}\n", LinkedHashMap.class);
    LinkedHashMap<String, Object> index = (LinkedHashMap<String, Object>) mapping.get("mishmash");
    LinkedHashMap<String, Object> mappings = (LinkedHashMap<String, Object>) index.get("mappings");
    when(mockClient.getMapping(anyString())).thenReturn(mappings);
    // Get real mapping.
    when(domainMapProvider.getDomainMap(null)).thenReturn(ImmutableMap.of("movies", "https://search-movies-ne3fcqzfipy6jcrew2wca6kyqu.us-east-1.es.amazonaws.com"));
    handler = new ElasticsearchMetadataHandler(awsGlue, new LocalKeyFactory(), awsSecretsManager, amazonAthena, "spill-bucket", "spill-prefix", domainMapProvider, clientFactory, 10);
    GetTableRequest req = new GetTableRequest(fakeIdentity(), "queryId", "elasticsearch", new TableName("movies", "mishmash"));
    GetTableResponse res = handler.doGetTable(allocator, req);
    Schema realMapping = res.getSchema();
    logger.info("doGetTable - {}", res);
    // Test1 - Real mapping must NOT be empty.
    assertTrue("Real mapping is empty!", realMapping.getFields().size() > 0);
    // Test2 - Real and mocked mappings must have the same fields.
    assertTrue("Real and mocked mappings are different!", ElasticsearchSchemaUtils.mappingsEqual(realMapping, mockMapping));
    logger.info("doGetTable - exit");
}

Example 13 with LocalKeyFactory

use of com.amazonaws.athena.connector.lambda.security.LocalKeyFactory in project aws-athena-query-federation by awslabs.

the class HbaseMetadataHandlerTest method setUp.

@Before
public void setUp() throws Exception {
    logger.info("{}: enter", testName.getMethodName());
    handler = new HbaseMetadataHandler(awsGlue, new LocalKeyFactory(), secretsManager, athena, mockConnFactory, "spillBucket", "spillPrefix");
    when(mockConnFactory.getOrCreateConn(anyString())).thenReturn(mockClient);
    allocator = new BlockAllocatorImpl();
}

Example 14 with LocalKeyFactory

use of com.amazonaws.athena.connector.lambda.security.LocalKeyFactory in project aws-athena-query-federation by awslabs.

the class ElasticsearchMetadataHandlerTest method doGetSplits.

/**
 * Used to test the doGetSplits() functionality in the ElasticsearchMetadataHandler class.
 */
@Test
public void doGetSplits() throws Exception {
    logger.info("doGetSplits: enter");
    List<String> partitionCols = new ArrayList<>();
    Block partitions = BlockUtils.newBlock(allocator, "partitionId", Types.MinorType.INT.getType(), 0);
    String continuationToken = null;
    GetSplitsRequest originalReq = new GetSplitsRequest(fakeIdentity(), "queryId", "elasticsearch", new TableName("movies", "customer"), partitions, partitionCols, new Constraints(new HashMap<>()), null);
    GetSplitsRequest req = new GetSplitsRequest(originalReq, continuationToken);
    logger.info("doGetSplits: req[{}]", req);
    // Setup domain and endpoint
    String domain = "movies";
    String endpoint = "https://search-movies-ne3fcqzfipy6jcrew2wca6kyqu.us-east-1.es.amazonaws.com";
    when(domainMapProvider.getDomainMap(null)).thenReturn(ImmutableMap.of(domain, endpoint));
    when(mockClient.getShardIds(anyString(), anyLong())).thenReturn(ImmutableSet.of(new Integer(0), new Integer(1), new Integer(2)));
    // Instantiate handler
    handler = new ElasticsearchMetadataHandler(awsGlue, new LocalKeyFactory(), awsSecretsManager, amazonAthena, "spill-bucket", "spill-prefix", domainMapProvider, clientFactory, 10);
    // Call doGetSplits()
    MetadataResponse rawResponse = handler.doGetSplits(allocator, req);
    assertEquals(MetadataRequestType.GET_SPLITS, rawResponse.getRequestType());
    GetSplitsResponse response = (GetSplitsResponse) rawResponse;
    continuationToken = response.getContinuationToken();
    logger.info("doGetSplits: continuationToken[{}] - numSplits[{}]", new Object[] { continuationToken, response.getSplits().size() });
    // Response should contain 2 splits.
    assertEquals("Response has invalid number of splits", 3, response.getSplits().size());
    Set<String> shardIds = new HashSet<>(2);
    shardIds.add("_shards:0");
    shardIds.add("_shards:1");
    shardIds.add("_shards:2");
    response.getSplits().forEach(split -> {
        assertEquals(endpoint, split.getProperty(domain));
        String shard = split.getProperty(ElasticsearchMetadataHandler.SHARD_KEY);
        assertTrue("Split contains invalid shard: " + shard, shardIds.contains(shard));
        shardIds.remove(shard);
    });
    assertTrue("Continuation criteria violated", response.getContinuationToken() == null);
    logger.info("doGetSplits: exit");
}

Example 15 with LocalKeyFactory

use of com.amazonaws.athena.connector.lambda.security.LocalKeyFactory in project aws-athena-query-federation by awslabs.

the class ElasticsearchMetadataHandlerTest method convertFieldTest.

@Test
public void convertFieldTest() {
    logger.info("convertFieldTest: enter");
    handler = new ElasticsearchMetadataHandler(awsGlue, new LocalKeyFactory(), awsSecretsManager, amazonAthena, "spill-bucket", "spill-prefix", domainMapProvider, clientFactory, 10);
    Field field = handler.convertField("myscaled", "SCALED_FLOAT(10.51)");
    assertEquals("myscaled", field.getName());
    assertEquals("10.51", field.getMetadata().get("scaling_factor"));
    field = handler.convertField("myscaledlist", "ARRAY<SCALED_FLOAT(100)>");
    assertEquals("myscaledlist", field.getName());
    assertEquals(Types.MinorType.LIST.getType(), field.getType());
    assertEquals("100", field.getChildren().get(0).getMetadata().get("scaling_factor"));
    field = handler.convertField("myscaledstruct", "STRUCT<myscaledstruct:SCALED_FLOAT(10.0)>");
    assertEquals(Types.MinorType.STRUCT.getType(), field.getType());
    assertEquals("myscaledstruct", field.getChildren().get(0).getName());
    assertEquals("10.0", field.getChildren().get(0).getMetadata().get("scaling_factor"));
    logger.info("convertFieldTest: exit");
}