massive update

.gitignore

@@ -2,3 +2,4 @@
 *.swp
 __pycache__
 bookworm.egg-info
+books.db

bookworm/batch_unpacker.py

@@ -0,0 +1,94 @@
+import os
+import json
+import socket
+import sys
+import time
+import tempfile
+from threading import Thread
+from subprocess import check_output
+from bookworm import s3
+from bookworm.logger import log, setup_logger
+from bookworm.constants import UNPACKABLE_EXTENSIONS, REDIS
+import redis
+
+def should_unpack(fname):
+    fname = fname.lower()
+    return fname.endswith('rar') or fname.endswith('zip')
+
+def archive_contents(fd):
+    to_extract = {}
+    contents = check_output(['lsar', '-j', fd.name]).decode('utf-8')
+    contents = json.loads(contents)
+    
+    log.debug(contents['lsarContents'])
+    for f in contents['lsarContents']:
+        fname = f['XADFileName']
+        log.info('Found file %s in the archive', fname)
+        if any([extension in fname.lower() for extension in ['.txt'] + UNPACKABLE_EXTENSIONS]):
+            log.info('Extracting %s from the archive', fname)
+            to_extract[f['XADIndex']] = fname
+    return to_extract
+
+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_store(job_key, s3key, s3client, redis, meta):
+    redis.hset(job_key, REDIS.STEP_KEY, 'UNPACKING')
+    unpacked_files = unpack(s3key, s3client, redis, meta)
+    redis.hset(job_key, REDIS.STEP_KEY, 'UNPACK_DONE')
+    log.info('Done unpacking job %s', job_key)
+
+    for fname, data in unpacked_files:
+        log.info('Batch process file: %s', fname)
+        books = bookworm.parse.lines_to_dicts(data.decode('utf-8'))
+        bookworm.parse.insert_books(books)
+
+    #delete_raw_file(s3client, s3key, meta)
+    redis.delete(job_key)
+
+def unpack(s3key, s3client, redis, meta):
+    log.info('Got a request to unpack %s', s3key)
+    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)
+        fd.flush()
+
+        if not should_unpack(s3key):
+            log.info("Not unpacking %s", s3key)
+            return [(s3key, raw_file_contents)]
+
+        to_extract = archive_contents(fd)
+        if not to_extract:
+            log.error("Could not find any valid file")
+            return []
+
+        ret = []
+        for index, fname in to_extract.items():
+            log.info('Processing %s %s', index, fname)
+            file_contents = check_output(['unar', '-o', '-', '-i', fd.name, str(index)])
+            log.info('Got %d bytes', len(file_contents))
+            ret.append((fname, file_contents))
+        return ret
+
+
+def main():
+    r = redis.StrictRedis(host='localhost', port=6379)
+    setup_logger()
+    s3client = s3.client()
+    while True:
+        log.info('Waiting for message on %s', REDIS.Q_PROCESS_BATCH_FILE)
+        topic, message = r.blpop(REDIS.Q_PROCESS_BATCH_FILE)
+
+        log.info('got message: %s', message)
+        params = json.loads(message.decode('utf-8'))
+        log.info('params for unpacker: %s', params)
+        params['s3client'] = s3client
+        params['redis'] = r
+
+        t = Thread(target=unpack_and_convert, kwargs=params)
+        t.daemon = True
+        t.start()
+
+main()

bookworm/constants.py

@@ -1,10 +1,19 @@
-RAW_FILE_BUCKET = 'rawfiles'
-IRC_TIME_TO_FIRST_COMMAND = 30
-IRC_CHANNEL = "#ebooks"
-PROCESSED_FILE_BUCKET = 'files'
+from enum import Enum
 UNPACKABLE_EXTENSIONS = ['epub', 'mobi', 'azw3']
