<reponame>antmicro/buildgrid
# Copyright (C) 2019 Bloomberg LP
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#  <http://www.apache.org/licenses/LICENSE-2.0>
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.


from contextlib import contextmanager
import logging
import os
import time
from tempfile import NamedTemporaryFile

from alembic import command
from alembic.config import Config
from sqlalchemy import and_, create_engine, event, or_
from sqlalchemy.orm.session import sessionmaker
from sqlalchemy.exc import OperationalError
from sqlalchemy.pool import StaticPool

from ...._enums import OperationStage
from ....settings import MAX_JOB_BLOCK_TIME
from ..interface import DataStoreInterface
from .models import digest_to_string, Job, Lease, Operation, PlatformRequirement


Session = sessionmaker()


def sqlite_on_connect(conn, record):
    conn.execute("PRAGMA journal_mode=WAL")
    conn.execute("PRAGMA synchronous=NORMAL")


class SQLDataStore(DataStoreInterface):

    def __init__(self, storage, *, connection_string=None, automigrate=False,
                 **kwargs):
        self.__logger = logging.getLogger(__name__)
        self.__logger.info("Creating SQL data store interface with: "
                           "automigrate=[%s], kwargs=[%s]",
                           automigrate, kwargs)

        self.storage = storage
        self.response_cache = {}

        # Set-up temporary SQLite Database when connection string is not specified
        if not connection_string:
            tmpdbfile = NamedTemporaryFile(prefix='bgd-', suffix='.db')
            self._tmpdbfile = tmpdbfile  # Make sure to keep this tempfile for the lifetime of this object
            self.__logger.warn("No connection string specified for the DataStore, "
                               "will use SQLite with tempfile: [%s]" % tmpdbfile.name)
            automigrate = True  # since this is a temporary database, we always need to create it
            connection_string = "sqlite:///{}".format(tmpdbfile.name)

        self._create_sqlalchemy_engine(connection_string, automigrate, **kwargs)

    def _create_sqlalchemy_engine(self, connection_string, automigrate, **kwargs):
        self.automigrate = automigrate

        # Disallow sqlite in-memory because multi-threaded access to it is
        # complex and potentially problematic at best
        # ref: https://docs.sqlalchemy.org/en/13/dialects/sqlite.html#threading-pooling-behavior
        if self._is_sqlite_inmemory_connection_string(connection_string):
            raise ValueError("Cannot use SQLite in-memory with BuildGrid (connection_string=[%s]). "
                             "Use a file or leave the connection_string empty for a tempfile." %
                             connection_string)

        # Only pass the (known) kwargs that have been explicitly set by the user
        available_options = set(['pool_size', 'max_overflow', 'pool_timeout', 'connect_args'])
        kwargs_keys = set(kwargs.keys())
        if not kwargs_keys.issubset(available_options):
            unknown_options = kwargs_keys - available_options
            raise TypeError("Unknown keyword arguments: [%s]" % unknown_options)

        self.__logger.debug("SQLAlchemy additional kwargs: [%s]", kwargs)

        self.engine = create_engine(connection_string, echo=False, **kwargs)
        Session.configure(bind=self.engine)

        if self.engine.dialect.name == "sqlite":
            event.listen(self.engine, "connect", sqlite_on_connect)

        if self.automigrate:
            self._create_or_migrate_db(connection_string)

    def _is_sqlite_connection_string(self, connection_string):
        if connection_string:
            return connection_string.startswith("sqlite")
        return False

    def _is_sqlite_inmemory_connection_string(self, full_connection_string):
        if self._is_sqlite_connection_string(full_connection_string):
            # Valid connection_strings for in-memory SQLite which we don't support could look like:
            # "sqlite:///file:memdb1?option=value&cache=shared&mode=memory",
            # "sqlite:///file:memdb1?mode=memory&cache=shared",
            # "sqlite:///file:memdb1?cache=shared&mode=memory",
            # "sqlite:///file::memory:?cache=shared",
            # "sqlite:///file::memory:",
            # "sqlite:///:memory:",
            # "sqlite:///",
            # "sqlite://"
            # ref: https://www.sqlite.org/inmemorydb.html
            # Note that a user can also specify drivers, so prefix could become 'sqlite+driver:///'
            connection_string = full_connection_string

            uri_split_index = connection_string.find("?")
            if uri_split_index != -1:
                connection_string = connection_string[0:uri_split_index]

            if connection_string.endswith((":memory:", ":///", "://")):
                return True
            elif uri_split_index != -1:
                opts = full_connection_string[uri_split_index + 1:].split("&")
                if "mode=memory" in opts:
                    return True

        return False

    def __repr__(self):
        return "SQL data store interface for `%s`" % repr(self.engine.url)

    def activate_monitoring(self):
        # Don't do anything. This function needs to exist but there's no
        # need to actually toggle monitoring in this implementation.
        pass

    def deactivate_monitoring(self):
        # Don't do anything. This function needs to exist but there's no
        # need to actually toggle monitoring in this implementation.
        pass

    def _create_or_migrate_db(self, connection_string):
        self.__logger.warn("Will attempt migration to latest version if needed.")

        config = Config()
        config.set_main_option("script_location", os.path.join(os.path.dirname(__file__), "alembic"))

        with self.engine.begin() as connection:
            config.attributes['connection'] = connection
            command.upgrade(config, "head")

    @contextmanager
    def session(self, sqlite_lock_immediately=False):
        session = Session()
        try:
            if sqlite_lock_immediately and session.bind.name == "sqlite":
                session.execute("BEGIN IMMEDIATE")
            yield session
            session.commit()
        except:
            session.rollback()
            self.__logger.error("Error in database session.", exc_info=True)
        finally:
            session.close()

    def _get_job(self, job_name, session, with_for_update=False):
        jobs = session.query(Job)
        if with_for_update:
            jobs = jobs.with_for_update()
        jobs = jobs.filter_by(name=job_name)
        return jobs.first()

    def get_job_by_action(self, action_digest):
        with self.session() as session:
            jobs = session.query(Job).filter_by(action_digest=digest_to_string(action_digest))
            jobs = jobs.filter(Job.stage != OperationStage.COMPLETED.value)
            job = jobs.first()
            if not job:
                return None
            return job.to_internal_job(self)

    def get_job_by_name(self, name):
        with self.session() as session:
            job = self._get_job(name, session)
            return job.to_internal_job(self)

    def get_job_by_operation(self, operation_name):
        with self.session() as session:
            operation = self._get_operation(operation_name, session)
            job = operation.job
            return job.to_internal_job(self)

    def get_all_jobs(self):
        with self.session() as session:
            jobs = session.query(Job).filter(Job.stage != OperationStage.COMPLETED.value)
            return [j.to_internal_job(self) for j in jobs]

    def create_job(self, job):
        with self.session() as session:
            if self._get_job(job.name, session) is None:
                requirements = [
                    PlatformRequirement(key=k, value=v)
                    for k, values in job.platform_requirements.items()
                    for v in values
                ]
                session.add(Job(
                    name=job.name,
                    action_digest=digest_to_string(job.action_digest),
                    do_not_cache=job.do_not_cache,
                    priority=job.priority,
                    operations=[],
                    platform_requirements=requirements,
                    stage=job.operation_stage.value,
                    queued_timestamp=job.queued_timestamp.ToDatetime(),
                    queued_time_duration=job.queued_time_duration.seconds,
                    worker_start_timestamp=job.worker_start_timestamp.ToDatetime(),
                    worker_completed_timestamp=job.worker_completed_timestamp.ToDatetime()
                ))
                self.response_cache[job.name] = job.execute_response

    def queue_job(self, job_name):
        with self.session(sqlite_lock_immediately=True) as session:
            job = self._get_job(job_name, session, with_for_update=True)
            job.assigned = False

    def update_job(self, job_name, changes):
        if "result" in changes:
            changes["result"] = digest_to_string(changes["result"])
        if "action_digest" in changes:
            changes["action_digest"] = digest_to_string(changes["action_digest"])

        with self.session() as session:
            job = self._get_job(job_name, session)
            job.update(changes)

    def delete_job(self, job_name):
        if job_name in self.response_cache:
            del self.response_cache[job_name]

    def store_response(self, job):
        digest = self.storage.put_message(job.execute_response)
        self.update_job(job.name, {"result": digest})
        self.response_cache[job.name] = job.execute_response

    def _get_operation(self, operation_name, session):
        operations = session.query(Operation).filter_by(name=operation_name)
        return operations.first()

    def get_operations_by_stage(self, operation_stage):
        with self.session() as session:
            operations = session.query(Operation)
            operations = operations.filter(Operation.job.has(stage=operation_stage.value))
            operations = operations.all()
            # Return a set of job names here for now, to match the `MemoryDataStore`
            # implementation's behaviour
            return set(op.job.name for op in operations)

    def get_all_operations(self):
        with self.session() as session:
            operations = session.query(Operation)
            return [op.name for op in operations]

    def create_operation(self, operation, job_name):
        with self.session() as session:
            session.add(Operation(
                name=operation.name,
                job_name=job_name,
                done=operation.done
            ))

    def update_operation(self, operation_name, changes):
        with self.session() as session:
            operation = self._get_operation(operation_name, session)
            operation.update(changes)

    def delete_operation(self, operation_name):
        # Don't do anything. This function needs to exist but there's no
        # need to actually delete operations in this implementation.
        pass

    def get_leases_by_state(self, lease_state):
        with self.session() as session:
            leases = session.query(Lease).filter_by(state=lease_state.value)
            leases = leases.all()
            # `lease.job_name` is the same as `lease.id` for a Lease protobuf
            return set(lease.job_name for lease in leases)

    def _create_lease(self, lease, session, job=None):
        if job is None:
            job = self._get_job(lease.id, session)
            job = job.to_internal_job(self)
        session.add(Lease(
            job_name=lease.id,
            state=lease.state,
            status=None,
            worker_name=job.worker_name
        ))

    def create_lease(self, lease):
        with self.session() as session:
            self._create_lease(lease, session)

    def update_lease(self, job_name, changes):
        with self.session() as session:
            job = self._get_job(job_name, session)
            lease = job.active_leases[0]
            lease.update(changes)

    def load_unfinished_jobs(self):
        with self.session() as session:
            jobs = session.query(Job)
            jobs = jobs.filter(Job.stage != OperationStage.COMPLETED.value)
            jobs = jobs.order_by(Job.priority)
            return [j.to_internal_job(self) for j in jobs.all()]

    def assign_lease_for_next_job(self, capabilities, callback, timeout=None):
        """Return a list of leases for the highest priority jobs that can be run by a worker.

        NOTE: Currently the list only ever has one or zero leases.

        Query the jobs table to find queued jobs which match the capabilities of
        a given worker, and return the one with the highest priority. Takes a
        dictionary of worker capabilities to compare with job requirements.

        :param capabilities: Dictionary of worker capabilities to compare
            with job requirements when finding a job.
        :type capabilities: dict
        :param callback: Function to run on the next runnable job, should return
            a list of leases.
        :type callback: function
        :param timeout: time to wait for new jobs, caps if longer
            than MAX_JOB_BLOCK_TIME.
        :type timeout: int
        :returns: List of leases

        """
        if not timeout:
            return self._assign_job_leases(capabilities, callback)

        # Cap the timeout if it's larger than MAX_JOB_BLOCK_TIME
        if timeout:
            timeout = min(timeout, MAX_JOB_BLOCK_TIME)

        start = time.time()
        while time.time() + 1 - start < timeout:
            leases = self._assign_job_leases(capabilities, callback)
            if leases:
                return leases
            time.sleep(0.5)

    def _assign_job_leases(self, capabilities, callback):
        with self.session(sqlite_lock_immediately=True) as session:
            jobs = session.query(Job).with_for_update(skip_locked=True)
            jobs = jobs.filter(Job.stage == OperationStage.QUEUED.value)
            jobs = jobs.filter(Job.assigned != True)  # noqa

            # Filter to find just jobs that either don't have requirements, or
            # jobs with requirement keys that are a subset of the worker's
            # advertised capabilities keys
            jobs = jobs.filter(or_(
                ~Job.platform_requirements.any(),
                and_(
                    Job.platform_requirements.any(
                        PlatformRequirement.key.in_(capabilities.keys())
                    ),
                    ~Job.platform_requirements.any(
                        ~PlatformRequirement.key.in_(capabilities.keys())
                    )
                )
            ))

            # For each advertised worker capability, filter to find jobs which
            # either don't require that capability or can be satisfied by the
            # capability
            for key, value in capabilities.items():
                jobs = jobs.filter(or_(
                    ~Job.platform_requirements.any(),
                    ~Job.platform_requirements.any(PlatformRequirement.key == key),
                    and_(
                        ~Job.platform_requirements.any(and_(
                            PlatformRequirement.key == key,
                            ~PlatformRequirement.value.in_(value)
                        )),
                        Job.platform_requirements.any(and_(
                            PlatformRequirement.key == key,
                            PlatformRequirement.value.in_(value)
                        ))
                    )
                ))

            jobs = jobs.order_by(Job.priority)
            job = jobs.first()
            if job:
                internal_job = job.to_internal_job(self)
                leases = callback(internal_job)
                if leases:
                    job.assigned = True
                    job.worker_start_timestamp = internal_job.worker_start_timestamp.ToDatetime()
                for lease in leases:
                    self._create_lease(lease, session, job=internal_job)
                return leases
            return []

