Data Aggregator Rewrite

.pre-commit-config.yaml

@@ -9,3 +9,4 @@ repos:
       - id: black
         language_version: python3 # Should be a command that runs python3.6+
 
+

.travis.yml

@@ -6,11 +6,11 @@ env:
 # Once we add and remove tests, this distribution may become unbalanced.
 # Feel free to move tests around to make the running time of the jobs
 # as close as possible.
-  - TESTS=test_[a-e]*
-  - TESTS=test_[f-h]*
-  - TESTS=test_[i-r,t-z]*
+  - TESTS=test/test_[a-e]*
+  - TESTS=test/test_[f-h]*
+  - TESTS=test/test_[i-r,t-z]*
 # test_simple_commands.py is slow due to parametrization.
-  - TESTS=test_[s]*
+  - TESTS=test/test_[s]*
   - TESTS=webextension
 git:
   depth: 3
@@ -36,15 +36,11 @@ after_success:
 
 jobs:
   include:
-    - language:
-      python:
+    - language: "python"
       env:
         - TESTS="Docker"
       services:
         - docker
-      before_install:
-      before_script:
-      install:
       script:
         - docker build -f Dockerfile -t openwpm .
         - ./scripts/deploy-to-dockerhub.sh

automation/BrowserManager.py

@@ -20,7 +20,7 @@
 from .Commands.Types import ShutdownCommand
 from .DeployBrowsers import deploy_browser
 from .Errors import BrowserConfigError, BrowserCrashError, ProfileLoadError
