ソースを参照

pass queues and buckets as metadata instead of hardcoding

david 7 年 前
コミット
7bd6b3f8dc
4 ファイル変更41 行追加27 行削除
  1. 5 5
      bookworm/file_fetcher.py
  2. 11 6
      bookworm/ircclient.py
  3. 15 15
      bookworm/unpacker.py
  4. 10 1
      bookworm/web.py

+ 5 - 5
bookworm/file_fetcher.py

@@ -3,13 +3,13 @@ import socket
 import sys
 import time
 from threading import Thread
-from bookworm.constants import RAW_FILE_BUCKET, REDIS_UNPACK_FILE, REDIS_FETCH_FILE, REDIS_STATE_KEY, REDIS_STEP_KEY
+from bookworm.constants import REDIS_FETCH_FILE, REDIS_STATE_KEY, REDIS_STEP_KEY
 from bookworm.logger import log, setup_logger
 from bookworm import s3
 
 import redis
 
-def netcat(filename, ip, port, size, job_key, s3client, redis):
+def netcat(filename, ip, port, size, job_key, s3client, redis, meta):
     log.info('netcat: %s %d %d %s', ip, port, size, filename)
     s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
     log.info("Fetching %s", filename)
@@ -38,16 +38,16 @@ def netcat(filename, ip, port, size, job_key, s3client, redis):
     s.close()
     log.info("Putting file in s3")
     redis.hset(job_key, REDIS_STATE_KEY, '100')
-    s3client.put_object(Body=buff, Bucket=RAW_FILE_BUCKET, Key=filename)
+    s3client.put_object(Body=buff, Bucket=meta['raw_file_bucket'], Key=filename)
     log.info("File %s in s3 with key %s", filename, job_key)
-    redis.rpush(REDIS_UNPACK_FILE, json.dumps({'job_key': job_key, 's3key': filename}))
+    redis.rpush(meta['unpack_file_queue'], json.dumps({'job_key': job_key, 's3key': filename, 'meta': meta}))
 
 def main():
     r = redis.StrictRedis(host='localhost', port=6379)
     setup_logger()
 
     while True:
-        log.info('Waiting for message...')
+        log.info('Waiting for message on %s', REDIS_FETCH_FILE)
         topic, message = r.blpop(REDIS_FETCH_FILE)
 
         log.info('got message: %s', message)

+ 11 - 6
bookworm/ircclient.py

@@ -5,7 +5,7 @@ import time
 import shlex
 
 from threading import Thread
-from bookworm.constants import IRC_TIME_TO_FIRST_COMMAND, IRC_CHANNEL, REDIS_BOOK_COMMANDS, REDIS_FETCH_FILE, REDIS_STEP_KEY
+from bookworm.constants import IRC_TIME_TO_FIRST_COMMAND, IRC_CHANNEL, REDIS_BOOK_COMMANDS, REDIS_STEP_KEY
 from bookworm.logger import log, setup_logger
 from bookworm.utils import random_hash
 
@@ -36,7 +36,7 @@ class IRCClient(irc.client.SimpleIRCClient):
 
     def wait_for_commands(self):
         while True:
-            log.info('Waiting for message...')
+            log.info('Waiting for message on %s', REDIS_BOOK_COMMANDS)
             topic, message = self.r.blpop(REDIS_BOOK_COMMANDS)
             log.info('got message: %s', message)
 
@@ -51,9 +51,13 @@ class IRCClient(irc.client.SimpleIRCClient):
             command = json.loads(data)
             log.info('Command: %s', command)
 
+            self.meta = command['meta']
             bot = command['bot'].strip()
             book = command['book'].strip()
-            self.r.hset('book_'+book, REDIS_STEP_KEY, 'REQUESTED')
+            self.job_key = 'book_'+book
+            self.fetch_queue = command['meta']['fetch_file_queue']
+
+            self.r.hset(self.job_key, REDIS_STEP_KEY, 'REQUESTED')
             self.connection.privmsg(self.target, f'!{bot} {book}')
 
     def on_pubmsg(self, connection, event):
@@ -89,9 +93,10 @@ class IRCClient(irc.client.SimpleIRCClient):
                            'port': peer_port,
                            'size': int(size),
                            'filename': filename,
-                           "job_key": job_key})
-        log.info('Publishing to FETCH_FILE: %s', data)
-        self.r.rpush(REDIS_FETCH_FILE, data)
+                           "job_key": job_key,
+                           "meta": self.meta})
+        log.info('Publishing to %s: %s', self.fetch_queue, data)
+        self.r.rpush(self.fetch_queue, data)
 
     def on_disconnect(self, connection, event):
         sys.exit(0)

+ 15 - 15
bookworm/unpacker.py

