Ejemplo n.º 1
0
class FilesystemEventLogStorage(WatchableEventLogStorage):
    def __init__(self, base_dir):
        self._base_dir = check.str_param(base_dir, 'base_dir')
        mkdir_p(self._base_dir)

        self._known_run_ids = set([])
        self._watchers = {}
        self._obs = Observer()
        self._obs.start()

    @contextmanager
    def _connect(self, run_id):
        try:
            with sqlite3.connect(self.filepath_for_run_id(run_id)) as conn:
                yield conn
        finally:
            conn.close()

    def filepath_for_run_id(self, run_id):
        check.str_param(run_id, 'run_id')
        return os.path.join(self._base_dir, '{run_id}.db'.format(run_id=run_id))

    def store_event(self, event):
        check.inst_param(event, 'event', EventRecord)
        run_id = event.run_id
        if not run_id in self._known_run_ids:
            with self._connect(run_id) as conn:
                conn.cursor().execute(CREATE_EVENT_LOG_SQL)
                self._known_run_ids.add(run_id)
        with self._connect(run_id) as conn:
            conn.cursor().execute(INSERT_EVENT_SQL, (serialize_dagster_namedtuple(event),))

    def get_logs_for_run(self, run_id, cursor=-1):
        check.str_param(run_id, 'run_id')
        check.int_param(cursor, 'cursor')
        check.invariant(
            cursor >= -1,
            'Don\'t know what to do with negative cursor {cursor}'.format(cursor=cursor),
        )

        events = []
        if not os.path.exists(self.filepath_for_run_id(run_id)):
            return events

        try:
            with self._connect(run_id) as conn:
                results = conn.cursor().execute(FETCH_EVENTS_SQL, (str(cursor),)).fetchall()
        except sqlite3.Error as err:
            six.raise_from(EventLogInvalidForRun(run_id=run_id), err)

        try:
            for (json_str,) in results:
                events.append(
                    check.inst_param(
                        deserialize_json_to_dagster_namedtuple(json_str), 'event', EventRecord
                    )
                )
        except (seven.JSONDecodeError, check.CheckError) as err:
            six.raise_from(EventLogInvalidForRun(run_id=run_id), err)

        return events

    def wipe(self):
        for filename in glob.glob(os.path.join(self._base_dir, '*.db')):
            os.unlink(filename)

    @property
    def is_persistent(self):
        return True

    def watch(self, run_id, start_cursor, callback):
        watchdog = EventLogStorageWatchdog(self, run_id, callback, start_cursor)
        self._watchers[run_id] = self._obs.schedule(watchdog, self._base_dir, True)

    def end_watch(self, run_id, handler):
        self._obs.remove_handler_for_watch(handler, self._watchers[run_id])
        del self._watchers[run_id]
Ejemplo n.º 2
0
class SqliteEventLogStorage(WatchableEventLogStorage, ConfigurableClass):
    def __init__(self, base_dir, inst_data=None):
        '''Note that idempotent initialization of the SQLite database is done on a per-run_id
        basis in the body of store_event, since each run is stored in a separate database.'''
        self._base_dir = check.str_param(base_dir, 'base_dir')
        mkdir_p(self._base_dir)

        self._known_run_ids = set([])
        self._watchers = {}
        self._obs = Observer()
        self._obs.start()
        self._inst_data = check.opt_inst_param(inst_data, 'inst_data', ConfigurableClassData)

    @property
    def inst_data(self):
        return self._inst_data

    @classmethod
    def config_type(cls):
        return SystemNamedDict('SqliteEventLogStorageConfig', {'base_dir': Field(String)})

    @staticmethod
    def from_config_value(inst_data, config_value, **kwargs):
        return SqliteEventLogStorage(inst_data=inst_data, **dict(config_value, **kwargs))

    @contextmanager
    def _connect(self, run_id):
        try:
            with sqlite3.connect(self.filepath_for_run_id(run_id)) as conn:
                yield conn
        finally:
            conn.close()

    def filepath_for_run_id(self, run_id):
        check.str_param(run_id, 'run_id')
        return os.path.join(self._base_dir, '{run_id}.db'.format(run_id=run_id))

    def store_event(self, event):
        check.inst_param(event, 'event', EventRecord)
        run_id = event.run_id
        if not run_id in self._known_run_ids:
            with self._connect(run_id) as conn:
                conn.cursor().execute(CREATE_EVENT_LOG_SQL)
                conn.cursor().execute('PRAGMA journal_mode=WAL;')
                self._known_run_ids.add(run_id)
        with self._connect(run_id) as conn:
            dagster_event_type = None
            if event.is_dagster_event:
                dagster_event_type = event.dagster_event.event_type_value

            conn.cursor().execute(
                INSERT_EVENT_SQL,
                (serialize_dagster_namedtuple(event), dagster_event_type, event.timestamp),
            )

    def get_logs_for_run(self, run_id, cursor=-1):
        check.str_param(run_id, 'run_id')
        check.int_param(cursor, 'cursor')
        check.invariant(
            cursor >= -1,
            'Don\'t know what to do with negative cursor {cursor}'.format(cursor=cursor),
        )

        events = []
        if not os.path.exists(self.filepath_for_run_id(run_id)):
            return events

        cursor += 1  # adjust from 0 based offset to 1
        try:
            with self._connect(run_id) as conn:
                results = conn.cursor().execute(FETCH_EVENTS_SQL, (str(cursor),)).fetchall()
        except sqlite3.Error as err:
            six.raise_from(EventLogInvalidForRun(run_id=run_id), err)

        try:
            for (json_str,) in results:
                events.append(
                    check.inst_param(
                        deserialize_json_to_dagster_namedtuple(json_str), 'event', EventRecord
                    )
                )
        except (seven.JSONDecodeError, check.CheckError) as err:
            six.raise_from(EventLogInvalidForRun(run_id=run_id), err)

        return events

    def get_stats_for_run(self, run_id):
        if not os.path.exists(self.filepath_for_run_id(run_id)):
            return None

        try:
            with self._connect(run_id) as conn:
                results = conn.cursor().execute(FETCH_STATS_SQL).fetchall()
        except sqlite3.Error as err:
            six.raise_from(EventLogInvalidForRun(run_id=run_id), err)

        try:
            counts = {}
            times = {}
            for result in results:
                if result[0]:
                    counts[result[0]] = result[1]
                    times[result[0]] = result[2]

            return PipelineRunStatsSnapshot(
                run_id=run_id,
                steps_succeeded=counts.get(DagsterEventType.STEP_SUCCESS.value, 0),
                steps_failed=counts.get(DagsterEventType.STEP_FAILURE.value, 0),
                materializations=counts.get(DagsterEventType.STEP_MATERIALIZATION.value, 0),
                expectations=counts.get(DagsterEventType.STEP_EXPECTATION_RESULT.value, 0),
                start_time=float(times.get(DagsterEventType.PIPELINE_START.value, 0.0)),
                end_time=float(
                    times.get(
                        DagsterEventType.PIPELINE_SUCCESS.value,
                        times.get(DagsterEventType.PIPELINE_FAILURE.value, 0.0),
                    )
                ),
            )
        except (seven.JSONDecodeError, check.CheckError) as err:
            six.raise_from(EventLogInvalidForRun(run_id=run_id), err)

    def wipe(self):
        for filename in glob.glob(os.path.join(self._base_dir, '*.db')):
            os.unlink(filename)

    def delete_events(self, run_id):
        path = self.filepath_for_run_id(run_id)
        if os.path.exists(path):
            os.unlink(path)

    @property
    def is_persistent(self):
        return True

    def watch(self, run_id, start_cursor, callback):
        watchdog = EventLogStorageWatchdog(self, run_id, callback, start_cursor)
        self._watchers[run_id] = self._obs.schedule(watchdog, self._base_dir, True)

    def end_watch(self, run_id, handler):
        self._obs.remove_handler_for_watch(handler, self._watchers[run_id])
        del self._watchers[run_id]