# Copyright (C) 2019 Bloomberg LP
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#  <http://www.apache.org/licenses/LICENSE-2.0>
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.


from contextlib import contextmanager
import logging
import os
import time
from tempfile import NamedTemporaryFile

from alembic import command
from alembic.config import Config
from sqlalchemy import and_, create_engine, event, or_
from sqlalchemy.orm.session import sessionmaker
from sqlalchemy.exc import OperationalError
from sqlalchemy.pool import StaticPool

from ...._enums import OperationStage
from ....settings import MAX_JOB_BLOCK_TIME
from ..interface import DataStoreInterface
from .models import digest_to_string, Job, Lease, Operation, PlatformRequirement


Session = sessionmaker()


def sqlite_on_connect(conn, record):
    conn.execute("PRAGMA journal_mode=WAL")
    conn.execute("PRAGMA synchronous=NORMAL")


class SQLDataStore(DataStoreInterface):

    def __init__(self, storage, *, connection_string=None, automigrate=False,
                 **kwargs):
        self.__logger = logging.getLogger(__name__)
        self.__logger.info("Creating SQL data store interface with: "
                           "automigrate=[%s], kwargs=[%s]",
                           automigrate, kwargs)

        self.storage = storage
        self.response_cache = {}

        # Set-up temporary SQLite Database when connection string is not specified
        if not connection_string:
            tmpdbfile = NamedTemporaryFile(prefix='bgd-', suffix='.db')
            self._tmpdbfile = tmpdbfile  # Make sure to keep this tempfile for the lifetime of this object
            self.__logger.warn("No connection string specified for the DataStore, "
                               "will use SQLite with tempfile: [%s]" % tmpdbfile.name)
            automigrate = True  # since this is a temporary database, we always need to create it
            connection_string = "sqlite:///{}".format(tmpdbfile.name)

        self._create_sqlalchemy_engine(connection_string, automigrate, **kwargs)

    def _create_sqlalchemy_engine(self, connection_string, automigrate, **kwargs):
        self.automigrate = automigrate

        # Disallow sqlite in-memory because multi-threaded access to it is
        # complex and potentially problematic at best
        # ref: https://docs.sqlalchemy.org/en/13/dialects/sqlite.html#threading-pooling-behavior
        if self._is_sqlite_inmemory_connection_string(connection_string):
            raise ValueError("Cannot use SQLite in-memory with BuildGrid (connection_string=[%s]). "
                             "Use a file or leave the connection_string empty for a tempfile." %
                             connection_string)

        # Only pass the (known) kwargs that have been explicitly set by the user
        available_options = set(['pool_size', 'max_overflow', 'pool_timeout', 'connect_args'])
        kwargs_keys = set(kwargs.keys())
        if not kwargs_keys.issubset(available_options):
            unknown_options = kwargs_keys - available_options
            raise TypeError("Unknown keyword arguments: [%s]" % unknown_options)

        self.__logger.debug("SQLAlchemy additional kwargs: [%s]", kwargs)

        self.engine = create_engine(connection_string, echo=False, **kwargs)
        Session.configure(bind=self.engine)

        if self.engine.dialect.name == "sqlite":
            event.listen(self.engine, "connect", sqlite_on_connect)

        if self.automigrate:
            self._create_or_migrate_db(connection_string)

    def _is_sqlite_connection_string(self, connection_string):
        if connection_string:
            return connection_string.startswith("sqlite")
        return False

    def _is_sqlite_inmemory_connection_string(self, full_connection_string):
        if self._is_sqlite_connection_string(full_connection_string):
            # Valid connection_strings for in-memory SQLite which we don't support could look like:
            # "sqlite:///file:memdb1?option=value&cache=shared&mode=memory",
            # "sqlite:///file:memdb1?mode=memory&cache=shared",
            # "sqlite:///file:memdb1?cache=shared&mode=memory",
            # "sqlite:///file::memory:?cache=shared",
            # "sqlite:///file::memory:",
            # "sqlite:///:memory:",
            # "sqlite:///",
            # "sqlite://"
            # ref: https://www.sqlite.org/inmemorydb.html
            # Note that a user can also specify drivers, so prefix could become 'sqlite+driver:///'
            connection_string = full_connection_string

            uri_split_index = connection_string.find("?")
            if uri_split_index != -1:
                connection_string = connection_string[0:uri_split_index]

            if connection_string.endswith((":memory:", ":///", "://")):
                return True
            elif uri_split_index != -1:
                opts = full_connection_string[uri_split_index + 1:].split("&")
                if "mode=memory" in opts:
                    return True

        return False

    def __repr__(self):
        return "SQL data store interface for `%s`" % repr(self.engine.url)

    def activate_monitoring(self):
        # Don't do anything. This function needs to exist but there's no
        # need to actually toggle monitoring in this implementation.
        pass

    def deactivate_monitoring(self):
        # Don't do anything. This function needs to exist but there's no
        # need to actually toggle monitoring in this implementation.
        pass

    def _create_or_migrate_db(self, connection_string):
        self.__logger.warn("Will attempt migration to latest version if needed.")

        config = Config()
        config.set_main_option("script_location", os.path.join(os.path.dirname(__file__), "alembic"))

        with self.engine.begin() as connection:
            config.attributes['connection'] = connection
            command.upgrade(config, "head")

    @contextmanager
    def session(self, sqlite_lock_immediately=False):
        session = Session()
        try:
            if sqlite_lock_immediately and session.bind.name == "sqlite":
                session.execute("BEGIN IMMEDIATE")
            yield session
            session.commit()
        except:
            session.rollback()
            self.__logger.error("Error in database session.", exc_info=True)
        finally:
            session.close()

    def _get_job(self, job_name, session, with_for_update=False):
        jobs = session.query(Job)
        if with_for_update:
            jobs = jobs.with_for_update()
        jobs = jobs.filter_by(name=job_name)
        return jobs.first()

    def get_job_by_action(self, action_digest):
        with self.session() as session:
            jobs = session.query(Job).filter_by(action_digest=digest_to_string(action_digest))
            jobs = jobs.filter(Job.stage != OperationStage.COMPLETED.value)
            job = jobs.first()
            if not job:
                return None
            return job.to_internal_job(self)

    def get_job_by_name(self, name):
        with self.session() as session:
            job = self._get_job(name, session)
            return job.to_internal_job(self)

    def get_job_by_operation(self, operation_name):
        with self.session() as session:
            operation = self._get_operation(operation_name, session)
            job = operation.job
            return job.to_internal_job(self)

    def get_all_jobs(self):
        with self.session() as session:
            jobs = session.query(Job).filter(Job.stage != OperationStage.COMPLETED.value)
            return [j.to_internal_job(self) for j in jobs]

    def create_job(self, job):
        with self.session() as session:
            if self._get_job(job.name, session) is None:
                requirements = [
                    PlatformRequirement(key=k, value=v)
                    for k, values in job.platform_requirements.items()
                    for v in values
                ]
                session.add(Job(
                    name=job.name,
                    action_digest=digest_to_string(job.action_digest),
                    do_not_cache=job.do_not_cache,
                    priority=job.priority,
                    operations=[],
                    platform_requirements=requirements,
                    stage=job.operation_stage.value,
                    queued_timestamp=job.queued_timestamp.ToDatetime(),
                    queued_time_duration=job.queued_time_duration.seconds,
                    worker_start_timestamp=job.worker_start_timestamp.ToDatetime(),
                    worker_completed_timestamp=job.worker_completed_timestamp.ToDatetime()
                ))
                self.response_cache[job.name] = job.execute_response

    def queue_job(self, job_name):
        with self.session(sqlite_lock_immediately=True) as session:
            job = self._get_job(job_name, session, with_for_update=True)
            job.assigned = False

    def update_job(self, job_name, changes):
        if "result" in changes:
            changes["result"] = digest_to_string(changes["result"])
        if "action_digest" in changes:
            changes["action_digest"] = digest_to_string(changes["action_digest"])

        with self.session() as session:
            job = self._get_job(job_name, session)
            job.update(changes)

    def delete_job(self, job_name):
        if job_name in self.response_cache:
            del self.response_cache[job_name]

    def store_response(self, job):
        digest = self.storage.put_message(job.execute_response)
        self.update_job(job.name, {"result": digest})
        self.response_cache[job.name] = job.execute_response

    def _get_operation(self, operation_name, session):
        operations = session.query(Operation).filter_by(name=operation_name)
        return operations.first()

    def get_operations_by_stage(self, operation_stage):
        with self.session() as session:
            operations = session.query(Operation)
            operations = operations.filter(Operation.job.has(stage=operation_stage.value))
            operations = operations.all()
            # Return a set of job names here for now, to match the `MemoryDataStore`
            # implementation's behaviour
            return set(op.job.name for op in operations)

    def get_all_operations(self):
        with self.session() as session:
            operations = session.query(Operation)
            return [op.name for op in operations]

    def create_operation(self, operation, job_name):
        with self.session() as session:
            session.add(Operation(
                name=operation.name,
                job_name=job_name,
                done=operation.done
            ))

    def update_operation(self, operation_name, changes):
        with self.session() as session:
            operation = self._get_operation(operation_name, session)
            operation.update(changes)

    def delete_operation(self, operation_name):
        # Don't do anything. This function needs to exist but there's no
        # need to actually delete operations in this implementation.
        pass

    def get_leases_by_state(self, lease_state):
        with self.session() as session:
            leases = session.query(Lease).filter_by(state=lease_state.value)
            leases = leases.all()
            # `lease.job_name` is the same as `lease.id` for a Lease protobuf
            return set(lease.job_name for lease in leases)

    def _create_lease(self, lease, session, job=None):
        if job is None:
            job = self._get_job(lease.id, session)
            job = job.to_internal_job(self)
        session.add(Lease(
            job_name=lease.id,
            state=lease.state,
            status=None,
            worker_name=job.worker_name
        ))

    def create_lease(self, lease):
        with self.session() as session:
            self._create_lease(lease, session)

    def update_lease(self, job_name, changes):
        with self.session() as session:
            job = self._get_job(job_name, session)
            lease = job.active_leases[0]
            lease.update(changes)

    def load_unfinished_jobs(self):
        with self.session() as session:
            jobs = session.query(Job)
            jobs = jobs.filter(Job.stage != OperationStage.COMPLETED.value)
            jobs = jobs.order_by(Job.priority)
            return [j.to_internal_job(self) for j in jobs.all()]

    def assign_lease_for_next_job(self, capabilities, callback, timeout=None):
        """Return a list of leases for the highest priority jobs that can be run by a worker.

        NOTE: Currently the list only ever has one or zero leases.

        Query the jobs table to find queued jobs which match the capabilities of
        a given worker, and return the one with the highest priority. Takes a
        dictionary of worker capabilities to compare with job requirements.

        :param capabilities: Dictionary of worker capabilities to compare
            with job requirements when finding a job.
        :type capabilities: dict
        :param callback: Function to run on the next runnable job, should return
            a list of leases.
        :type callback: function
        :param timeout: time to wait for new jobs, caps if longer
            than MAX_JOB_BLOCK_TIME.
        :type timeout: int
        :returns: List of leases

        """
        if not timeout:
            return self._assign_job_leases(capabilities, callback)

        # Cap the timeout if it's larger than MAX_JOB_BLOCK_TIME
        if timeout:
            timeout = min(timeout, MAX_JOB_BLOCK_TIME)

        start = time.time()
        while time.time() + 1 - start < timeout:
            leases = self._assign_job_leases(capabilities, callback)
            if leases:
                return leases
            time.sleep(0.5)

    def _assign_job_leases(self, capabilities, callback):
        with self.session(sqlite_lock_immediately=True) as session:
            jobs = session.query(Job).with_for_update(skip_locked=True)
            jobs = jobs.filter(Job.stage == OperationStage.QUEUED.value)
            jobs = jobs.filter(Job.assigned != True)  # noqa

            # Filter to find just jobs that either don't have requirements, or
            # jobs with requirement keys that are a subset of the worker's
            # advertised capabilities keys
            jobs = jobs.filter(or_(
                ~Job.platform_requirements.any(),
                and_(
                    Job.platform_requirements.any(
                        PlatformRequirement.key.in_(capabilities.keys())
                    ),
                    ~Job.platform_requirements.any(
                        ~PlatformRequirement.key.in_(capabilities.keys())
                    )
                )
            ))

            # For each advertised worker capability, filter to find jobs which
            # either don't require that capability or can be satisfied by the
            # capability
            for key, value in capabilities.items():
                jobs = jobs.filter(or_(
                    ~Job.platform_requirements.any(),
                    ~Job.platform_requirements.any(PlatformRequirement.key == key),
                    and_(
                        ~Job.platform_requirements.any(and_(
                            PlatformRequirement.key == key,
                            ~PlatformRequirement.value.in_(value)
                        )),
                        Job.platform_requirements.any(and_(
                            PlatformRequirement.key == key,
                            PlatformRequirement.value.in_(value)
                        ))
                    )
                ))

            jobs = jobs.order_by(Job.priority)
            job = jobs.first()
            if job:
                internal_job = job.to_internal_job(self)
                leases = callback(internal_job)
                if leases:
                    job.assigned = True
                    job.worker_start_timestamp = internal_job.worker_start_timestamp.ToDatetime()
                for lease in leases:
                    self._create_lease(lease, session, job=internal_job)
                return leases
            return []