-from .SocketInterface import clientsocket
+from .SocketInterface import ClientSocket
 from .utilities.multiprocess_utils import (
     Process,
     kill_process_and_children,
@@ -478,7 +478,7 @@ def BrowserManager(
                 "BROWSER %i: Connecting to extension on port %i"
                 % (browser_params["browser_id"], port)
             )
-            extension_socket = clientsocket(serialization="json")
+            extension_socket = ClientSocket(serialization="json")
             extension_socket.connect("127.0.0.1", int(port))
         else:
             extension_socket = None

automation/CommandSequence.py

@@ -190,7 +190,7 @@ def run_custom_function(self, function_handle, func_args=(), timeout=30):
         command = RunCustomFunctionCommand(function_handle, func_args)
         self._commands_with_timeout.append((command, timeout))
 
-    def mark_done(self, success: bool):
+    def mark_done(self, success: bool) -> None:
         if self.callback is not None:
             self.callback(success)
 

automation/Commands/browser_commands.py

@@ -20,7 +20,7 @@
 from selenium.webdriver.support import expected_conditions as EC
 from selenium.webdriver.support.ui import WebDriverWait
 
-from ..SocketInterface import clientsocket
+from ..SocketInterface import ClientSocket
 from .utils.webdriver_utils import (
     execute_in_all_frames,
     execute_script_with_retry,
@@ -112,7 +112,7 @@ def tab_restart_browser(webdriver):
 
 
 def get_website(
-    url, sleep, visit_id, webdriver, browser_params, extension_socket: clientsocket
+    url, sleep, visit_id, webdriver, browser_params, extension_socket: ClientSocket
 ):
     """
     goes to <url> using the given <webdriver> instance
@@ -372,7 +372,7 @@ def collect_source(driver, frame_stack, rv={}):
 
 
 def finalize(
-    visit_id: int, webdriver: WebDriver, extension_socket: clientsocket, sleep: int
+    visit_id: int, webdriver: WebDriver, extension_socket: ClientSocket, sleep: int
 ) -> None:
     """ Informs the extension that a visit is done """
     tab_restart_browser(webdriver)
@@ -383,6 +383,6 @@ def finalize(
     extension_socket.send(msg)
 
 
-def initialize(visit_id: int, extension_socket: clientsocket) -> None:
+def initialize(visit_id: int, extension_socket: ClientSocket) -> None:
     msg = {"action": "Initialize", "visit_id": visit_id}
     extension_socket.send(msg)

automation/DataAggregator/BaseAggregator.py

@@ -7,7 +7,7 @@
 
 from multiprocess import Queue
 
-from ..SocketInterface import serversocket
+from ..SocketInterface import ServerSocket
 from ..utilities.multiprocess_utils import Process
 
 RECORD_TYPE_CONTENT = "page_content"
@@ -61,7 +61,7 @@ def __init__(
         self.record_queue: Queue = None  # Initialized on `startup`
         self.logger = logging.getLogger("openwpm")
         self.curent_visit_ids: List[int] = list()  # All visit_ids in flight
-        self.sock: Optional[serversocket] = None
+        self.sock: Optional[ServerSocket] = None
 
     @abc.abstractmethod
     def process_record(self, record):
@@ -84,7 +84,9 @@ def process_content(self, record):
             for (content, content_hash)"""
 
     @abc.abstractmethod
-    def run_visit_completion_tasks(self, visit_id: int, interrupted: bool = False):
+    def run_visit_completion_tasks(
+        self, visit_id: int, interrupted: bool = False
+    ) -> None:
         """Will be called once a visit_id will receive no new records
 
         Parameters
@@ -98,7 +100,7 @@ def startup(self):
         """Run listener startup tasks
 
         Note: Child classes should call this method"""
-        self.sock = serversocket(name=type(self).__name__)
+        self.sock = ServerSocket(name=type(self).__name__)
         self.status_queue.put(self.sock.sock.getsockname())
         self.sock.start_accepting()
         self.record_queue = self.sock.queue
@@ -162,7 +164,7 @@ def mark_visit_complete(self, visit_id: int) -> None:
         relating to a certain visit_id have been saved"""
         self.completion_queue.put((visit_id, False))
 
-    def mark_visit_incomplete(self, visit_id: int):
+    def mark_visit_incomplete(self, visit_id: int) -> None:
         """This function should be called to indicate that a certain visit
         has been interrupted and will forever be incomplete
         """
@@ -224,6 +226,10 @@ def get_next_visit_id(self):
     def get_next_browser_id(self):
         """Return a unique crawl ID used as a key for a browser instance"""
 
+    @abc.abstractmethod
+    def launch(self):
+        """Launch the aggregator listener process"""
+
     def get_most_recent_status(self):
         """Return the most recent queue size sent from the listener process"""
 
@@ -273,15 +279,15 @@ def get_new_completed_visits(self) -> List[Tuple[int, bool]]:
             finished_visit_ids.append(self.completion_queue.get())
         return finished_visit_ids
 
-    def launch(self, listener_process_runner, *args):
+    def _launch(self, listener_process_runner, *args):
         """Launch the aggregator listener process"""
         args = ((self.status_queue, self.completion_queue, self.shutdown_queue),) + args
         self.listener_process = Process(target=listener_process_runner, args=args)
         self.listener_process.daemon = True
         self.listener_process.start()
         self.listener_address = self.status_queue.get()
 
-    def shutdown(self, relaxed: bool = True):
+    def shutdown(self, relaxed: bool = True) -> None:
         """ Terminate the aggregator listener process"""
         self.logger.debug(
             "Sending the shutdown signal to the %s listener process..."

automation/DataAggregator/LocalAggregator.py

@@ -19,7 +19,9 @@
 SQL_BATCH_SIZE = 1000
 LDB_BATCH_SIZE = 100
 MIN_TIME = 5  # seconds
-SCHEMA_FILE = os.path.join(os.path.dirname(__file__), "schema.sql")
+SCHEMA_FILE = os.path.join(
+    os.path.dirname(__file__), "..", "data_aggregator", "schema.sql"
+)
 LDB_NAME = "content.ldb"
 
 
@@ -87,7 +89,7 @@ def _generate_insert(self, table, data):
         statement = statement + ") " + value_str + ")"
         return statement, values
 
-    def process_record(self, record: Tuple[str, Union[str, Dict[str, Any]]]):
+    def process_record(self, record: Tuple[str, Union[str, Dict[str, Any]]]) -> None:
         """Add `record` to database"""
 
         if len(record) != 2:
@@ -181,7 +183,9 @@ def maybe_commit_records(self):
             self._ldb_counter = 0
             self._ldb_commit_time = time.time()
 
-    def run_visit_completion_tasks(self, visit_id: int, interrupted: bool = False):
+    def run_visit_completion_tasks(
+        self, visit_id: int, interrupted: bool = False
+    ) -> None:
         if interrupted:
             self.logger.warning("Visit with visit_id %d got interrupted", visit_id)
             self.cur.execute("INSERT INTO incomplete_visits VALUES (?)", (visit_id,))
@@ -283,7 +287,7 @@ def get_next_browser_id(self):
 
     def launch(self):
         """Launch the aggregator listener process"""
-        super(LocalAggregator, self).launch(
+        super(LocalAggregator, self)._launch(
             listener_process_runner, self.manager_params, self.ldb_enabled
         )
 

automation/DataAggregator/S3Aggregator.py

@@ -18,6 +18,7 @@
 from botocore.exceptions import ClientError, EndpointConnectionError
 from pyarrow.filesystem import S3FSWrapper  # noqa
 
+from ..data_aggregator.parquet_schema import PQ_SCHEMAS
 from .BaseAggregator import (
     RECORD_TYPE_CONTENT,
     RECORD_TYPE_CREATE,
@@ -26,7 +27,6 @@
     BaseListener,
     BaseParams,
 )
-from .parquet_schema import PQ_SCHEMAS
 
 CACHE_SIZE = 500
 SITE_VISITS_INDEX = "_site_visits_index"
@@ -76,7 +76,7 @@ def __init__(
         def factory_function():
             return defaultdict(list)
 
-        self._records: Dict[int, DefaultDict[str, List[Any]]] = defaultdict(
+        self._records: DefaultDict[int, DefaultDict[str, List[Any]]] = defaultdict(
             factory_function
         )  # maps visit_id and table to records
         self._batches: DefaultDict[str, List[pa.RecordBatch]] = defaultdict(
@@ -302,7 +302,9 @@ def drain_queue(self):
         super(S3Listener, self).drain_queue()
         self._send_to_s3(force=True)
 
-    def run_visit_completion_tasks(self, visit_id: int, interrupted: bool = False):
+    def run_visit_completion_tasks(
+        self, visit_id: int, interrupted: bool = False
+    ) -> None:
         if interrupted:
             self.logger.error("Visit with visit_id %d got interrupted", visit_id)
             self._write_record("incomplete_visits", {"visit_id": visit_id}, visit_id)
@@ -406,6 +408,6 @@ def get_next_browser_id(self):
 
     def launch(self):
         """Launch the aggregator listener process"""
-        super(S3Aggregator, self).launch(
+        super(S3Aggregator, self)._launch(
             listener_process_runner, self.manager_params, self._instance_id
         )

automation/MPLogger.py

@@ -17,7 +17,7 @@
 from tblib import pickling_support
 
 from .Commands.utils.webdriver_utils import parse_neterror
-from .SocketInterface import serversocket
+from .SocketInterface import ServerSocket
 
 pickling_support.install()
 
@@ -218,7 +218,7 @@ def _initialize_sentry(self):
 
     def _start_listener(self):
         """Start listening socket for remote logs from extension"""
-        socket = serversocket(name="loggingserver")
+        socket = ServerSocket(name="loggingserver")
         self._status_queue.put(socket.sock.getsockname())
         socket.start_accepting()
         self._status_queue.join()  # block to allow parent to retrieve address

automation/SocketInterface.py

@@ -11,7 +11,7 @@
 # see: https://stackoverflow.com/a/1148237
 
 
-class serversocket:
+class ServerSocket:
     """
     A server socket to receive and process string messages
     from client sockets to a central queue
@@ -111,7 +111,7 @@ def close(self):
         self.sock.close()
 
 
-class clientsocket:
+class ClientSocket:
     """A client socket for sending messages"""
 
     def __init__(self, serialization="json", verbose=False):
@@ -171,17 +171,17 @@ def main():
 
     # Just for testing
     if sys.argv[1] == "s":
-        sock = serversocket(verbose=True)
-        sock.start_accepting()
+        ssock: ServerSocket = ServerSocket(verbose=True)
+        ssock.start_accepting()
         input("Press enter to exit...")
-        sock.close()
+        ssock.close()
     elif sys.argv[1] == "c":
         host = input("Enter the host name:\n")
         port = input("Enter the port:\n")
         serialization = input("Enter the serialization type (default: 'json'):\n")
         if serialization == "":
             serialization = "json"
-        sock = clientsocket(serialization=serialization)
+        sock: ClientSocket = ClientSocket(serialization=serialization)
         sock.connect(host, int(port))
         msg = None