Ejemplo n.º 3
0
class SqliteEventLogStorage(SqlEventLogStorage, ConfigurableClass):
    """SQLite-backed event log storage.

    Users should not directly instantiate this class; it is instantiated by internal machinery when
    ``dagit`` and ``dagster-graphql`` load, based on the values in the ``dagster.yaml`` file in
    ``$DAGSTER_HOME``. Configuration of this class should be done by setting values in that file.

    This is the default event log storage when none is specified in the ``dagster.yaml``.

    To explicitly specify SQLite for event log storage, you can add a block such as the following
    to your ``dagster.yaml``:

    .. code-block:: YAML

        run_storage:
          module: dagster.core.storage.event_log
          class: SqliteEventLogStorage
          config:
            base_dir: /path/to/dir

    The ``base_dir`` param tells the event log storage where on disk to store the databases. To
    improve concurrent performance, event logs are stored in a separate SQLite database for each
    run.
    """

    def __init__(self, base_dir, inst_data=None):
        """Note that idempotent initialization of the SQLite database is done on a per-run_id
        basis in the body of connect, since each run is stored in a separate database."""
        self._base_dir = os.path.abspath(check.str_param(base_dir, "base_dir"))
        mkdir_p(self._base_dir)

        self._watchers = defaultdict(dict)
        self._obs = Observer()
        self._obs.start()
        self._inst_data = check.opt_inst_param(inst_data, "inst_data", ConfigurableClassData)

    def upgrade(self):
        all_run_ids = self.get_all_run_ids()
        print(  # pylint: disable=print-call
            "Updating event log storage for {n_runs} runs on disk...".format(
                n_runs=len(all_run_ids)
            )
        )
        alembic_config = get_alembic_config(__file__)
        for run_id in tqdm(all_run_ids):
            with self.connect(run_id) as conn:
                run_alembic_upgrade(alembic_config, conn, run_id)

    @property
    def inst_data(self):
        return self._inst_data

    @classmethod
    def config_type(cls):
        return {"base_dir": str}

    @staticmethod
    def from_config_value(inst_data, config_value):
        return SqliteEventLogStorage(inst_data=inst_data, **config_value)

    def get_all_run_ids(self):
        all_filenames = glob.glob(os.path.join(self._base_dir, "*.db"))
        return [os.path.splitext(os.path.basename(filename))[0] for filename in all_filenames]

    def path_for_run_id(self, run_id):
        return os.path.join(self._base_dir, "{run_id}.db".format(run_id=run_id))

    def conn_string_for_run_id(self, run_id):
        check.str_param(run_id, "run_id")
        return create_db_conn_string(self._base_dir, run_id)

    def _initdb(self, engine):

        alembic_config = get_alembic_config(__file__)

        try:
            SqlEventLogStorageMetadata.create_all(engine)
            engine.execute("PRAGMA journal_mode=WAL;")
            stamp_alembic_rev(alembic_config, engine)
        except (db.exc.DatabaseError, sqlite3.DatabaseError, sqlite3.OperationalError) as exc:
            # This is SQLite-specific handling for concurrency issues that can arise when, e.g.,
            # the root nodes of a pipeline execute simultaneously on Airflow with SQLite storage
            # configured and contend with each other to init the db. When we hit the following
            # errors, we know that another process is on the case and it's safe to continue:
            err_msg = str(exc)
            if not (
                "table event_logs already exists" in err_msg
                or "database is locked" in err_msg
                or "table alembic_version already exists" in err_msg
                or "UNIQUE constraint failed: alembic_version.version_num" in err_msg
            ):
                raise
            else:
                logging.info(
                    "SqliteEventLogStorage._initdb: Encountered apparent concurrent init, "
                    "swallowing {str_exc}".format(str_exc=err_msg)
                )

    @contextmanager
    def connect(self, run_id=None):
        check.str_param(run_id, "run_id")

        conn_string = self.conn_string_for_run_id(run_id)
        engine = create_engine(conn_string, poolclass=NullPool)

        if not os.path.exists(self.path_for_run_id(run_id)):
            self._initdb(engine)

        conn = engine.connect()
        try:
            with handle_schema_errors(
                conn,
                get_alembic_config(__file__),
                msg="SqliteEventLogStorage for run {run_id}".format(run_id=run_id),
            ):
                yield conn
        finally:
            conn.close()
        engine.dispose()

    def has_secondary_index(self, name, run_id=None):
        return False

    def enable_secondary_index(self, name, run_id=None):
        pass

    def wipe(self):
        for filename in (
            glob.glob(os.path.join(self._base_dir, "*.db"))
            + glob.glob(os.path.join(self._base_dir, "*.db-wal"))
            + glob.glob(os.path.join(self._base_dir, "*.db-shm"))
        ):
            os.unlink(filename)

    def watch(self, run_id, start_cursor, callback):
        watchdog = SqliteEventLogStorageWatchdog(self, run_id, callback, start_cursor)
        self._watchers[run_id][callback] = (
            watchdog,
            self._obs.schedule(watchdog, self._base_dir, True),
        )

    def end_watch(self, run_id, handler):
        if handler in self._watchers[run_id]:
            event_handler, watch = self._watchers[run_id][handler]
            self._obs.remove_handler_for_watch(event_handler, watch)
            del self._watchers[run_id][handler]