-REDIS_BOOK_COMMANDS = 'BOOK_COMMANDS'
-REDIS_FETCH_FILE = 'FETCH_FILE'
-REDIS_UNPACK_FILE = 'UNPACK_FILE'
-REDIS_STEP_KEY = 'STEP'
-REDIS_STATE_KEY = 'STATE'
+JOB_KEY_PREFIX = 'job_'
+
+class REDIS():
+    Q_BOOK_COMMANDS = 'BOOK_COMMANDS'
+    Q_FETCH_FILE = 'FETCH_FILE'
+    Q_UNPACK_FILE = 'UNPACK_FILE'
+    Q_PROCESS_BATCH_FILE = 'PROCESS_BATCH_FILE'
+    STEP_KEY = 'STEP'
+    STATE_KEY = 'STATE'
+
+class IRC():
+    TIME_TO_FIRST_COMMAND = 32
+    CHANNEL = "#ebooks"
+
+class BUCKET():
+    RAW_FILE = 'rawfiles'
+    PROCESSED_FILE = 'files'

bookworm/file_fetcher.py

@@ -3,7 +3,7 @@ import socket
 import sys
 import time
 from threading import Thread
-from bookworm.constants import REDIS_FETCH_FILE, REDIS_STATE_KEY, REDIS_STEP_KEY
+from bookworm.constants import REDIS
 from bookworm.logger import log, setup_logger
 from bookworm import s3
 
@@ -16,7 +16,7 @@ def netcat(filename, ip, port, size, job_key, s3client, redis, meta):
     s.connect((ip, port))
     log.info("Receiving file %s", filename)
 
-    redis.hset(job_key, REDIS_STEP_KEY, 'DOWNLOADING')
+    redis.hset(job_key, REDIS.STEP_KEY, 'DOWNLOADING')
     buff = b''
     count = 0
     last_perc = 0
@@ -30,14 +30,14 @@ def netcat(filename, ip, port, size, job_key, s3client, redis, meta):
         perc = int(100 * count / size)
         if perc % 5 == 0 and perc != last_perc:
             log.info("Download percentage: %d", perc)
-            redis.hset(job_key, REDIS_STATE_KEY, str(perc))
+            redis.hset(job_key, REDIS.STATE_KEY, str(perc))
             last_perc = perc
         if count >= size:
             break
     log.info("Download complete")
     s.close()
     log.info("Putting file in s3")
-    redis.hset(job_key, REDIS_STATE_KEY, '100')
+    redis.hset(job_key, REDIS.STATE_KEY, '100')
     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(meta['unpack_file_queue'], json.dumps({'job_key': job_key, 's3key': filename, 'meta': meta}))
@@ -47,8 +47,8 @@ def main():
     setup_logger()
 
     while True:
-        log.info('Waiting for message on %s', REDIS_FETCH_FILE)
-        topic, message = r.blpop(REDIS_FETCH_FILE)
+        log.info('Waiting for message on %s', REDIS.Q_FETCH_FILE)
+        topic, message = r.blpop(REDIS.Q_FETCH_FILE)
 
         log.info('got message: %s', message)
         params = json.loads(message.decode('utf-8'))

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_STEP_KEY
+from bookworm.constants import IRC, REDIS, JOB_KEY_PREFIX
 from bookworm.logger import log, setup_logger
 from bookworm.utils import random_hash
 
@@ -36,13 +36,13 @@ class IRCClient(irc.client.SimpleIRCClient):
 
     def wait_for_commands(self):
         while True:
-            log.info('Waiting for message on %s', REDIS_BOOK_COMMANDS)
-            topic, message = self.r.blpop(REDIS_BOOK_COMMANDS)
+            log.info('Waiting for message on %s', REDIS.Q_BOOK_COMMANDS)
+            topic, message = self.r.blpop(REDIS.Q_BOOK_COMMANDS)
             log.info('got message: %s', message)
 
-            delta = IRC_TIME_TO_FIRST_COMMAND - (time.time() - self.startup)
+            delta = IRC.TIME_TO_FIRST_COMMAND - (time.time() - self.startup)
             while delta > 0:
-                delta = IRC_TIME_TO_FIRST_COMMAND - (time.time() - self.startup)
+                delta = IRC.TIME_TO_FIRST_COMMAND - (time.time() - self.startup)
                 log.info("I am not ready yet, still %d to go", delta)
                 time.sleep(max(min(delta, 2), 0))
 
