<commit_before>from Queue import Queue, Empty
import contextlib
import datetime
import json
import logging
import os
import pickle
import select
import threading
import time

import psutil
import psycopg2.pool
import psycopg2.extras
import psycopg2 as pg

import aque.utils as utils
from .base import Broker


log = logging.getLogger(__name__)


class PostgresBroker(Broker):
    """A :class:`.Broker` which uses Postgresql_ as a data store and event dispatcher.

    .. _Postgresql: http://www.postgresql.org/
    """

    @classmethod
    def from_url(cls, parts):
        return cls(host=parts.netloc, database=parts.path.strip('/').lower())

    def __init__(self, **kwargs):

        self._notify_conn = None
        self._notify_lock = threading.Lock()
        self._notify_thread = None
        
        super(PostgresBroker, self).__init__()

        self._kwargs = kwargs
        self._pool = self._open_pool()

        self._reflect()


        self._heartbeat_lock = threading.Lock()
        self._heartbeat_thread = None
        self._captured_tids = []

    def _open_pool(self):
        return pg.pool.ThreadedConnectionPool(0, 4 * psutil.cpu_count(), **self._kwargs)

    def after_fork(self):
        self._pool = self._open_pool()

    def __del__(self):
        self.close()

    def close(self):
        pass
    
    def _notify_start(self):
        with self._notify_lock:
            if not self._notify_thread:
                self._notify_thread = threading.Thread(target=self._notify_target)
                self._notify_thread.daemon=True
                self._notify_thread.start()

    @contextlib.contextmanager
    def _connect(self):
        conn = self._pool.getconn()
        # conn.initialize(logging.getLogger('aque.sql'))
        try:
            yield conn
            conn.commit()
        finally:
            self._pool.putconn(conn)

    @contextlib.contextmanager
    def _cursor(self):
        with self._connect() as conn:
            with conn.cursor() as cur:
                yield cur

    def update_schema(self):
        with self._cursor() as cur:
            cur.execute('''CREATE TABLE IF NOT EXISTS tasks (

                id SERIAL PRIMARY KEY,
                name TEXT,

                dependencies INTEGER[],

                -- requirements
                cpus REAL, -- fractional CPUs are okay
                memory INTEGER, -- in bytes
                platform TEXT,
                host TEXT,

                status TEXT NOT NULL DEFAULT 'creating',
                last_active TIMESTAMP,

                -- priority
                priority INTEGER NOT NULL DEFAULT 1000,
                io_paths TEXT[],
                duration INTEGER,

                -- execution
                pattern BYTEA,
                func BYTEA,
                args BYTEA,
                kwargs BYTEA,

                -- environmental
                "user" TEXT,
                "group" TEXT,
                cwd TEXT,
                environ BYTEA,

                -- the final result, or exception (depending on status)
                result BYTEA,

                -- everything else
                extra BYTEA

            )''')

            cur.execute('''CREATE TABLE IF NOT EXISTS output_logs (
                task_id SERIAL NOT NULL,
                ctime TIMESTAMP NOT NULL DEFAULT localtimestamp,
                fd SMALLINT NOT NULL,
                content TEXT NOT NULL
            )''')
            # TODO: create an index

    def _reflect(self):
        with self._cursor() as cur:
            # Determine what rows we actually have.
            cur.execute('''SELECT column_name, data_type FROM information_schema.columns WHERE table_name = 'tasks' ORDER BY ordinal_position''')
            self._task_fields = tuple(cur)
            self._task_field_types = dict(self._task_fields)

    def destroy_schema(self):
        with self._cursor() as cur:
            cur.execute('''DROP TABLE IF EXISTS dependencies''')
            cur.execute('''DROP TABLE IF EXISTS tasks''')<commit_msg>PGBroker uses very basic schema migrations<commit_after>from Queue import Queue, Empty
import contextlib
import datetime
import functools
import json
import logging
import os
import pickle
import select
import threading
import time

import psutil
import psycopg2.pool
import psycopg2.extras
import psycopg2 as pg

import aque.utils as utils
from aque.brokers.base import Broker


log = logging.getLogger(__name__)


schema_migrations = []
def patch(func=None, name=None):
    if func is None:
        return functools.partial(patch, name=name)

    name = name or func.__name__
    name = name.strip('_')
    schema_migrations.append((name, func))

    return func