Ejemplo n.º 4
0
class ConsolidatedSqliteEventLogStorage(SqlEventLogStorage, ConfigurableClass):
    """SQLite-backed consolidated event log storage intended for test cases only.

    Users should not directly instantiate this class; it is instantiated by internal machinery when
    ``dagit`` and ``dagster-graphql`` load, based on the values in the ``dagster.yaml`` file in
    ``$DAGSTER_HOME``. Configuration of this class should be done by setting values in that file.

    To explicitly specify the consolidated SQLite for event log storage, you can add a block such as
    the following to your ``dagster.yaml``:

    .. code-block:: YAML

        run_storage:
          module: dagster.core.storage.event_log
          class: ConsolidatedSqliteEventLogStorage
          config:
            base_dir: /path/to/dir

    The ``base_dir`` param tells the event log storage where on disk to store the database.
    """

    def __init__(self, base_dir, inst_data=None):
        self._base_dir = check.str_param(base_dir, "base_dir")
        self._conn_string = create_db_conn_string(base_dir, SQLITE_EVENT_LOG_FILENAME)
        self._secondary_index_cache = {}
        self._inst_data = check.opt_inst_param(inst_data, "inst_data", ConfigurableClassData)
        self._watchdog = None
        self._watchers = defaultdict(dict)
        self._obs = Observer()
        self._obs.start()

        if not os.path.exists(self.get_db_path()):
            self._init_db()

        super().__init__()

    @property
    def inst_data(self):
        return self._inst_data

    @classmethod
    def config_type(cls):
        return {"base_dir": StringSource}

    @staticmethod
    def from_config_value(inst_data, config_value):
        return ConsolidatedSqliteEventLogStorage(inst_data=inst_data, **config_value)

    def _init_db(self):
        mkdir_p(self._base_dir)
        engine = create_engine(self._conn_string, poolclass=NullPool)
        alembic_config = get_alembic_config(__file__)

        should_mark_indexes = False
        with engine.connect() as connection:
            db_revision, head_revision = check_alembic_revision(alembic_config, connection)
            if not (db_revision and head_revision):
                SqlEventLogStorageMetadata.create_all(engine)
                engine.execute("PRAGMA journal_mode=WAL;")
                stamp_alembic_rev(alembic_config, connection)
                should_mark_indexes = True

        if should_mark_indexes:
            # mark all secondary indexes
            self.reindex_events()
            self.reindex_assets()

    @contextmanager
    def _connect(self):
        engine = create_engine(self._conn_string, poolclass=NullPool)
        conn = engine.connect()
        try:
            with handle_schema_errors(
                conn,
                get_alembic_config(__file__),
                msg="ConsolidatedSqliteEventLogStorage requires migration",
            ):
                yield conn
        finally:
            conn.close()

    def run_connection(self, run_id):
        return self._connect()

    def index_connection(self):
        return self._connect()

    def get_db_path(self):
        return os.path.join(self._base_dir, "{}.db".format(SQLITE_EVENT_LOG_FILENAME))

    def upgrade(self):
        alembic_config = get_alembic_config(__file__)
        with self._connect() as conn:
            run_alembic_upgrade(alembic_config, conn)

    def has_secondary_index(self, name):
        if name not in self._secondary_index_cache:
            self._secondary_index_cache[name] = super(
                ConsolidatedSqliteEventLogStorage, self
            ).has_secondary_index(name)
        return self._secondary_index_cache[name]

    def enable_secondary_index(self, name):
        super(ConsolidatedSqliteEventLogStorage, self).enable_secondary_index(name)
        if name in self._secondary_index_cache:
            del self._secondary_index_cache[name]

    def watch(self, run_id, start_cursor, callback):
        if not self._watchdog:
            self._watchdog = ConsolidatedSqliteEventLogStorageWatchdog(self)

        watch = self._obs.schedule(self._watchdog, self._base_dir, True)
        cursor = start_cursor if start_cursor is not None else -1
        self._watchers[run_id][callback] = (cursor, watch)

    def on_modified(self):
        keys = [
            (run_id, callback)
            for run_id, callback_dict in self._watchers.items()
            for callback, _ in callback_dict.items()
        ]
        for run_id, callback in keys:
            cursor, watch = self._watchers[run_id][callback]

            # fetch events
            events = self.get_logs_for_run(run_id, cursor)

            # update cursor
            self._watchers[run_id][callback] = (cursor + len(events), watch)

            for event in events:
                status = callback(event)
                if (
                    status == PipelineRunStatus.SUCCESS
                    or status == PipelineRunStatus.FAILURE
                    or status == PipelineRunStatus.CANCELED
                ):
                    self.end_watch(run_id, callback)

    def end_watch(self, run_id, handler):
        if run_id in self._watchers:
            _cursor, watch = self._watchers[run_id][handler]
            self._obs.remove_handler_for_watch(self._watchdog, watch)
            del self._watchers[run_id][handler]