@@ -52,13 +52,12 @@ class IRCClient(irc.client.SimpleIRCClient):
             log.info('Command: %s', command)
 
             self.meta = command['meta']
-            bot = command['bot'].strip()
-            book = command['book'].strip()
-            self.job_key = 'book_'+book
+            irc_command = command['command'].strip()
+            self.job_key = JOB_KEY_PREFIX + irc_command
             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}')
+            self.r.hset(self.job_key, REDIS.STEP_KEY, 'REQUESTED')
+            self.connection.privmsg(self.target, irc_command)
 
     def on_pubmsg(self, connection, event):
         log.debug('pubmsg %s', event)
@@ -107,6 +106,6 @@ def main():
     setup_logger()
     name = "bookbot" + random_hash()
 
-    c = IRCClient(IRC_CHANNEL, name)
+    c = IRCClient(IRC.CHANNEL, name)
     c.wait_for_commands()
 main()

bookworm/logger.py

@@ -1,7 +1,7 @@
 import sys
 import logging
 
-FORMAT = '%(asctime)s - %(name)s - %(levelname)s - %(message)s'
+FORMAT = '%(asctime)s - %(threadName)s - %(levelname)s - %(message)s'
 
 log = logging.getLogger(__name__)
 def setup_logger():

bookworm/parse.py

@@ -1,44 +1,42 @@
-import time
 import re
-from collections import namedtuple
-import pg_simple
+import sqlite3
+import time
 
 ONE_GB = 2**30
 ONE_MB = 2**20
 ONE_KB = 2**10
 
-connection_pool = pg_simple.config_pool(dsn='dbname=david user=david')
 
 r = re.compile(r'^!(?P<bot>[a-z0-9-]+?) (?P<book>.+)\s(::INFO::|-+)\s+(?P<size>[0-9.]+\s*[BKMG]+)\s*$', re.I)
-tags_re = re.compile(r'[\[\(](?P<tag>.*?)[\]\)]')
-books = []
-Book = namedtuple('Book', 'raw bot book size')
-
-for line in open('all_books_2.txt', 'r'):
-    line = line.strip()
-    match = r.match(line)
-    bot = match.group('bot').strip()
-    book = match.group('book').strip()
-    size = match.group('size').strip().lower()
-
-    if 'gb' in size:
-        size = float(size.replace('gb', '')) * ONE_GB
-    elif 'mb' in size:
-        size = float(size.replace('mb', '')) * ONE_MB
-    elif 'kb' in size:
-        size = float(size.replace('kb', '')) * ONE_KB
-    elif 'b' in size:
-        size = float(size.replace('b', ''))
-
-    books.append(Book(line, bot, book, size)._asdict())
-
-print('db..')
-with pg_simple.PgSimple(connection_pool) as db:
-    t = time.time()
-    vals = [tuple(book.values()) for book in books]
-    print('making tuples', time.time() - t)
-
-    db.insert_many('books', keys=books[0].keys(), values=vals, page_size=1000)
-    print('commiting..')
+# tags_re = re.compile(r'[\[\(](?P<tag>.*?)[\]\)]')
+
+def lines_to_dicts(lines):
+    books = []
+
+    for line in lines:
+        line = line.strip()
+        match = r.match(line)
+        bot = match.group('bot').strip()
+        book = match.group('book').strip()
+        size = match.group('size').strip().lower()
+    
+        if 'gb' in size:
+            size = int(float(size.replace('gb', '')) * ONE_GB)
+        elif 'mb' in size:
+            size = int(float(size.replace('mb', '')) * ONE_MB)
+        elif 'kb' in size:
+            size = int(float(size.replace('kb', '')) * ONE_KB)
+        elif 'b' in size:
+            size = int(float(size.replace('b', '')))
+    
+        books.append((bot, book, size))
+    return books
+    
+def insert_books(books):
+    db = sqlite3.connect('books.db')
+    c = db.cursor()
+    entries = []
+    c.executemany('INSERT OR IGNORE INTO books(bot, book, size) values (?,?,?)',
+            books)
     db.commit()