@patch
def create_task_table(cur):
    cur.execute('''CREATE TABLE tasks (

        id SERIAL PRIMARY KEY,
        name TEXT,

        dependencies INTEGER[],

        -- requirements
        cpus REAL, -- fractional CPUs are okay
        memory INTEGER, -- in bytes
        platform TEXT,
        host TEXT,

        status TEXT NOT NULL DEFAULT 'creating',
        creation_time TIMESTAMP NOT NULL DEFAULT localtimestamp,
        first_active TIMESTAMP,
        last_active TIMESTAMP,

        -- priority
        priority INTEGER NOT NULL DEFAULT 1000,
        io_paths TEXT[],
        duration INTEGER,

        -- execution
        pattern BYTEA,
        func BYTEA,
        args BYTEA,
        kwargs BYTEA,

        -- environmental
        "user" TEXT,
        "group" TEXT,
        cwd TEXT,
        environ BYTEA,

        -- the final result, or exception (depending on status)
        result BYTEA,

        -- everything else
        extra BYTEA

    )''')
    cur.execute('CREATE INDEX tasks_status_index ON tasks (id, status)')


@patch
def create_output_log_table(cur):
    cur.execute('''CREATE TABLE output_logs (
        task_id SERIAL NOT NULL,
        ctime TIMESTAMP NOT NULL DEFAULT localtimestamp,
        fd SMALLINT NOT NULL,
        content TEXT NOT NULL
    )''')
    cur.execute('CREATE INDEX output_logs_index ON output_logs (task_id)')


class PostgresBroker(Broker):
    """A :class:`.Broker` which uses Postgresql_ as a data store and event dispatcher.

    .. _Postgresql: http://www.postgresql.org/
    """

    @classmethod
    def from_url(cls, parts):
        return cls(host=parts.netloc, database=parts.path.strip('/').lower())

    def __init__(self, **kwargs):

        self._notify_conn = None
        self._notify_lock = threading.Lock()
        self._notify_thread = None
        
        super(PostgresBroker, self).__init__()

        self._kwargs = kwargs
        self._pool = self._open_pool()

        self._reflect()


        self._heartbeat_lock = threading.Lock()
        self._heartbeat_thread = None
        self._captured_tids = []

    def _open_pool(self):
        return pg.pool.ThreadedConnectionPool(0, 4 * psutil.cpu_count(), **self._kwargs)

    def after_fork(self):
        self._pool = self._open_pool()

    def __del__(self):
        self.close()

    def close(self):
        pass
    
    def _notify_start(self):
        with self._notify_lock:
            if not self._notify_thread:
                self._notify_thread = threading.Thread(target=self._notify_target)
                self._notify_thread.daemon=True
                self._notify_thread.start()

    @contextlib.contextmanager
    def _connect(self):
        conn = self._pool.getconn()
        # conn.initialize(logging.getLogger('aque.sql'))
        try:
            yield conn
            conn.commit()
        finally:
            self._pool.putconn(conn)

    @contextlib.contextmanager
    def _cursor(self):
        with self._connect() as conn:
            with conn.cursor() as cur:
                yield cur

    def update_schema(self):

        with self._cursor() as cur:
            cur.execute('''CREATE TABLE IF NOT EXISTS schema_migrations (
                id SERIAL PRIMARY KEY,
                ctime TIMESTAMP NOT NULL DEFAULT localtimestamp,
                name TEXT NOT NULL
            )''')
            cur.execute('SELECT name from schema_migrations')
            applied_patches = set(row[0] for row in cur)

        for name, func in schema_migrations:
            if name not in applied_patches:
                log.info('applying schema migration %s' % name)
                with self._cursor() as cur:
                    func(cur)
                    cur.execute('INSERT INTO schema_migrations (name) VALUES (%s)', [name])

    def _reflect(self):
        with self._cursor() as cur:
            cur.execute('''SELECT column_name, data_type FROM information_schema.columns WHERE table_name = 'tasks' ORDER BY ordinal_position''')
            self._task_field_types = dict(cur)

    def destroy_schema(self):
        with self._cursor() as cur:
            cur.execute('''DROP TABLE IF EXISTS schema_migrations''')
            cur.execute('''DROP TABLE IF EXISTS tasks''')
{'fraction_non_alphanumeric': 0.08587999107740353, 'fraction_numerical': 0.0021191166629489183, 'mean_word_length': 1.9574538258575198, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 3, 'https://': 0, 'lorem ipsum': 0, 'www.': 2, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '6852995', 'n_tokens_mistral': 2552, 'n_tokens_neox': 2405, 'n_words': 687}