Ejemplo n.º 5
0
class SqliteEventLogStorage(SqlEventLogStorage, ConfigurableClass):
    """SQLite-backed event log storage.

    Users should not directly instantiate this class; it is instantiated by internal machinery when
    ``dagit`` and ``dagster-graphql`` load, based on the values in the ``dagster.yaml`` file in
    ``$DAGSTER_HOME``. Configuration of this class should be done by setting values in that file.

    This is the default event log storage when none is specified in the ``dagster.yaml``.

    To explicitly specify SQLite for event log storage, you can add a block such as the following
    to your ``dagster.yaml``:

    .. code-block:: YAML

        event_log_storage:
          module: dagster.core.storage.event_log
          class: SqliteEventLogStorage
          config:
            base_dir: /path/to/dir

    The ``base_dir`` param tells the event log storage where on disk to store the databases. To
    improve concurrent performance, event logs are stored in a separate SQLite database for each
    run.
    """
    def __init__(self, base_dir, inst_data=None):
        """Note that idempotent initialization of the SQLite database is done on a per-run_id
        basis in the body of connect, since each run is stored in a separate database."""
        self._base_dir = os.path.abspath(check.str_param(base_dir, "base_dir"))
        mkdir_p(self._base_dir)

        self._obs = None

        self._watchers = defaultdict(dict)
        self._inst_data = check.opt_inst_param(inst_data, "inst_data",
                                               ConfigurableClassData)

        # Used to ensure that each run ID attempts to initialize its DB the first time it connects,
        # ensuring that the database will be created if it doesn't exist
        self._initialized_dbs = set()

        # Ensure that multiple threads (like the event log watcher) interact safely with each other
        self._db_lock = threading.Lock()

        if not os.path.exists(self.path_for_shard(INDEX_SHARD_NAME)):
            conn_string = self.conn_string_for_shard(INDEX_SHARD_NAME)
            engine = create_engine(conn_string, poolclass=NullPool)
            self._initdb(engine)
            self.reindex()

        super().__init__()

    def upgrade(self):
        all_run_ids = self.get_all_run_ids()
        print(  # pylint: disable=print-call
            f"Updating event log storage for {len(all_run_ids)} runs on disk..."
        )
        alembic_config = get_alembic_config(__file__)
        if all_run_ids:
            for run_id in tqdm(all_run_ids):
                with self.run_connection(run_id) as conn:
                    run_alembic_upgrade(alembic_config, conn, run_id)

        print("Updating event log storage for index db on disk...")  # pylint: disable=print-call
        with self.index_connection() as conn:
            run_alembic_upgrade(alembic_config, conn, "index")

        self._initialized_dbs = set()

    @property
    def inst_data(self):
        return self._inst_data

    @classmethod
    def config_type(cls):
        return {"base_dir": StringSource}

    @staticmethod
    def from_config_value(inst_data, config_value):
        return SqliteEventLogStorage(inst_data=inst_data, **config_value)

    def get_all_run_ids(self):
        all_filenames = glob.glob(os.path.join(self._base_dir, "*.db"))
        return [
            os.path.splitext(os.path.basename(filename))[0]
            for filename in all_filenames if
            os.path.splitext(os.path.basename(filename))[0] != INDEX_SHARD_NAME
        ]

    def path_for_shard(self, run_id):
        return os.path.join(self._base_dir,
                            "{run_id}.db".format(run_id=run_id))

    def conn_string_for_shard(self, shard_name):
        check.str_param(shard_name, "shard_name")
        return create_db_conn_string(self._base_dir, shard_name)

    def _initdb(self, engine):
        alembic_config = get_alembic_config(__file__)

        retry_limit = 10

        while True:
            try:

                with engine.connect() as connection:
                    db_revision, head_revision = check_alembic_revision(
                        alembic_config, connection)

                    if not (db_revision and head_revision):
                        SqlEventLogStorageMetadata.create_all(engine)
                        engine.execute("PRAGMA journal_mode=WAL;")
                        stamp_alembic_rev(alembic_config, connection)

                break
            except (db.exc.DatabaseError, sqlite3.DatabaseError,
                    sqlite3.OperationalError) as exc:
                # This is SQLite-specific handling for concurrency issues that can arise when
                # multiple processes (e.g. the dagit process and user code process) contend with
                # each other to init the db. When we hit the following errors, we know that another
                # process is on the case and we should retry.
                err_msg = str(exc)

                if not ("table asset_keys already exists" in err_msg
                        or "table secondary_indexes already exists" in err_msg
                        or "table event_logs already exists" in err_msg
                        or "database is locked" in err_msg
                        or "table alembic_version already exists" in err_msg or
                        "UNIQUE constraint failed: alembic_version.version_num"
                        in err_msg):
                    raise

                if retry_limit == 0:
                    raise
                else:
                    logging.info(
                        "SqliteEventLogStorage._initdb: Encountered apparent concurrent init, "
                        "retrying ({retry_limit} retries left). Exception: {str_exc}"
                        .format(retry_limit=retry_limit, str_exc=err_msg))
                    time.sleep(0.2)
                    retry_limit -= 1

    @contextmanager
    def _connect(self, shard):
        with self._db_lock:
            check.str_param(shard, "shard")

            conn_string = self.conn_string_for_shard(shard)
            engine = create_engine(conn_string, poolclass=NullPool)

            if not shard in self._initialized_dbs:
                self._initdb(engine)
                self._initialized_dbs.add(shard)

            conn = engine.connect()

            try:
                with handle_schema_errors(
                        conn,
                        get_alembic_config(__file__),
                        msg="SqliteEventLogStorage for shard {shard}".format(
                            shard=shard),
                ):
                    yield conn
            finally:
                conn.close()
            engine.dispose()

    def run_connection(self, run_id=None):
        return self._connect(run_id)

    def index_connection(self):
        return self._connect(INDEX_SHARD_NAME)

    def store_event(self, event):
        """
        Overridden method to replicate asset events in a central assets.db sqlite shard, enabling
        cross-run asset queries.

        Args:
            event (EventRecord): The event to store.
        """
        check.inst_param(event, "event", EventRecord)
        insert_event_statement = self.prepare_insert_event(event)
        run_id = event.run_id

        with self.run_connection(run_id) as conn:
            conn.execute(insert_event_statement)

        if event.is_dagster_event and event.dagster_event.asset_key:
            # mirror the event in the cross-run index database
            with self.index_connection() as conn:
                conn.execute(insert_event_statement)

            self.store_asset_key(event)

    def delete_events(self, run_id):
        with self.run_connection(run_id) as conn:
            self.delete_events_for_run(conn, run_id)

        # delete the mirrored event in the cross-run index database
        with self.index_connection() as conn:
            self.delete_events_for_run(conn, run_id)

    def wipe(self):
        # should delete all the run-sharded dbs as well as the index db
        for filename in (glob.glob(os.path.join(self._base_dir, "*.db")) +
                         glob.glob(os.path.join(self._base_dir, "*.db-wal")) +
                         glob.glob(os.path.join(self._base_dir, "*.db-shm"))):
            os.unlink(filename)

        self._initialized_dbs = set()

    def _delete_mirrored_events_for_asset_key(self, asset_key):
        with self.index_connection() as conn:
            conn.execute(SqlEventLogStorageTable.delete().where(  # pylint: disable=no-value-for-parameter
                db.or_(
                    SqlEventLogStorageTable.c.asset_key ==
                    asset_key.to_string(),
                    SqlEventLogStorageTable.c.asset_key == asset_key.to_string(
                        legacy=True),
                )))

    def wipe_asset(self, asset_key):
        # default implementation will update the event_logs in the sharded dbs, and the asset_key
        # table in the asset shard, but will not remove the mirrored event_log events in the asset
        # shard
        super(SqliteEventLogStorage, self).wipe_asset(asset_key)
        self._delete_mirrored_events_for_asset_key(asset_key)

    def watch(self, run_id, start_cursor, callback):
        if not self._obs:
            self._obs = Observer()
            self._obs.start()

        watchdog = SqliteEventLogStorageWatchdog(self, run_id, callback,
                                                 start_cursor)
        self._watchers[run_id][callback] = (
            watchdog,
            self._obs.schedule(watchdog, self._base_dir, True),
        )

    def end_watch(self, run_id, handler):
        if handler in self._watchers[run_id]:
            event_handler, watch = self._watchers[run_id][handler]
            self._obs.remove_handler_for_watch(event_handler, watch)
            del self._watchers[run_id][handler]

    def dispose(self):
        if self._obs:
            self._obs.stop()
            self._obs.join(timeout=15)
Ejemplo n.º 6
0
class Schedule:
    def __init__(self, config):
        self.observer = Observer()
        self.handlers = {}
        self.watchers = {}
        self.counter = Counter()
        self.notifier = Notifier(config)
        self.offset_db = OffsetPersistence(config)

    def __make_key(self, filename):
        return path.abspath(urlsafe_b64decode(filename).decode())

    def add_watcher(self, filename):
        #如果监控的文件都在同一个目录下,下面这种方法只会启动一个inotify来进行监控这个目录下的所有文件,handler.start实际上调用了monitor的start
        filename = self.__make_key(filename)
        if handler.filename not in self.handlers.keys():
            handler = Watcher(filename, counter=self.counter, notifier=self.notifier, offset_db=self.offset_db)
            if path.dirname(handler.filename) not in self.watchers.keys():
                self.watchers[path.dirname(handler.filename)] = self.observer.schedule(handler, path.dirname(handler.filename), recursive=False)
            else:
                watch = self.watchers[path.dirname(handler.filename)]
                self.observer.add_handler_for_watch(handler, watch)
            self.handlers[handler.filename] = handler
            handler.start()


    def remove_watcher(self, filename):
        #判断key是否在watchers里面,有则从observer中剔除,然后从watchers字典内删除,并把handlers剔除并stop
        key = self.__make_key(filename)
        handler = self.handlers.pop(key)
        if handler is not None:
            watch = self.watchers[path.dirname(key)]
            self.observer.remove_handler_for_watch(handler, watch)
            handler.stop()
            if not self.observer._handlers[watch]:
                self.observer.unschedule(watch)
                self.watchers.pop(path.dirname(handler.filename))


    #添加和移除monitor,实质上就是调用Monitor类的add和remove方法来操作
    def add_monitor(self, filename, name, src):
        key = self.__make_key(filename)
        handler = self.handlers.get(key)
        if handler is None:
            logging.warning('watcher {0} not found, auto add it'.format(filename))
            handler.add_watcher(filename)
            handler = self.handlers.get(key)
        handler.monitor.add(filename, name, src)

    def remove_monitor(self, filename, name):
        key = self.__make_key(filename)
        handler = self.handlers.get(key)
        if handler is None:
            logging.warning('watcher {0} not found'.format(filename))
            return
        handler.monitor.remove(name)

    def start(self):
        self.observer.start()
        self.notifier.start()

    def join(self):
        self.observer.join()

    def stop(self):
        self.observer.stop()
        for handler in self.handlers.values():
            handler.stop()
        self.notifier.stop()
        self.offset_db.close()