-    print('done..')
+    db.close()

bookworm/s3.py

@@ -6,8 +6,8 @@ def client():
     session = boto3.session.Session()
     return session.client(
         service_name='s3',
-        aws_access_key_id='3207900NM6AZ01AN02O1',
-        aws_secret_access_key='pmWuVRye20yiPbRk5tau1L6ggSueeaU5KXh2n0aZ',
+        aws_access_key_id='LBSBQ46DB9S160XIW438',
+        aws_secret_access_key='7nVT+TwGzPoC02AIIY4+1fVNdj07RQt+ntd55S+I',
         endpoint_url='http://localhost:9000',
         config=config
     )

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 UNPACKABLE_EXTENSIONS, REDIS_UNPACK_FILE, REDIS_STATE_KEY, REDIS_STEP_KEY
+from bookworm.constants import UNPACKABLE_EXTENSIONS, REDIS
 import redis
 
 def should_unpack(fname):
@@ -46,15 +46,15 @@ def delete_raw_file(s3client, s3key, meta):
     s3client.delete_object(Bucket=meta['raw_file_bucket'], Key=s3key)
 
 def unpack_and_convert(job_key, s3key, s3client, redis, meta):
-    redis.hset(job_key, REDIS_STEP_KEY, 'UNPACKING')
+    redis.hset(job_key, REDIS.STEP_KEY, 'UNPACKING')
     unpacked_files = unpack(s3key, s3client, redis, meta)
-    redis.hset(job_key, REDIS_STEP_KEY, 'UNPACK_DONE')
+    redis.hset(job_key, REDIS.STEP_KEY, 'UNPACK_DONE')
     log.info('Done unpacking job %s', job_key)
 
     for fname, data in unpacked_files:
         if not fname.lower().endswith('mobi'):
-            redis.hset(job_key, REDIS_STEP_KEY, 'CONVERTING')
-            redis.hset(job_key, REDIS_STATE_KEY, fname)
+            redis.hset(job_key, REDIS.STEP_KEY, 'CONVERTING')
+            redis.hset(job_key, REDIS.STATE_KEY, fname)
             log.info('Asked to store %s, need to convert first', fname)
             fname_no_ext, ext = os.path.splitext(fname)
             with tempfile.NamedTemporaryFile(suffix=ext) as original:
@@ -82,7 +82,6 @@ def unpack(s3key, s3client, redis, meta):
 
         to_extract = archive_contents(fd)
         if not to_extract:
-            log.info(contents['lsarContents'])
             log.error("Could not find any valid file")
             return []
 
@@ -100,8 +99,8 @@ def main():
     setup_logger()
     s3client = s3.client()
     while True:
-        log.info('Waiting for message on %s', REDIS_UNPACK_FILE)
-        topic, message = r.blpop(REDIS_UNPACK_FILE)
+        log.info('Waiting for message on %s', REDIS.Q_UNPACK_FILE)
+        topic, message = r.blpop(REDIS.Q_UNPACK_FILE)
 
         log.info('got message: %s', message)
         params = json.loads(message.decode('utf-8'))

docs/services.dot

@@ -6,9 +6,11 @@ digraph Services {
    Unpacker
    WebAPI
    BulkInserter
-   S3
    DB
 
-   WebAPI -> IRCClient -> FileFetcher -> Unpacker -> S3
-   Unpacker -> BulkInserter -> DB
+   WebAPI -> IRCClient -> FileFetcher -> Unpacker
+   FileFetcher -> BulkInserter
+
+   WebAPI -> DB[label="Search" style=dashed];
+   BulkInserter -> DB[label="Save cache" style=dashed];
 }

services/bookworm@.service

@@ -0,0 +1,11 @@
+[Unit]
+Description=Bookworm %i service
+
+[Service]
+User=david
+Type=simple
+ExecStart=/home/david/git/irc-books/.venv/bin/%i
+Nice=5
+
+[Install]
+WantedBy=multi-user.target