RethinkDB Process Runner

rethink_runner.py

import os
import subprocess
import shutil
import signal
import tempfile
import time


class RethinkProcess(object):
    def __init__(self):
        self.data_dir = None
        self.pid_file = None
        self.log_file = None
        self.driver_port = 2815
        self.cluster_port = 2915
        self._process = None

    def _wait_until_available(self):
        available = False

        while not os.path.exists(self.log_file):
            time.sleep(0.01)

        while not available:
            with open(self.log_file, 'r') as fp:
                for line in fp:
                    if 'Listening for client driver' in line:
                        available = True
                        break
                    elif 'error: TCP socket creation failed' in line:
                        while os.path.exists(self.pid_file):
                            time.sleep(0.01)

                        raise OSError('Could not bind to ports')

    def start(self):
        self.data_dir = tempfile.mkdtemp()
        self.pid_file = os.path.join(self.data_dir, 'pidfile')
        self.log_file = os.path.join(self.data_dir, 'logfile')

        cmd = [
            'rethinkdb', '--no-http-admin', '--no-update-check', '--daemon',
            '--directory', self.data_dir,
            '--pid-file', self.pid_file,
            '--log-file', self.log_file,
            '--driver-port', str(self.driver_port),
            '--cluster-port', str(self.cluster_port)
        ]

        self._process = subprocess.Popen(cmd, stdout=subprocess.PIPE)
        self._process.communicate()

        if self._process.returncode != 0:
            raise Exception('Failed to start RethinkDB...')

        self._wait_until_available()

    def stop(self):
        if self._process:
            with open(self.pid_file, 'r') as fp:
                pid = int(fp.read())
                os.kill(pid, signal.SIGTERM)

        if self.data_dir:
            shutil.rmtree(self.data_dir)