Ejemplo n.º 7
0
class WatchdogTests(unittest.TestCase):

    def setUp(self):
        self.events = queue.Queue()
        self.observer = Observer()
        self.handler = EventHandler(self.events, channel_id=5)
        self.watch = self.observer.schedule(event_handler=self.handler, path=TMP_DIR)
        self.observer.start()

    def testEventChannelID(self):
        with open(TMP_FILE, "w") as local_file:
            local_file.write(string_common)
        event = self.events.get(block=True, timeout=1)
        self.assertEqual(event.channel_id, 5)

    # Modifying the file should add a new task to the queue
    def testTaskInsertion(self):
        new_string = "This is a whole new line in file "+filename
        with open (TMP_FILE, "w") as local_file:
            local_file.write(new_string)
        while True:
            try:
                task = self.events.get(block=True, timeout=1)
            except queue.Empty:
                break
            self.assertEqual(task.src_path,local_dir+filename)
            self.assertEqual(task.dest_path, remote_dir+filename)

    # It should be possible to add a new Handler to an already scheduled watch
    def testHandlerAdding(self):
        fake_handle="I am a fake handler, don't mind me!"
        class FakeHandler(EventHandler):
            def handle_modification(self, event):
                self.tasks.put(fake_handle)
        fake_queue = queue.Queue()
        fake_handler = FakeHandler(fake_queue)
        self.observer.add_handler_for_watch(fake_handler, self.watch)
        new_string = "This is a whole new line in file " + filename
        with open(local_dir + filename, "w") as local_file:
            local_file.write(new_string)
        while True:
            try:
                task = fake_queue.get(block=True, timeout=1)
            except queue.Empty:
                break
            self.assertEqual(task, fake_handle)

    # A Handler blocking on one event (ex: inserting into a busy queue)
    # should not prevent handling of further events
    def testHandlerHanging(self):
        class HangingHandler(EventHandler):
            def on_any_event(self, event):
                print("Handler hanging...")
                time.sleep(1)
                print("Handler dispatching")
                super().on_any_event(event)

        self.observer.remove_handler_for_watch(self.handler, self.watch)
        slow_handler = HangingHandler(self.events)
        self.observer.add_handler_for_watch(slow_handler, self.watch)
        for i in range(5):
            with open(os.path.join(TMP_DIR, "f"+str(i)), "w") as local_file:
                local_file.write("Write #"+str(i))
            time.sleep(0.1)
        time.sleep(6)
        num = len(empty_tasks(self.events))
        print("[Hanging] "+str(num)+" events")
        self.assertTrue(num >= 5)

    # Scheduling a new watch while another one is running
    # In this test, each write should set off 2 events (open+close) as seen on the next test
    # However watchdog is nice enough to avoid adding similar events to its internal queue
    def testNewScheduling(self):
        self.assertTrue(self.events.empty()) # Queue starts off
        with open(remote_dir + filename, "w") as remote_file:
            remote_file.write("Going once")
        self.assertRaises(queue.Empty, self.events.get, timeout=1) # No Handler is watching the file yet

        handler2 = EventHandler(self.events)
        self.observer.schedule(event_handler=handler2, path=remote_dir, recursive=True)

        with open(remote_dir + filename, "w") as remote_file:
            remote_file.write("Going twice")
        l1 = empty_tasks(self.events) # Now it should work
        self.assertEqual(len(l1), 1) # Single event
        for task in l1:
            self.assertEqual(task.src_path, remote_dir + filename)
            self.assertEqual(task.dest_path, local_dir + filename)

        with open(local_dir + filename, "w") as local_file:
            local_file.write("Going thrice") # Writing to the local file still works
        l2 = empty_tasks(self.events)
        self.assertEqual(len(l2), 1)  # Single event
        for task in l2:
            self.assertEqual(task.src_path, local_dir + filename)
            self.assertEqual(task.dest_path, remote_dir + filename)

    # It should be possible to remove a scheduled watch
    def testWatchRemoval(self):
        handler2 = EventHandler(self.events)
        watch2 = self.observer.schedule(event_handler=handler2, path=remote_dir, recursive=True)
        for client in [ {"path":local_dir+filename, "watch":self.watch},
                      {"path":remote_dir+filename, "watch":watch2} ]:
            with open(client["path"], "w") as file:
                file.write("This will make an event")
            time.sleep(0.5)
            task = self.events.get(timeout=1)
            self.assertEquals(task.src_path, client["path"])

            self.observer.unschedule(client["watch"])
            with open(local_dir+filename, "w") as local_file:
                local_file.write("This won't")
            self.assertRaises(queue.Empty, self.events.get, timeout=1)

    # Each open() and each close() should produce an event
    # They are sometimes squashed into a single event if done
    # Quickly enough (i.e. "with open(file) as f: f.write()")
    def testEventsPerWrite(self):
        local_file = open(local_dir+filename, "w")
        self.assertTrue(len(empty_tasks(self.events)) == 1) # Opening sets off an event
        local_file.write("First")
        self.assertTrue(len(empty_tasks(self.events)) == 0) # Writing doesn't set off an event
        local_file.write("Second")
        self.assertTrue(len(empty_tasks(self.events)) == 0) # Writing doesn't set off an event
        local_file.close()
        self.assertTrue(len(empty_tasks(self.events)) == 1) # Closing sets off an event

    def testEventsForBigFileCopy(self):
        self.observer.schedule(self.handler, TMP_DIR, recursive=True)
        try:
            os.mkdir(TMP_DIR)
        except FileExistsError:
            pass
        try:
            shutil.copy(BIG_FILE, BIG_IN)
        except OSError:
            self.fail()
        num = empty_tasks(self.events)
        print("Used "+str(len(num))+" events")