# Copyright (C) 2019 Bloomberg LP # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # #  <http://www.apache.org/licenses/LICENSE-2.0> # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # Set-up temporary SQLite Database when connection string is not specified # Make sure to keep this tempfile for the lifetime of this object # since this is a temporary database, we always need to create it # Disallow sqlite in-memory because multi-threaded access to it is # complex and potentially problematic at best # ref: https://docs.sqlalchemy.org/en/13/dialects/sqlite.html#threading-pooling-behavior # Only pass the (known) kwargs that have been explicitly set by the user # Valid connection_strings for in-memory SQLite which we don't support could look like: # "sqlite:///file:memdb1?option=value&cache=shared&mode=memory", # "sqlite:///file:memdb1?mode=memory&cache=shared", # "sqlite:///file:memdb1?cache=shared&mode=memory", # "sqlite:///file::memory:?cache=shared", # "sqlite:///file::memory:", # "sqlite:///:memory:", # "sqlite:///", # "sqlite://" # ref: https://www.sqlite.org/inmemorydb.html # Note that a user can also specify drivers, so prefix could become 'sqlite+driver:///' # Don't do anything. This function needs to exist but there's no # need to actually toggle monitoring in this implementation. # Don't do anything. This function needs to exist but there's no # need to actually toggle monitoring in this implementation. # Return a set of job names here for now, to match the `MemoryDataStore` # implementation's behaviour # Don't do anything. This function needs to exist but there's no # need to actually delete operations in this implementation. # `lease.job_name` is the same as `lease.id` for a Lease protobuf Return a list of leases for the highest priority jobs that can be run by a worker.          NOTE: Currently the list only ever has one or zero leases.          Query the jobs table to find queued jobs which match the capabilities of         a given worker, and return the one with the highest priority. Takes a         dictionary of worker capabilities to compare with job requirements.          :param capabilities: Dictionary of worker capabilities to compare             with job requirements when finding a job.         :type capabilities: dict         :param callback: Function to run on the next runnable job, should return             a list of leases.         :type callback: function         :param timeout: time to wait for new jobs, caps if longer             than MAX_JOB_BLOCK_TIME.         :type timeout: int         :returns: List of leases # Cap the timeout if it's larger than MAX_JOB_BLOCK_TIME # noqa # Filter to find just jobs that either don't have requirements, or # jobs with requirement keys that are a subset of the worker's # advertised capabilities keys # For each advertised worker capability, filter to find jobs which # either don't require that capability or can be satisfied by the # capability