1 files changed, 36 insertions, 3 deletions
diff --git a/python/fatcat_tools/workers/elasticsearch.py b/python/fatcat_tools/workers/elasticsearch.py
index 71c4dcf6..bfadea64 100644
--- a/python/fatcat_tools/workers/elasticsearch.py
+++ b/python/fatcat_tools/workers/elasticsearch.py
@@ -5,12 +5,19 @@ from typing import Any, Callable, List, Optional
 import elasticsearch
 import requests
 from confluent_kafka import Consumer, KafkaException
-from fatcat_openapi_client import ApiClient, ChangelogEntry, ContainerEntity, ReleaseEntity
+from fatcat_openapi_client import (
+    ApiClient,
+    ChangelogEntry,
+    ContainerEntity,
+    FileEntity,
+    ReleaseEntity,
+)
 
 from fatcat_tools import entity_from_json, public_api
 from fatcat_tools.transforms import (
     changelog_to_elasticsearch,
     container_to_elasticsearch,
+    file_to_elasticsearch,
     release_to_elasticsearch,
 )
 from fatcat_web.search import get_elastic_container_stats
@@ -34,7 +41,7 @@ class ElasticsearchReleaseWorker(FatcatWorker):
         offset: Optional[int] = None,
         elasticsearch_backend: str = "http://localhost:9200",
         elasticsearch_index: str = "fatcat",
-        elasticsearch_release_index: str = "fatcat_releases",
+        elasticsearch_release_index: str = "fatcat_release",
         batch_size: int = 200,
         api_host: str = "https://api.fatcat.wiki/v0",
         query_stats: bool = False,
@@ -213,7 +220,7 @@ class ElasticsearchContainerWorker(ElasticsearchReleaseWorker):
         query_stats: bool = False,
         elasticsearch_release_index: str = "fatcat_release",
         elasticsearch_backend: str = "http://localhost:9200",
-        elasticsearch_index: str = "fatcat",
+        elasticsearch_index: str = "fatcat_container",
         batch_size: int = 200,
     ):
         super().__init__(
@@ -233,6 +240,32 @@ class ElasticsearchContainerWorker(ElasticsearchReleaseWorker):
         self.transform_func = container_to_elasticsearch
 
 
+class ElasticsearchFileWorker(ElasticsearchReleaseWorker):
+    def __init__(
+        self,
+        kafka_hosts: str,
+        consume_topic: str,
+        poll_interval: float = 10.0,
+        offset: Optional[int] = None,
+        elasticsearch_backend: str = "http://localhost:9200",
+        elasticsearch_index: str = "fatcat_file",
+        batch_size: int = 200,
+    ):
+        super().__init__(
+            kafka_hosts=kafka_hosts,
+            consume_topic=consume_topic,
+            poll_interval=poll_interval,
+            offset=offset,
+            elasticsearch_backend=elasticsearch_backend,
+            elasticsearch_index=elasticsearch_index,
+            batch_size=batch_size,
+        )
+        # previous group got corrupted (by pykafka library?)
+        self.consumer_group = "elasticsearch-updates3"
+        self.entity_type = FileEntity
+        self.transform_func = file_to_elasticsearch
+
+
 class ElasticsearchChangelogWorker(ElasticsearchReleaseWorker):
     """
     Pulls changelog messages from Kafka, runs transformations and indexes them.