Ejemplo n.º 8
0
class SqliteEventLogStorage(SqlEventLogStorage, ConfigurableClass):
    """SQLite-backed event log storage.

    Users should not directly instantiate this class; it is instantiated by internal machinery when
    ``dagit`` and ``dagster-graphql`` load, based on the values in the ``dagster.yaml`` file in
    ``$DAGSTER_HOME``. Configuration of this class should be done by setting values in that file.

    This is the default event log storage when none is specified in the ``dagster.yaml``.

    To explicitly specify SQLite for event log storage, you can add a block such as the following
    to your ``dagster.yaml``:

    .. code-block:: YAML

        event_log_storage:
          module: dagster.core.storage.event_log
          class: SqliteEventLogStorage
          config:
            base_dir: /path/to/dir

    The ``base_dir`` param tells the event log storage where on disk to store the databases. To
    improve concurrent performance, event logs are stored in a separate SQLite database for each
    run.
    """
    def __init__(self, base_dir, inst_data=None):
        """Note that idempotent initialization of the SQLite database is done on a per-run_id
        basis in the body of connect, since each run is stored in a separate database."""
        self._base_dir = os.path.abspath(check.str_param(base_dir, "base_dir"))
        mkdir_p(self._base_dir)

        self._obs = None

        self._watchers = defaultdict(dict)
        self._inst_data = check.opt_inst_param(inst_data, "inst_data",
                                               ConfigurableClassData)

        # Used to ensure that each run ID attempts to initialize its DB the first time it connects,
        # ensuring that the database will be created if it doesn't exist
        self._initialized_dbs = set()

        # Ensure that multiple threads (like the event log watcher) interact safely with each other
        self._db_lock = threading.Lock()

        if not os.path.exists(self.path_for_shard(INDEX_SHARD_NAME)):
            conn_string = self.conn_string_for_shard(INDEX_SHARD_NAME)
            engine = create_engine(conn_string, poolclass=NullPool)
            self._initdb(engine)
            self.reindex_events()
            self.reindex_assets()

        super().__init__()

    def upgrade(self):
        all_run_ids = self.get_all_run_ids()
        print(  # pylint: disable=print-call
            f"Updating event log storage for {len(all_run_ids)} runs on disk..."
        )
        alembic_config = get_alembic_config(__file__)
        if all_run_ids:
            for run_id in tqdm(all_run_ids):
                with self.run_connection(run_id) as conn:
                    run_alembic_upgrade(alembic_config, conn, run_id)

        print("Updating event log storage for index db on disk...")  # pylint: disable=print-call
        with self.index_connection() as conn:
            run_alembic_upgrade(alembic_config, conn, "index")

        self._initialized_dbs = set()

    @property
    def inst_data(self):
        return self._inst_data

    @classmethod
    def config_type(cls):
        return {"base_dir": StringSource}

    @staticmethod
    def from_config_value(inst_data, config_value):
        return SqliteEventLogStorage(inst_data=inst_data, **config_value)

    def get_all_run_ids(self):
        all_filenames = glob.glob(os.path.join(self._base_dir, "*.db"))
        return [
            os.path.splitext(os.path.basename(filename))[0]
            for filename in all_filenames if
            os.path.splitext(os.path.basename(filename))[0] != INDEX_SHARD_NAME
        ]

    def path_for_shard(self, run_id):
        return os.path.join(self._base_dir,
                            "{run_id}.db".format(run_id=run_id))

    def conn_string_for_shard(self, shard_name):
        check.str_param(shard_name, "shard_name")
        return create_db_conn_string(self._base_dir, shard_name)

    def _initdb(self, engine):
        alembic_config = get_alembic_config(__file__)

        retry_limit = 10

        while True:
            try:

                with engine.connect() as connection:
                    db_revision, head_revision = check_alembic_revision(
                        alembic_config, connection)

                    if not (db_revision and head_revision):
                        SqlEventLogStorageMetadata.create_all(engine)
                        engine.execute("PRAGMA journal_mode=WAL;")
                        stamp_alembic_rev(alembic_config, connection)

                break
            except (db.exc.DatabaseError, sqlite3.DatabaseError,
                    sqlite3.OperationalError) as exc:
                # This is SQLite-specific handling for concurrency issues that can arise when
                # multiple processes (e.g. the dagit process and user code process) contend with
                # each other to init the db. When we hit the following errors, we know that another
                # process is on the case and we should retry.
                err_msg = str(exc)

                if not ("table asset_keys already exists" in err_msg
                        or "table secondary_indexes already exists" in err_msg
                        or "table event_logs already exists" in err_msg
                        or "database is locked" in err_msg
                        or "table alembic_version already exists" in err_msg or
                        "UNIQUE constraint failed: alembic_version.version_num"
                        in err_msg):
                    raise

                if retry_limit == 0:
                    raise
                else:
                    logging.info(
                        "SqliteEventLogStorage._initdb: Encountered apparent concurrent init, "
                        "retrying ({retry_limit} retries left). Exception: {str_exc}"
                        .format(retry_limit=retry_limit, str_exc=err_msg))
                    time.sleep(0.2)
                    retry_limit -= 1

    @contextmanager
    def _connect(self, shard):
        with self._db_lock:
            check.str_param(shard, "shard")

            conn_string = self.conn_string_for_shard(shard)
            engine = create_engine(conn_string, poolclass=NullPool)

            if not shard in self._initialized_dbs:
                self._initdb(engine)
                self._initialized_dbs.add(shard)

            conn = engine.connect()

            try:
                with handle_schema_errors(conn, get_alembic_config(__file__)):
                    yield conn
            finally:
                conn.close()
            engine.dispose()

    def run_connection(self, run_id=None):
        return self._connect(run_id)

    def index_connection(self):
        return self._connect(INDEX_SHARD_NAME)

    def store_event(self, event):
        """
        Overridden method to replicate asset events in a central assets.db sqlite shard, enabling
        cross-run asset queries.

        Args:
            event (EventLogEntry): The event to store.
        """
        check.inst_param(event, "event", EventLogEntry)
        insert_event_statement = self.prepare_insert_event(event)
        run_id = event.run_id

        with self.run_connection(run_id) as conn:
            conn.execute(insert_event_statement)

        if event.is_dagster_event and event.dagster_event.asset_key:
            check.invariant(
                event.dagster_event_type
                == DagsterEventType.ASSET_MATERIALIZATION or
                event.dagster_event_type == DagsterEventType.ASSET_OBSERVATION
                or event.dagster_event_type
                == DagsterEventType.ASSET_MATERIALIZATION_PLANNED,
                "Can only store asset materializations, materialization_planned, and observations in index database",
            )

            # mirror the event in the cross-run index database
            with self.index_connection() as conn:
                conn.execute(insert_event_statement)

            if (event.dagster_event.is_step_materialization
                    or event.dagster_event.is_asset_observation
                    or event.dagster_event.is_asset_materialization_planned):
                self.store_asset_event(event)

    def get_event_records(
        self,
        event_records_filter: Optional[EventRecordsFilter] = None,
        limit: Optional[int] = None,
        ascending: bool = False,
    ) -> Iterable[EventLogRecord]:
        """Overridden method to enable cross-run event queries in sqlite.

        The record id in sqlite does not auto increment cross runs, so instead of fetching events
        after record id, we only fetch events whose runs updated after update_timestamp.
        """
        check.opt_inst_param(event_records_filter, "event_records_filter",
                             EventRecordsFilter)
        check.opt_int_param(limit, "limit")
        check.bool_param(ascending, "ascending")

        is_asset_query = event_records_filter and (
            event_records_filter.event_type
            == DagsterEventType.ASSET_MATERIALIZATION
            or event_records_filter.event_type
            == DagsterEventType.ASSET_OBSERVATION)
        if is_asset_query:
            # asset materializations and observations get mirrored into the index shard, so no
            # custom run shard-aware cursor logic needed
            return super(SqliteEventLogStorage, self).get_event_records(
                event_records_filter=event_records_filter,
                limit=limit,
                ascending=ascending)

        query = db.select(
            [SqlEventLogStorageTable.c.id, SqlEventLogStorageTable.c.event])
        if event_records_filter and event_records_filter.asset_key:
            asset_details = next(
                iter(self._get_assets_details([event_records_filter.asset_key
                                               ])))
        else:
            asset_details = None

        if not event_records_filter or not (isinstance(
                event_records_filter.after_cursor, RunShardedEventsCursor)):
            warnings.warn("""
                Called `get_event_records` on a run-sharded event log storage with a query that
                is not run aware (e.g. not using a RunShardedEventsCursor).  This likely has poor
                performance characteristics.  Consider adding a RunShardedEventsCursor to your query
                or switching your instance configuration to use a non-run sharded event log storage
                (e.g. PostgresEventLogStorage, ConsolidatedSqliteEventLogStorage)
            """)

        query = self._apply_filter_to_query(
            query=query,
            event_records_filter=event_records_filter,
            asset_details=asset_details,
            apply_cursor_filters=
            False,  # run-sharded cursor filters don't really make sense
        )
        if limit:
            query = query.limit(limit)
        if ascending:
            query = query.order_by(SqlEventLogStorageTable.c.timestamp.asc())
        else:
            query = query.order_by(SqlEventLogStorageTable.c.timestamp.desc())

        # workaround for the run-shard sqlite to enable cross-run queries: get a list of run_ids
        # whose events may qualify the query, and then open run_connection per run_id at a time.
        run_updated_after = (
            event_records_filter.after_cursor.run_updated_after
            if event_records_filter and isinstance(
                event_records_filter.after_cursor, RunShardedEventsCursor) else
            None)
        run_records = self._instance.get_run_records(
            filters=RunsFilter(updated_after=run_updated_after),
            order_by="update_timestamp",
            ascending=ascending,
        )

        event_records = []
        for run_record in run_records:
            run_id = run_record.pipeline_run.run_id
            with self.run_connection(run_id) as conn:
                results = conn.execute(query).fetchall()

            for row_id, json_str in results:
                try:
                    event_record = deserialize_json_to_dagster_namedtuple(
                        json_str)
                    if not isinstance(event_record, EventLogEntry):
                        logging.warning(
                            "Could not resolve event record as EventLogEntry for id `{}`."
                            .format(row_id))
                        continue
                    else:
                        event_records.append(
                            EventLogRecord(storage_id=row_id,
                                           event_log_entry=event_record))
                    if limit and len(event_records) >= limit:
                        break
                except seven.JSONDecodeError:
                    logging.warning(
                        "Could not parse event record id `{}`.".format(row_id))

            if limit and len(event_records) >= limit:
                break

        return event_records[:limit]

    def delete_events(self, run_id):
        with self.run_connection(run_id) as conn:
            self.delete_events_for_run(conn, run_id)

        # delete the mirrored event in the cross-run index database
        with self.index_connection() as conn:
            self.delete_events_for_run(conn, run_id)

    def wipe(self):
        # should delete all the run-sharded dbs as well as the index db
        for filename in (glob.glob(os.path.join(self._base_dir, "*.db")) +
                         glob.glob(os.path.join(self._base_dir, "*.db-wal")) +
                         glob.glob(os.path.join(self._base_dir, "*.db-shm"))):
            os.unlink(filename)

        self._initialized_dbs = set()

    def _delete_mirrored_events_for_asset_key(self, asset_key):
        with self.index_connection() as conn:
            conn.execute(SqlEventLogStorageTable.delete().where(  # pylint: disable=no-value-for-parameter
                db.or_(
                    SqlEventLogStorageTable.c.asset_key ==
                    asset_key.to_string(),
                    SqlEventLogStorageTable.c.asset_key == asset_key.to_string(
                        legacy=True),
                )))

    def wipe_asset(self, asset_key):
        # default implementation will update the event_logs in the sharded dbs, and the asset_key
        # table in the asset shard, but will not remove the mirrored event_log events in the asset
        # shard
        super(SqliteEventLogStorage, self).wipe_asset(asset_key)
        self._delete_mirrored_events_for_asset_key(asset_key)

    def watch(self, run_id, start_cursor, callback):
        if not self._obs:
            self._obs = Observer()
            self._obs.start()

        watchdog = SqliteEventLogStorageWatchdog(self, run_id, callback,
                                                 start_cursor)
        self._watchers[run_id][callback] = (
            watchdog,
            self._obs.schedule(watchdog, self._base_dir, True),
        )

    def end_watch(self, run_id, handler):
        if handler in self._watchers[run_id]:
            event_handler, watch = self._watchers[run_id][handler]
            self._obs.remove_handler_for_watch(event_handler, watch)
            del self._watchers[run_id][handler]

    def dispose(self):
        if self._obs:
            self._obs.stop()
            self._obs.join(timeout=15)
