file elasticsearch index worker

3 files changed, 63 insertions, 1 deletions

diff --git a/python/fatcat_tools/workers/__init__.py b/python/fatcat_tools/workers/__init__.py
index 18b35b2b..d3c0978f 100644
--- a/python/fatcat_tools/workers/__init__.py
+++ b/python/fatcat_tools/workers/__init__.py
@@ -2,6 +2,7 @@ from .changelog import ChangelogWorker, EntityUpdatesWorker
 from .elasticsearch import (
     ElasticsearchChangelogWorker,
     ElasticsearchContainerWorker,
+    ElasticsearchFileWorker,
     ElasticsearchReleaseWorker,
 )
 from .worker_common import FatcatWorker, most_recent_message
diff --git a/python/fatcat_tools/workers/elasticsearch.py b/python/fatcat_tools/workers/elasticsearch.py
index 71c4dcf6..21e1f41b 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
@@ -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.
diff --git a/python/fatcat_worker.py b/python/fatcat_worker.py
index 0e2d03b5..67ea8451 100755
--- a/python/fatcat_worker.py
+++ b/python/fatcat_worker.py
@@ -10,6 +10,7 @@ from fatcat_tools.workers import (
     ChangelogWorker,
     ElasticsearchChangelogWorker,
     ElasticsearchContainerWorker,
+    ElasticsearchFileWorker,
     ElasticsearchReleaseWorker,
     EntityUpdatesWorker,
 )
@@ -70,6 +71,17 @@ def run_elasticsearch_container(args: argparse.Namespace) -> None:
     worker.run()
 
 
+def run_elasticsearch_file(args: argparse.Namespace) -> None:
+    consume_topic = "fatcat-{}.file-updates".format(args.env)
+    worker = ElasticsearchFileWorker(
+        args.kafka_hosts,
+        consume_topic,
+        elasticsearch_backend=args.elasticsearch_backend,
+        elasticsearch_index=args.elasticsearch_index,
+    )
+    worker.run()
+
+
 def run_elasticsearch_changelog(args: argparse.Namespace) -> None:
     consume_topic = "fatcat-{}.changelog".format(args.env)
     worker = ElasticsearchChangelogWorker(
@@ -150,6 +162,22 @@ def main() -> None:
         help="whether to query release search index for container stats",
     )
 
+    sub_elasticsearch_file = subparsers.add_parser(
+        "elasticsearch-file",
+        help="consume kafka feed of new/updated files, transform and push to search",
+    )
+    sub_elasticsearch_file.set_defaults(func=run_elasticsearch_file)
+    sub_elasticsearch_file.add_argument(
+        "--elasticsearch-backend",
+        help="elasticsearch backend to connect to",
+        default="http://localhost:9200",
+    )
+    sub_elasticsearch_file.add_argument(
+        "--elasticsearch-index",
+        help="elasticsearch index to push into",
+        default="fatcat_file",
+    )
+
     sub_elasticsearch_changelog = subparsers.add_parser(
         "elasticsearch-changelog",
         help="consume changelog kafka feed, transform and push to search",