embulk-output-elasticsearch5 0.1.0

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
@@ -0,0 +1,482 @@
1
+ package org.embulk.output.elasticsearch;
2
+
3
+ import com.google.common.collect.ImmutableList;
4
+ import com.google.common.collect.ImmutableMap;
5
+ import com.google.common.collect.Lists;
6
+ import org.elasticsearch.action.admin.cluster.state.ClusterStateRequest;
7
+ import org.elasticsearch.action.admin.indices.create.CreateIndexRequest;
8
+ import org.elasticsearch.action.admin.indices.delete.DeleteIndexRequest;
9
+ import org.elasticsearch.action.get.GetResponse;
10
+ import org.elasticsearch.client.Client;
11
+ import org.embulk.EmbulkTestRuntime;
12
+ import org.embulk.config.ConfigException;
13
+ import org.embulk.config.ConfigSource;
14
+ import org.embulk.config.TaskReport;
15
+ import org.embulk.config.TaskSource;
16
+ import org.embulk.output.elasticsearch.ElasticsearchOutputPlugin.PluginTask;
17
+ import org.embulk.spi.Exec;
18
+ import org.embulk.spi.OutputPlugin;
19
+ import org.embulk.spi.Page;
20
+ import org.embulk.spi.PageTestUtils;
21
+ import org.embulk.spi.Schema;
22
+ import org.embulk.spi.TestPageBuilderReader.MockPageOutput;
23
+ import org.embulk.spi.TransactionalPageOutput;
24
+ import org.embulk.spi.time.Timestamp;
25
+ import org.embulk.standards.CsvParserPlugin;
26
+ import org.junit.Before;
27
+ import org.junit.BeforeClass;
28
+ import org.junit.Rule;
29
+ import org.junit.Test;
30
+
31
+ import java.io.ByteArrayOutputStream;
32
+ import java.io.IOException;
33
+ import java.io.InputStream;
34
+ import java.lang.reflect.InvocationTargetException;
35
+ import java.lang.reflect.Method;
36
+ import java.net.UnknownHostException;
37
+ import java.security.GeneralSecurityException;
38
+ import java.text.ParseException;
39
+ import java.text.SimpleDateFormat;
40
+ import java.util.Arrays;
41
+ import java.util.List;
42
+ import java.util.Map;
43
+ import static org.junit.Assert.assertEquals;
44
+ import static org.junit.Assert.assertTrue;
45
+ import static org.junit.Assume.assumeNotNull;
46
+
47
+ public class TestElasticsearchOutputPlugin
48
+ {
49
+ private static String ES_HOST;
50
+ private static int ES_PORT;
51
+ private static List ES_NODES;
52
+ private static String ES_CLUSTER_NAME;
53
+ private static String ES_INDEX;
54
+ private static String ES_INDEX_TYPE;
55
+ private static String ES_ID;
56
+ private static int ES_BULK_ACTIONS;
57
+ private static int ES_BULK_SIZE;
58
+ private static int ES_CONCURRENT_REQUESTS;
59
+ private static String PATH_PREFIX;
60
+
61
+ private MockPageOutput pageOutput;
62
+
63
+ private static String ES_TEST_INDEX = "index_for_unittest";
64
+ private static String ES_TEST_INDEX2 = "index_for_unittest2";
65
+ private static String ES_TEST_ALIAS = "alias_for_unittest";
66
+
67
+ /*
68
+ * This test case requires environment variables
69
+ * ES_HOST
70
+ * ES_INDEX
71
+ * ES_INDEX_TYPE
72
+ */
73
+ @BeforeClass
74
+ public static void initializeConstant()
75
+ {
76
+ ES_HOST = System.getenv("ES_HOST") != null ? System.getenv("ES_HOST") : "";
77
+ ES_PORT = System.getenv("ES_PORT") != null ? Integer.valueOf(System.getenv("ES_PORT")) : 9300;
78
+
79
+ ES_CLUSTER_NAME = System.getenv("ES_CLUSTER_NAME") != null ? System.getenv("ES_CLUSTER_NAME") : "";
80
+ ES_INDEX = System.getenv("ES_INDEX");
81
+ ES_INDEX_TYPE = System.getenv("ES_INDEX_TYPE");
82
+ ES_ID = "id";
83
+ ES_BULK_ACTIONS = System.getenv("ES_BULK_ACTIONS") != null ? Integer.valueOf(System.getenv("ES_BULK_ACTIONS")) : 1000;
84
+ ES_BULK_SIZE = System.getenv("ES_BULK_SIZE") != null ? Integer.valueOf(System.getenv("ES_BULK_SIZE")) : 5242880;
85
+ ES_CONCURRENT_REQUESTS = System.getenv("ES_CONCURRENT_REQUESTS") != null ? Integer.valueOf(System.getenv("ES_CONCURRENT_REQUESTS")) : 5;
86
+
87
+ assumeNotNull(ES_HOST, ES_INDEX, ES_INDEX_TYPE);
88
+
89
+ ES_NODES = Arrays.asList(ImmutableMap.of("host", ES_HOST, "port", ES_PORT));
90
+
91
+ PATH_PREFIX = ElasticsearchOutputPlugin.class.getClassLoader().getResource("sample_01.csv").getPath();
92
+ }
93
+
94
+ @Rule
95
+ public EmbulkTestRuntime runtime = new EmbulkTestRuntime();
96
+ private ElasticsearchOutputPlugin plugin;
97
+
98
+ @Before
99
+ public void createResources()
100
+ throws GeneralSecurityException, NoSuchMethodException,
101
+ IllegalAccessException, InvocationTargetException
102
+ {
103
+ ConfigSource config = config();
104
+ plugin = new ElasticsearchOutputPlugin();
105
+ PluginTask task = config.loadConfig(PluginTask.class);
106
+ pageOutput = new MockPageOutput();
107
+
108
+ Method createClient = ElasticsearchOutputPlugin.class.getDeclaredMethod("createClient", PluginTask.class);
109
+ createClient.setAccessible(true);
110
+ try (Client client = (Client) createClient.invoke(plugin, task)) {
111
+ // Delete alias
112
+ if (client.admin().cluster().state(new ClusterStateRequest()).actionGet().getState().getMetaData().hasAlias(ES_TEST_ALIAS)) {
113
+ client.admin().indices().delete(new DeleteIndexRequest(ES_TEST_ALIAS)).actionGet();
114
+ }
115
+
116
+ // Delete index
117
+ if (client.admin().cluster().state(new ClusterStateRequest()).actionGet().getState().getMetaData().hasIndex(ES_TEST_INDEX)) {
118
+ client.admin().indices().delete(new DeleteIndexRequest(ES_TEST_INDEX)).actionGet();
119
+ }
120
+
121
+ if (client.admin().cluster().state(new ClusterStateRequest()).actionGet().getState().getMetaData().hasIndex(ES_TEST_INDEX2)) {
122
+ client.admin().indices().delete(new DeleteIndexRequest(ES_TEST_INDEX2)).actionGet();
123
+ }
124
+ }
125
+ }
126
+
127
+ @Test
128
+ public void testDefaultValues()
129
+ {
130
+ ConfigSource config = config();
131
+ ElasticsearchOutputPlugin.PluginTask task = config.loadConfig(PluginTask.class);
132
+ assertEquals(ES_INDEX, task.getIndex());
133
+ }
134
+
135
+ @Test
136
+ public void testDefaultValuesNull()
137
+ {
138
+ ConfigSource config = Exec.newConfigSource()
139
+ .set("in", inputConfig())
140
+ .set("parser", parserConfig(schemaConfig()))
141
+ .set("type", "elasticsearch")
142
+ .set("mode", "") // NULL
143
+ .set("nodes", ES_NODES)
144
+ .set("cluster_name", ES_CLUSTER_NAME)
145
+ .set("index", ES_INDEX)
146
+ .set("index_type", ES_INDEX_TYPE)
147
+ .set("id", ES_ID)
148
+ .set("bulk_actions", ES_BULK_ACTIONS)
149
+ .set("bulk_size", ES_BULK_SIZE)
150
+ .set("concurrent_requests", ES_CONCURRENT_REQUESTS
151
+ );
152
+ Schema schema = config.getNested("parser").loadConfig(CsvParserPlugin.PluginTask.class).getSchemaConfig().toSchema();
153
+ try {
154
+ plugin.transaction(config, schema, 0, new OutputPlugin.Control()
155
+ {
156
+ @Override
157
+ public List<TaskReport> run(TaskSource taskSource)
158
+ {
159
+ return Lists.newArrayList(Exec.newTaskReport());
160
+ }
161
+ });
162
+ }
163
+ catch (Throwable t) {
164
+ if (t instanceof RuntimeException) {
165
+ assertTrue(t.getCause().getCause() instanceof ConfigException);
166
+ }
167
+ }
168
+ }
169
+
170
+ @Test
171
+ public void testTransaction()
172
+ {
173
+ ConfigSource config = config();
174
+ Schema schema = config.getNested("parser").loadConfig(CsvParserPlugin.PluginTask.class).getSchemaConfig().toSchema();
175
+ plugin.transaction(config, schema, 0, new OutputPlugin.Control()
176
+ {
177
+ @Override
178
+ public List<TaskReport> run(TaskSource taskSource)
179
+ {
180
+ return Lists.newArrayList(Exec.newTaskReport());
181
+ }
182
+ });
183
+ // no error happens
184
+ }
185
+
186
+ @Test
187
+ public void testResume()
188
+ {
189
+ ConfigSource config = config();
190
+ Schema schema = config.getNested("parser").loadConfig(CsvParserPlugin.PluginTask.class).getSchemaConfig().toSchema();
191
+ PluginTask task = config.loadConfig(PluginTask.class);
192
+ plugin.resume(task.dump(), schema, 0, new OutputPlugin.Control()
193
+ {
194
+ @Override
195
+ public List<TaskReport> run(TaskSource taskSource)
196
+ {
197
+ return Lists.newArrayList(Exec.newTaskReport());
198
+ }
199
+ });
200
+ }
201
+
202
+ @Test
203
+ public void testCleanup()
204
+ {
205
+ ConfigSource config = config();
206
+ Schema schema = config.getNested("parser").loadConfig(CsvParserPlugin.PluginTask.class).getSchemaConfig().toSchema();
207
+ PluginTask task = config.loadConfig(PluginTask.class);
208
+ plugin.cleanup(task.dump(), schema, 0, Arrays.asList(Exec.newTaskReport()));
209
+ // no error happens
210
+ }
211
+
212
+ @Test
213
+ public void testOutputByOpen()
214
+ throws GeneralSecurityException, IOException, NoSuchMethodException,
215
+ IllegalAccessException, InvocationTargetException, ParseException
216
+ {
217
+ ConfigSource config = config();
218
+ Schema schema = config.getNested("parser").loadConfig(CsvParserPlugin.PluginTask.class).getSchemaConfig().toSchema();
219
+ PluginTask task = config.loadConfig(PluginTask.class);
220
+ plugin.transaction(config, schema, 0, new OutputPlugin.Control() {
221
+ @Override
222
+ public List<TaskReport> run(TaskSource taskSource)
223
+ {
224
+ return Lists.newArrayList(Exec.newTaskReport());
225
+ }
226
+ });
227
+ TransactionalPageOutput output = plugin.open(task.dump(), schema, 0);
228
+
229
+ List<Page> pages = PageTestUtils.buildPage(runtime.getBufferAllocator(), schema, 1L, 32864L, Timestamp.ofEpochSecond(1422386629), Timestamp.ofEpochSecond(1422316800), true, 123.45, "embulk");
230
+ assertEquals(1, pages.size());
231
+ for (Page page : pages) {
232
+ output.add(page);
233
+ }
234
+
235
+ output.finish();
236
+ output.commit();
237
+
238
+ Method createClient = ElasticsearchOutputPlugin.class.getDeclaredMethod("createClient", PluginTask.class);
239
+ createClient.setAccessible(true);
240
+ try (Client client = (Client) createClient.invoke(plugin, task)) {
241
+ GetResponse response = client.prepareGet(ES_INDEX, ES_INDEX_TYPE, "1").execute().actionGet();
242
+ assertTrue(response.isExists());
243
+ if (response.isExists()) {
244
+ Map<String, Object> map = response.getSourceAsMap();
245
+ assertEquals(1, map.get("id"));
246
+ assertEquals(32864, map.get("account"));
247
+ assertEquals("2015-01-27T19:23:49.000Z", map.get("time"));
248
+ assertEquals("2015-01-27T00:00:00.000Z", map.get("purchase"));
249
+ assertEquals(true, map.get("flg"));
250
+ assertEquals(123.45, map.get("score"));
251
+ assertEquals("embulk", map.get("comment"));
252
+ }
253
+ }
254
+ }
255
+
256
+ @Test
257
+ public void testOpenAbort()
258
+ {
259
+ ConfigSource config = config();
260
+ Schema schema = config.getNested("parser").loadConfig(CsvParserPlugin.PluginTask.class).getSchemaConfig().toSchema();
261
+ PluginTask task = config.loadConfig(PluginTask.class);
262
+ TransactionalPageOutput output = plugin.open(task.dump(), schema, 0);
263
+ output.abort();
264
+ // no error happens.
265
+ }
266
+
267
+ @Test
268
+ public void testCreateClientThrowsException()
269
+ throws GeneralSecurityException, IOException, NoSuchMethodException,
270
+ IllegalAccessException, InvocationTargetException
271
+ {
272
+ ConfigSource config = Exec.newConfigSource()
273
+ .set("in", inputConfig())
274
+ .set("parser", parserConfig(schemaConfig()))
275
+ .set("type", "elasticsearch")
276
+ .set("mode", "replace")
277
+ .set("nodes", Arrays.asList(ImmutableMap.of("host", "unknown-host", "port", 9300)))
278
+ .set("cluster_name", ES_CLUSTER_NAME)
279
+ .set("index", ES_INDEX)
280
+ .set("index_type", ES_INDEX_TYPE)
281
+ .set("id", ES_ID)
282
+ .set("bulk_actions", ES_BULK_ACTIONS)
283
+ .set("bulk_size", ES_BULK_SIZE)
284
+ .set("concurrent_requests", ES_CONCURRENT_REQUESTS
285
+ );
286
+ PluginTask task = config.loadConfig(PluginTask.class);
287
+
288
+ Method createClient = ElasticsearchOutputPlugin.class.getDeclaredMethod("createClient", PluginTask.class);
289
+ createClient.setAccessible(true);
290
+ try (Client client = (Client) createClient.invoke(plugin, task)) {
291
+ // do nothing
292
+ } catch (Throwable t) {
293
+ if (t instanceof InvocationTargetException) {
294
+ assertTrue(t.getCause().getCause() instanceof UnknownHostException);
295
+ }
296
+ }
297
+ }
298
+
299
+ @Test
300
+ public void testMode()
301
+ {
302
+ assertEquals(2, ElasticsearchOutputPlugin.Mode.values().length);
303
+ assertEquals(ElasticsearchOutputPlugin.Mode.INSERT, ElasticsearchOutputPlugin.Mode.valueOf("INSERT"));
304
+ }
305
+
306
+ @Test(expected = ConfigException.class)
307
+ public void testModeThrowsConfigException()
308
+ {
309
+ ElasticsearchOutputPlugin.Mode.fromString("non-exists-mode");
310
+ }
311
+
312
+ @Test
313
+ public void testDeleteIndex()
314
+ throws GeneralSecurityException, IOException, NoSuchMethodException,
315
+ IllegalAccessException, InvocationTargetException
316
+ {
317
+ ConfigSource config = config();
318
+ PluginTask task = config.loadConfig(PluginTask.class);
319
+
320
+ Method createClient = ElasticsearchOutputPlugin.class.getDeclaredMethod("createClient", PluginTask.class);
321
+ createClient.setAccessible(true);
322
+ try (Client client = (Client) createClient.invoke(plugin, task)) {
323
+ // Create Index
324
+ client.admin().indices().create(new CreateIndexRequest(ES_TEST_INDEX)).actionGet();
325
+
326
+ Method deleteIndex = ElasticsearchOutputPlugin.class.getDeclaredMethod("deleteIndex", String.class, Client.class);
327
+ deleteIndex.setAccessible(true);
328
+ deleteIndex.invoke(plugin, ES_TEST_INDEX, client);
329
+
330
+ assertEquals(false, client.admin().cluster().state(new ClusterStateRequest()).actionGet().getState().getMetaData().hasIndex(ES_TEST_INDEX));
331
+ }
332
+ }
333
+
334
+ @Test
335
+ public void testAlias()
336
+ throws GeneralSecurityException, IOException, NoSuchMethodException,
337
+ IllegalAccessException, InvocationTargetException
338
+ {
339
+ ConfigSource config = config();
340
+ PluginTask task = config.loadConfig(PluginTask.class);
341
+
342
+ Method createClient = ElasticsearchOutputPlugin.class.getDeclaredMethod("createClient", PluginTask.class);
343
+ createClient.setAccessible(true);
344
+ try (Client client = (Client) createClient.invoke(plugin, task)) {
345
+ Method isAlias = ElasticsearchOutputPlugin.class.getDeclaredMethod("isAlias", String.class, Client.class);
346
+ isAlias.setAccessible(true);
347
+
348
+ Method isExistsAlias = ElasticsearchOutputPlugin.class.getDeclaredMethod("isExistsAlias", String.class, Client.class);
349
+ isExistsAlias.setAccessible(true);
350
+
351
+ Method getIndexByAlias = ElasticsearchOutputPlugin.class.getDeclaredMethod("getIndexByAlias", String.class, Client.class);
352
+ getIndexByAlias.setAccessible(true);
353
+
354
+ Method reAssignAlias = ElasticsearchOutputPlugin.class.getDeclaredMethod("reAssignAlias", String.class, String.class, Client.class);
355
+ reAssignAlias.setAccessible(true);
356
+
357
+ assertEquals(false, isAlias.invoke(plugin, ES_TEST_ALIAS, client));
358
+ assertEquals(false, isExistsAlias.invoke(plugin, ES_TEST_ALIAS, client));
359
+ List<String> indicesBefore = (List<String>) getIndexByAlias.invoke(plugin, ES_TEST_ALIAS, client);
360
+ assertEquals(0, indicesBefore.size());
361
+
362
+ // Create Index
363
+ client.admin().indices().create(new CreateIndexRequest(ES_TEST_INDEX)).actionGet();
364
+ client.admin().indices().create(new CreateIndexRequest(ES_TEST_INDEX2)).actionGet();
365
+ // Assign Alias
366
+ reAssignAlias.invoke(plugin, ES_TEST_ALIAS, ES_TEST_INDEX, client);
367
+
368
+ assertEquals(true, isAlias.invoke(plugin, ES_TEST_ALIAS, client));
369
+ assertEquals(true, isExistsAlias.invoke(plugin, ES_TEST_ALIAS, client));
370
+ List<String> indicesAfter = (List<String>) getIndexByAlias.invoke(plugin, ES_TEST_ALIAS, client);
371
+ assertEquals(1, indicesAfter.size());
372
+
373
+ // ReAssginAlias
374
+ reAssignAlias.invoke(plugin, ES_TEST_ALIAS, ES_TEST_INDEX2, client);
375
+ List<String> indicesReassign = (List<String>) getIndexByAlias.invoke(plugin, ES_TEST_ALIAS, client);
376
+ assertEquals(1, indicesReassign.size());
377
+ }
378
+ }
379
+
380
+ @Test
381
+ public void testIsExistsIndex()
382
+ throws GeneralSecurityException, IOException, NoSuchMethodException,
383
+ IllegalAccessException, InvocationTargetException
384
+ {
385
+ ConfigSource config = config();
386
+ PluginTask task = config.loadConfig(PluginTask.class);
387
+
388
+ Method createClient = ElasticsearchOutputPlugin.class.getDeclaredMethod("createClient", PluginTask.class);
389
+ createClient.setAccessible(true);
390
+ try (Client client = (Client) createClient.invoke(plugin, task)) {
391
+ Method isExistsIndex = ElasticsearchOutputPlugin.class.getDeclaredMethod("isExistsIndex", String.class, Client.class);
392
+ isExistsIndex.setAccessible(true);
393
+
394
+ // Delete index
395
+ if (client.admin().cluster().state(new ClusterStateRequest()).actionGet().getState().getMetaData().hasIndex(ES_TEST_INDEX)) {
396
+ client.admin().indices().delete(new DeleteIndexRequest(ES_TEST_INDEX)).actionGet();
397
+ }
398
+ assertEquals(false, isExistsIndex.invoke(plugin, ES_TEST_INDEX, client));
399
+
400
+ // Create Index
401
+ client.admin().indices().create(new CreateIndexRequest(ES_TEST_INDEX)).actionGet();
402
+ assertEquals(true, isExistsIndex.invoke(plugin, ES_TEST_INDEX, client));
403
+ }
404
+ }
405
+
406
+ @Test
407
+ public void testGenerateNewIndex()
408
+ {
409
+ String newIndexName = plugin.generateNewIndexName(ES_INDEX);
410
+ Timestamp time = Exec.getTransactionTime();
411
+ assertEquals(ES_INDEX + new SimpleDateFormat("_yyyyMMdd-HHmmss").format(time.toEpochMilli()), newIndexName);
412
+ }
413
+
414
+ private byte[] convertInputStreamToByte(InputStream is) throws IOException
415
+ {
416
+ ByteArrayOutputStream bo = new ByteArrayOutputStream();
417
+ byte [] buffer = new byte[1024];
418
+ while (true) {
419
+ int len = is.read(buffer);
420
+ if (len < 0) {
421
+ break;
422
+ }
423
+ bo.write(buffer, 0, len);
424
+ }
425
+ return bo.toByteArray();
426
+ }
427
+
428
+ private ConfigSource config()
429
+ {
430
+ return Exec.newConfigSource()
431
+ .set("in", inputConfig())
432
+ .set("parser", parserConfig(schemaConfig()))
433
+ .set("type", "elasticsearch")
434
+ .set("mode", "insert")
435
+ .set("nodes", ES_NODES)
436
+ .set("cluster_name", ES_CLUSTER_NAME)
437
+ .set("index", ES_INDEX)
438
+ .set("index_type", ES_INDEX_TYPE)
439
+ .set("id", ES_ID)
440
+ .set("bulk_actions", ES_BULK_ACTIONS)
441
+ .set("bulk_size", ES_BULK_SIZE)
442
+ .set("concurrent_requests", ES_CONCURRENT_REQUESTS);
443
+ }
444
+
445
+ private ImmutableMap<String, Object> inputConfig()
446
+ {
447
+ ImmutableMap.Builder<String, Object> builder = new ImmutableMap.Builder<>();
448
+ builder.put("type", "file");
449
+ builder.put("path_prefix", PATH_PREFIX);
450
+ builder.put("last_path", "");
451
+ return builder.build();
452
+ }
453
+
454
+ private ImmutableMap<String, Object> parserConfig(ImmutableList<Object> schemaConfig)
455
+ {
456
+ ImmutableMap.Builder<String, Object> builder = new ImmutableMap.Builder<>();
457
+ builder.put("type", "csv");
458
+ builder.put("newline", "CRLF");
459
+ builder.put("delimiter", ",");
460
+ builder.put("quote", "\"");
461
+ builder.put("escape", "\"");
462
+ builder.put("trim_if_not_quoted", false);
463
+ builder.put("skip_header_lines", 1);
464
+ builder.put("allow_extra_columns", false);
465
+ builder.put("allow_optional_columns", false);
466
+ builder.put("columns", schemaConfig);
467
+ return builder.build();
468
+ }
469
+
470
+ private ImmutableList<Object> schemaConfig()
471
+ {
472
+ ImmutableList.Builder<Object> builder = new ImmutableList.Builder<>();
473
+ builder.add(ImmutableMap.of("name", "id", "type", "long"));
474
+ builder.add(ImmutableMap.of("name", "account", "type", "long"));
475
+ builder.add(ImmutableMap.of("name", "time", "type", "timestamp", "format", "%Y-%m-%d %H:%M:%S"));
476
+ builder.add(ImmutableMap.of("name", "purchase", "type", "timestamp", "format", "%Y%m%d"));
477
+ builder.add(ImmutableMap.of("name", "flg", "type", "boolean"));
478
+ builder.add(ImmutableMap.of("name", "score", "type", "double"));
479
+ builder.add(ImmutableMap.of("name", "comment", "type", "string"));
480
+ return builder.build();
481
+ }
482
+ }