@@ -8,7 +8,7 @@ from threading import Thread
 from subprocess import check_output
 from bookworm import s3
 from bookworm.logger import log, setup_logger
-from bookworm.constants import RAW_FILE_BUCKET, PROCESSED_FILE_BUCKET, UNPACKABLE_EXTENSIONS, REDIS_UNPACK_FILE, REDIS_STATE_KEY, REDIS_STEP_KEY
+from bookworm.constants import UNPACKABLE_EXTENSIONS, REDIS_UNPACK_FILE, REDIS_STATE_KEY, REDIS_STEP_KEY
 import redis
 
 def should_unpack(fname):
@@ -36,18 +36,18 @@ def convert_to_mobi(orig_fname) -> bytes:
         buff = open(fd.name, 'rb').read()
         return buff
 
-def store_file(s3client, fname, file_contents):
-    log.info('Puttin in s3 under bucket %s with key %s', PROCESSED_FILE_BUCKET, fname)
-    s3client.put_object(Body=file_contents, Bucket=PROCESSED_FILE_BUCKET, Key=fname)
-    log.info('Put in s3 under bucket %s with key %s', PROCESSED_FILE_BUCKET, fname)
+def store_file(s3client, fname, file_contents, meta):
+    log.info('Puttin in s3 under bucket %s with key %s', meta['processed_file_bucket'], fname)
+    s3client.put_object(Body=file_contents, Bucket=meta['processed_file_bucket'], Key=fname)
+    log.info('Put in s3 under bucket %s with key %s', meta['processed_file_bucket'], fname)
 
-def delete_raw_file(s3client, s3key):
-    log.info('Deleting %s from %s', s3key, RAW_FILE_BUCKET)
-    s3client.delete_object(Bucket=RAW_FILE_BUCKET, Key=s3key)
+def delete_raw_file(s3client, s3key, meta):
+    log.info('Deleting %s from %s', s3key, meta['raw_file_bucket'])
+    s3client.delete_object(Bucket=meta['raw_file_bucket'], Key=s3key)
 
-def unpack_and_convert(job_key, s3key, s3client, redis):
+def unpack_and_convert(job_key, s3key, s3client, redis, meta):
     redis.hset(job_key, REDIS_STEP_KEY, 'UNPACKING')
-    unpacked_files = unpack(s3key, s3client, redis)
+    unpacked_files = unpack(s3key, s3client, redis, meta)
     redis.hset(job_key, REDIS_STEP_KEY, 'UNPACK_DONE')
     log.info('Done unpacking job %s', job_key)
 
@@ -63,14 +63,14 @@ def unpack_and_convert(job_key, s3key, s3client, redis):
                 converted_data = convert_to_mobi(original.name)
                 data = converted_data
             fname = fname_no_ext + '.mobi'
-        store_file(s3client, fname, data)
+        store_file(s3client, fname, data, meta)
 
-    delete_raw_file(s3client, s3key)
+    delete_raw_file(s3client, s3key, meta)
     redis.delete(job_key)
 
-def unpack(s3key, s3client, redis):
+def unpack(s3key, s3client, redis, meta):
     log.info('Got a request to unpack %s', s3key)
-    data = s3client.get_object(Key=s3key, Bucket=RAW_FILE_BUCKET)
+    data = s3client.get_object(Key=s3key, Bucket=meta['raw_file_bucket'])
     with tempfile.NamedTemporaryFile() as fd:
         raw_file_contents = data['Body'].read()
         fd.write(raw_file_contents)
@@ -100,7 +100,7 @@ def main():
     setup_logger()
     s3client = s3.client()
     while True:
-        log.info('Waiting for message...')
+        log.info('Waiting for message on %s', REDIS_UNPACK_FILE)
         topic, message = r.blpop(REDIS_UNPACK_FILE)
 
         log.info('got message: %s', message)

+ 10 - 1
bookworm/web.py

@@ -66,7 +66,16 @@ def search_books():
 @app.route('/book/fetch', methods=['POST'])
 def fetch_books():
     fetch = request.json
-    r.rpush(constants.REDIS_BOOK_COMMANDS, json.dumps({'bot': fetch['bot'], 'book': fetch['book']}))
+    data = {'bot': fetch['bot'],
+            'book': fetch['book'],
+            'meta': {
+                'unpack_file_queue': constants.REDIS_UNPACK_FILE,
+                'fetch_file_queue': constants.REDIS_FETCH_FILE,
+                'raw_file_bucket': constants.RAW_FILE_BUCKET,
+                'processed_file_bucket': constants.PROCESSED_FILE_BUCKET,
+            }
+           }
+    r.rpush(constants.REDIS_BOOK_COMMANDS, json.dumps(data))
     return ''