Ejemplo n.º 9
0
class SqliteEventLogStorage(SqlEventLogStorage, ConfigurableClass):
    def __init__(self, base_dir, inst_data=None):
        '''Note that idempotent initialization of the SQLite database is done on a per-run_id
        basis in the body of connect, since each run is stored in a separate database.'''
        self._base_dir = os.path.abspath(check.str_param(base_dir, 'base_dir'))
        mkdir_p(self._base_dir)

        self._watchers = defaultdict(dict)
        self._obs = Observer()
        self._obs.start()
        self._inst_data = check.opt_inst_param(inst_data, 'inst_data',
                                               ConfigurableClassData)

    def upgrade(self):
        all_run_ids = self.get_all_run_ids()
        print('Updating event log storage for {n_runs} runs on disk...'.format(
            n_runs=len(all_run_ids)))
        alembic_config = get_alembic_config(__file__)
        for run_id in all_run_ids:
            with self.connect(run_id) as conn:
                run_alembic_upgrade(alembic_config, conn, run_id)

    @property
    def inst_data(self):
        return self._inst_data

    @classmethod
    def config_type(cls):
        return {'base_dir': str}

    @staticmethod
    def from_config_value(inst_data, config_value):
        return SqliteEventLogStorage(inst_data=inst_data, **config_value)

    def get_all_run_ids(self):
        all_filenames = glob.glob(os.path.join(self._base_dir, '*.db'))
        return [
            os.path.splitext(os.path.basename(filename))[0]
            for filename in all_filenames
        ]

    def path_for_run_id(self, run_id):
        return os.path.join(self._base_dir,
                            '{run_id}.db'.format(run_id=run_id))

    def conn_string_for_run_id(self, run_id):
        check.str_param(run_id, 'run_id')
        return 'sqlite:///{}'.format('/'.join(
            self.path_for_run_id(run_id).split(os.sep)))

    def _initdb(self, engine, run_id):

        alembic_config = get_alembic_config(__file__)

        try:
            SqlEventLogStorageMetadata.create_all(engine)
            engine.execute('PRAGMA journal_mode=WAL;')
            stamp_alembic_rev(alembic_config, engine)
        except (db.exc.DatabaseError, sqlite3.DatabaseError,
                sqlite3.OperationalError) as exc:
            # This is SQLite-specific handling for concurrency issues that can arise when, e.g.,
            # the root nodes of a pipeline execute simultaneously on Airflow with SQLite storage
            # configured and contend with each other to init the db. When we hit the following
            # errors, we know that another process is on the case and it's safe to continue:
            if not ('table event_logs already exists' in str(exc)
                    or 'database is locked' in str(exc)
                    or 'table alembic_version already exists' in str(exc)):
                raise
            else:
                logging.info(
                    'SqliteEventLogStorage._initdb: Encountered apparent concurrent init, '
                    'swallowing {str_exc}'.format(str_exc=str(exc)))

    @contextmanager
    def connect(self, run_id=None):
        check.str_param(run_id, 'run_id')

        conn_string = self.conn_string_for_run_id(run_id)
        engine = create_engine(conn_string, poolclass=NullPool)

        if not os.path.exists(self.path_for_run_id(run_id)):
            self._initdb(engine, run_id)

        conn = engine.connect()
        try:
            with handle_schema_errors(
                    conn,
                    get_alembic_config(__file__),
                    msg='SqliteEventLogStorage for run {run_id}'.format(
                        run_id=run_id),
            ):
                yield conn
        finally:
            conn.close()
        engine.dispose()

    def wipe(self):
        for filename in (glob.glob(os.path.join(self._base_dir, '*.db')) +
                         glob.glob(os.path.join(self._base_dir, '*.db-wal')) +
                         glob.glob(os.path.join(self._base_dir, '*.db-shm'))):
            os.unlink(filename)

    def watch(self, run_id, start_cursor, callback):
        watchdog = SqliteEventLogStorageWatchdog(self, run_id, callback,
                                                 start_cursor)
        self._watchers[run_id][callback] = (
            watchdog,
            self._obs.schedule(watchdog, self._base_dir, True),
        )

    def end_watch(self, run_id, handler):
        if handler in self._watchers[run_id]:
            event_handler, watch = self._watchers[run_id][handler]
            self._obs.remove_handler_for_watch(event_handler, watch)
            del self._watchers[run_id][handler]
Ejemplo n.º 10
0
class FilesystemEventLogStorage(WatchableEventLogStorage):
    def __init__(self, base_dir):
        self._base_dir = check.str_param(base_dir, 'base_dir')
        mkdir_p(self._base_dir)
        self.file_cursors = defaultdict(lambda: (0, 0))
        # Swap these out to use lockfiles
        self.file_lock = defaultdict(gevent.lock.Semaphore)
        self._watchers = {}
        self._obs = Observer()
        self._obs.start()

    def filepath_for_run_id(self, run_id):
        return os.path.join(self._base_dir,
                            '{run_id}.log'.format(run_id=run_id))

    def store_event(self, event):
        run_id = event.run_id

        with self.file_lock[run_id]:
            # Going to do the less error-prone, simpler, but slower strategy:
            # open, append, close for every log message for now.
            # Open the file for binary content and create if it doesn't exist.
            with open(self.filepath_for_run_id(run_id), 'ab') as file_handle:
                file_handle.seek(0, os.SEEK_END)
                pickle.dump(event, file_handle)

    def wipe(self):
        for filename in glob.glob(os.path.join(self._base_dir, '*.log')):
            os.unlink(filename)

        self.file_lock = defaultdict(gevent.lock.Semaphore)
        self.file_cursors = defaultdict(lambda: (0, 0))

    @property
    def is_persistent(self):
        return True

    def logs_ready(self, run_id):
        return os.path.exists(self.filepath_for_run_id(run_id))

    def get_logs_for_run(self, run_id, cursor=-1):
        events = []
        cursor = int(cursor) + 1
        with self.file_lock[run_id]:
            with open(self.filepath_for_run_id(run_id), 'rb') as fd:
                # There might be a path to make this more performant, at the expense of interop,
                # by using a modified file format: https://stackoverflow.com/a/8936927/324449
                # Alternatively, we could use .jsonl and linecache instead of pickle
                if cursor == self.file_cursors[run_id][0]:
                    fd.seek(self.file_cursors[run_id][1])
                else:
                    i = 0
                    while i < cursor:
                        pickle.load(fd)
                        i += 1
                    self.file_cursors[run_id] = (cursor, fd.tell())
                try:
                    while True:
                        events.append(pickle.load(fd))
                except EOFError:
                    pass
        return events

    def watch(self, run_id, start_cursor, callback):
        watchdog = EventLogStorageWatchdog(self, run_id, callback,
                                           start_cursor)
        self._watchers[run_id] = self._obs.schedule(watchdog, self._base_dir,
                                                    True)

    def end_watch(self, run_id, handler):
        self._obs.remove_handler_for_watch(handler, self._watchers[run_id])
        del self._watchers[run_id]
Ejemplo n.º 11
0
class SqliteEventLogStorage(SqlEventLogStorage, ConfigurableClass):
    def __init__(self, base_dir, inst_data=None):
        '''Note that idempotent initialization of the SQLite database is done on a per-run_id
        basis in the body of connect, since each run is stored in a separate database.'''
        self._base_dir = os.path.abspath(check.str_param(base_dir, 'base_dir'))
        mkdir_p(self._base_dir)

        self._watchers = defaultdict(dict)
        self._obs = Observer()
        self._obs.start()
        self._inst_data = check.opt_inst_param(inst_data, 'inst_data', ConfigurableClassData)

    def upgrade(self):
        all_run_ids = self.get_all_run_ids()
        print(
            'Updating event log storage for {n_runs} runs on disk...'.format(
                n_runs=len(all_run_ids)
            )
        )
        alembic_config = get_alembic_config(__file__)
        for run_id in all_run_ids:
            with self.connect(run_id) as conn:
                run_alembic_upgrade(alembic_config, conn, run_id)

    @property
    def inst_data(self):
        return self._inst_data

    @classmethod
    def config_type(cls):
        return {'base_dir': str}

    @staticmethod
    def from_config_value(inst_data, config_value):
        return SqliteEventLogStorage(inst_data=inst_data, **config_value)

    def get_all_run_ids(self):
        all_filenames = glob.glob(os.path.join(self._base_dir, '*.db'))
        return [os.path.splitext(os.path.basename(filename))[0] for filename in all_filenames]

    def path_for_run_id(self, run_id):
        return os.path.join(self._base_dir, '{run_id}.db'.format(run_id=run_id))

    def conn_string_for_run_id(self, run_id):
        check.str_param(run_id, 'run_id')
        return 'sqlite:///{}'.format('/'.join(self.path_for_run_id(run_id).split(os.sep)))

    def _initdb(self, engine, run_id):
        try:
            SqlEventLogStorageMetadata.create_all(engine)
            engine.execute('PRAGMA journal_mode=WAL;')
        except (db.exc.DatabaseError, sqlite3.DatabaseError) as exc:
            six.raise_from(DagsterEventLogInvalidForRun(run_id=run_id), exc)

        alembic_config = get_alembic_config(__file__)
        conn = engine.connect()
        try:
            stamp_alembic_rev(alembic_config, conn)
        finally:
            conn.close()

    @contextmanager
    def connect(self, run_id=None):
        check.str_param(run_id, 'run_id')

        conn_string = self.conn_string_for_run_id(run_id)
        engine = create_engine(conn_string, poolclass=NullPool)

        if not os.path.exists(self.path_for_run_id(run_id)):
            self._initdb(engine, run_id)

        conn = engine.connect()
        try:
            with handle_schema_errors(
                conn,
                get_alembic_config(__file__),
                msg='SqliteEventLogStorage for run {run_id}'.format(run_id=run_id),
            ):
                yield conn
        finally:
            conn.close()

    def wipe(self):
        for filename in (
            glob.glob(os.path.join(self._base_dir, '*.db'))
            + glob.glob(os.path.join(self._base_dir, '*.db-wal'))
            + glob.glob(os.path.join(self._base_dir, '*.db-shm'))
        ):
            os.unlink(filename)

    def watch(self, run_id, start_cursor, callback):
        watchdog = SqliteEventLogStorageWatchdog(self, run_id, callback, start_cursor)
        self._watchers[run_id][callback] = (
            watchdog,
            self._obs.schedule(watchdog, self._base_dir, True),
        )

    def end_watch(self, run_id, handler):
        if handler in self._watchers[run_id]:
            event_handler, watch = self._watchers[run_id][handler]
            self._obs.remove_handler_for_watch(event_handler, watch)
            del self._watchers[run_id][handler]