dump&restore script

DevAlone · DevAlone · commit 2f30c3023bc8 · 2018-04-08T18:31:53.000+03:00
diff --git a/collectors_list.py b/collectors_list.py
@@ -1,28 +1,20 @@
-import asyncio
-
-import peewee
-
 from models import db, CollectorState
 from proxy_py import settings
 
 import os
+import asyncio
 import importlib.util
 
 
 collectors = {}
 
 
-print("collectors dir: {}".format(settings.COLLECTORS_DIR))
 for root, dirs, files in os.walk(settings.COLLECTORS_DIR):
     for file in files:
         if file.endswith(".py"):
-            print("processing file: {}".format(file))
-
             file_path = os.path.join(root, file)
             module_name = os.path.splitext(file_path)[0].replace('/', '.')
-            print("module name: {}".format(module_name))
             spec = importlib.util.spec_from_file_location(module_name, file_path)
-
             collector_module = importlib.util.module_from_spec(spec)
             spec.loader.exec_module(collector_module)
 
@@ -61,7 +53,8 @@ def get_collector_state(module_name: str):
 def get_collector_of_module_name(module_name: str):
     if module_name not in collectors:
         raise CollectorNotFoundException(
-            "Probably some collector exists in database but not in filesystem. module_name = {}".format(module_name)
+            "Probably some collector exists in database but not in filesystem. "
+            "module_name = {}".format(module_name)
         )
 
     return collectors[module_name]
diff --git a/dump_db.py b/dump_db.py
@@ -0,0 +1,60 @@
+from proxy_py import settings
+
+import json
+import os
+import asyncio
+import aiopg
+import sys
+
+
+dsn = 'dbname={} user={} password={} host=127.0.0.1'.format(
+    settings.DATABASE_CONNECTION_KWARGS['database'],
+    settings.DATABASE_CONNECTION_KWARGS['user'],
+    settings.DATABASE_CONNECTION_KWARGS['password'],
+)
+
+
+async def dump_it(table_name, save_dir):
+    print("start dumping table {}".format(table_name))
+
+    pool = await aiopg.create_pool(dsn)
+    async with pool.acquire() as conn:
+        async with conn.cursor() as cur:
+            await cur.execute("SELECT * from {}".format(table_name))
+            # ret = []
+            column_names = [desc[0] for desc in cur.description]
+            print("found columns: {}".format(column_names))
+
+            file_path = os.path.join(save_dir, table_name + ".json")
+            print("writing to file {}".format(file_path))
+            with open(file_path, 'w') as file:
+                file.write("[\n")
+                async for row in cur:
+                    json_obj = {
+                        column_name: row[i]
+                        for i, column_name in enumerate(column_names)
+                    }
+
+                    file.write("{},\n".format(json.dumps(json_obj)))
+
+                file.write("]\n")
+
+    print()
+
+
+async def main():
+    if len(sys.argv) < 2:
+        raise BaseException("Please, specify directory where to save your dump")
+
+    save_to_dir = sys.argv[1]
+    if not os.path.isdir(save_to_dir):
+        raise BaseException("{} is not a directory".format(save_to_dir))
+
+    await dump_it('proxies', save_to_dir)
+    await dump_it('proxy_count_items', save_to_dir)
+    await dump_it('collector_states', save_to_dir)
+
+
+if __name__ == '__main__':
+    loop = asyncio.get_event_loop()
+    loop.run_until_complete(main())
diff --git a/fill_db.py b/fill_db.py
@@ -0,0 +1,76 @@
+from proxy_py import settings
+
+import json
+import os
+import asyncio
+import aiopg
+import sys
+
+
+dsn = 'dbname={} user={} password={} host=127.0.0.1'.format(
+    settings.DATABASE_CONNECTION_KWARGS['database'],
+    settings.DATABASE_CONNECTION_KWARGS['user'],
+    settings.DATABASE_CONNECTION_KWARGS['password'],
+)
+
+
+async def fill_it(dir_path, table_name):
+    print("start filling table {}".format(table_name))
+
+    pool = await aiopg.create_pool(dsn)
+    async with pool.acquire() as conn:
+        async with conn.cursor() as cur:
+            # await cur.execute()
+
+            file_path = os.path.join(dir_path, table_name + ".json")
+
+            print("reading file {}".format(file_path))
+            with open(file_path) as file:
+                print("filling database", end="")
+                for i, line in enumerate(file):
+                    print(".", end="")
+                    sys.stdout.flush()
+
+                    line = line.strip()
+                    if line and line not in '[]':
+                        line = line[:-1]
+                        json_obj = json.loads(line)
+                        # print(json_obj)
+                        keys = []
+                        values = []
+                        for key, val in json_obj.items():
+                            keys.append(key)
+                            values.append(val)
+
+                        sql_request = """
+                            INSERT INTO {} ({})
+                            VALUES ({})
+                            ON CONFLICT DO NOTHING; 
+                        """.format(table_name, ', '.join(keys), ', '.join(['%s' for _ in range(len(values))]))
+
+                        # print(sql_request.strip())
+
+                        await cur.execute(sql_request, values)
+
+                print()
+
+    print()
+
+
+async def main():
+    if len(sys.argv) < 2:
+        raise BaseException("Please, specify directory from which to load your dumps")
+
+    load_from_dir = sys.argv[1]
+    if not os.path.isdir(load_from_dir):
+        raise BaseException("{} is not a directory".format(load_from_dir))
+
+    for filename in [filename for filename in os.listdir(load_from_dir)
+                     if os.path.isfile(os.path.join(load_from_dir, filename))]:
+        table_name, _ = filename.split(".")
+        await fill_it(load_from_dir, table_name)
+
+
+if __name__ == '__main__':
+    loop = asyncio.get_event_loop()
+    loop.run_until_complete(main())
diff --git a/main.py b/main.py
@@ -3,7 +3,6 @@
 from proxy_py import settings
 from processor import Processor
 from server.proxy_provider_server import ProxyProviderServer
-
 from models import Proxy, ProxyCountItem, db
 
 import asyncio
diff --git a/processor.py b/processor.py
@@ -4,7 +4,6 @@
 
 import collectors_list
 import proxy_utils
-
 import asyncio
 import time
 import re
@@ -26,7 +25,6 @@ class Processor:
     """
 
     def __init__(self):
-        # TODO: find better logger
         self.logger = logging.getLogger("proxy_py/processor")
 
         if settings.DEBUG:
@@ -66,17 +64,23 @@ async def consumer(self):
         tasks = []
         while True:
             await asyncio.sleep(0.1)
-            i = 0
+            try:
+                i = 0
 
-            while not self.queue.empty() and i <= settings.CONCURRENT_TASKS_COUNT:
-                proxy_data = self.queue.get_nowait()
-                tasks.append(self.process_proxy(*proxy_data))
-                self.queue.task_done()
-                i += 1
+                while not self.queue.empty() and i <= settings.CONCURRENT_TASKS_COUNT:
+                    proxy_data = self.queue.get_nowait()
+                    tasks.append(self.process_proxy(*proxy_data))
+                    self.queue.task_done()
+                    i += 1
 
-            if tasks:
-                await asyncio.wait(tasks)
-                tasks.clear()
+                if tasks:
+                    await asyncio.wait(tasks)
+                    tasks.clear()
+            except KeyboardInterrupt:
+                raise
+            except BaseException as ex:
+                self.logger.exception(ex)
+                await asyncio.sleep(10)
 
     async def producer(self):
         while True:
