def test_delayed_iter():
events = []
def maker(n):
for i in range(n):
events.append(("add", i))
yield i
for i in delayed_iter(maker(4), 2):
events.append(("remove", i))
assert events == [
("add", 0),
("add", 1),
("add", 2),
("remove", 0),
("add", 3),
("remove", 1),
("remove", 2),
("remove", 3),
]
events.clear()
for i in delayed_iter(maker(2), None):
events.append(("remove", i))
assert events == [
("add", 0),
("add", 1),
("remove", 0),
("remove", 1),
]
def _join_events_by_job_id(
self, events: Iterable[Union[ReadEvent,
QueryEvent]]) -> Iterable[ReadEvent]:
# If caching eviction is enabled, we only store the most recently used query events,
# which are used when resolving job information within the read events.
query_jobs: MutableMapping[str, QueryEvent]
if self.config.query_log_delay:
query_jobs = cachetools.LRUCache(maxsize=5 *
self.config.query_log_delay)
else:
query_jobs = {}
def event_processor(
events: Iterable[Union[ReadEvent,
QueryEvent]]) -> Iterable[ReadEvent]:
for event in events:
if isinstance(event, QueryEvent):
if event.jobName:
query_jobs[event.jobName] = event
else:
yield event
# TRICKY: To account for the possibility that the query event arrives after
# the read event in the audit logs, we wait for at least `query_log_delay`
# additional events to be processed before attempting to resolve BigQuery
# job information from the logs. If `query_log_delay` is None, it gets treated
# as an unlimited delay, which prioritizes correctness at the expense of memory usage.
original_read_events = event_processor(events)
delayed_read_events = delayed_iter(original_read_events,
self.config.query_log_delay)
num_joined: int = 0
for event in delayed_read_events:
if (event.timestamp < self.config.start_time
or event.timestamp >= self.config.end_time
or not self._is_table_allowed(event.resource)):
continue
if event.jobName:
if event.jobName in query_jobs:
# Join the query log event into the table read log event.
num_joined += 1
event.query = query_jobs[event.jobName].query
# TODO also join into the query itself for column references
else:
self.report.report_warning(
str(event.resource),
f"Failed to match table read event {event.jobName} with job; try increasing `query_log_delay` or `max_query_duration`",
)
yield event
logger.info(
f"Number of read events joined with query events: {num_joined}")
def _join_events_by_job_id(
self, events: Iterable[Union[ReadEvent,
QueryEvent]]) -> Iterable[ReadEvent]:
# If caching eviction is enabled, we only store the most recently used query events,
# which are used when resolving job information within the read events.
query_jobs: MutableMapping[str, QueryEvent]
if self.config.query_log_delay:
query_jobs = cachetools.LRUCache(maxsize=5 *
self.config.query_log_delay)
else:
query_jobs = {}
def event_processor(
events: Iterable[Union[ReadEvent,
QueryEvent]]) -> Iterable[ReadEvent]:
for event in events:
if isinstance(event, QueryEvent):
query_jobs[event.jobName] = event
else:
yield event
# TRICKY: To account for the possibility that the query event arrives after
# the read event in the audit logs, we wait for at least `query_log_delay`
# additional events to be processed before attempting to resolve BigQuery
# job information from the logs. If `query_log_delay` is None, it gets treated
# as an unlimited delay, which prioritizes correctness at the expense of memory usage.
original_read_events = event_processor(events)
delayed_read_events = delayed_iter(original_read_events,
self.config.query_log_delay)
for event in delayed_read_events:
if event.jobName:
if event.jobName in query_jobs:
# Join the query log event into the table read log event.
event.query = query_jobs[event.jobName].query
# TODO also join into the query itself for column references
else:
self.report.report_warning(
"<general>",
"failed to match table read event with job; try increasing `query_log_delay`",
)
yield event
def _join_events_by_job_id(
self, events: Iterable[Union[ReadEvent, QueryEvent]]
) -> Iterable[ReadEvent]:
# We only store the most recently used query events, which are used when
# resolving job information within the read events.
query_jobs: cachetools.LRUCache[str, QueryEvent] = cachetools.LRUCache(
maxsize=2 * self.config.query_log_delay
)
def event_processor(
events: Iterable[Union[ReadEvent, QueryEvent]]
) -> Iterable[ReadEvent]:
for event in events:
if isinstance(event, QueryEvent):
query_jobs[event.jobName] = event
else:
yield event
# TRICKY: To account for the possibility that the query event arrives after
# the read event in the audit logs, we wait for at least `query_log_delay`
# additional events to be processed before attempting to resolve BigQuery
# job information from the logs.
original_read_events = event_processor(events)
delayed_read_events = delayed_iter(
original_read_events, self.config.query_log_delay
)
for event in delayed_read_events:
if event.jobName:
if event.jobName in query_jobs:
# Join the query log event into the table read log event.
event.query = query_jobs[event.jobName].query
# TODO also join into the query itself for column references
else:
self.report.report_warning(
"<general>",
"failed to match table read event with job; try increasing `query_log_delay`",
)
yield event