Commit 8b8177b4 authored by Louis Simoneau's avatar Louis Simoneau Committed by Tao Feng

Guard against log entry pages not containing 'entries' key (#170)

In practice, when requesting logs going back more than a few days, the
StackDriver API sometimes responds with an empty page with a next page
cursor. After polling the next page a few times you eventually start
getting entries. The empty pages have no `entries` key, which previously
resulted in a KeyError.

This change allows the extractor to continue polling instead of blowing up.
parent bb66d041
...@@ -158,7 +158,8 @@ class BigQueryTableUsageExtractor(Extractor): ...@@ -158,7 +158,8 @@ class BigQueryTableUsageExtractor(Extractor):
response = self.logging_service.entries().list(body=body).execute( response = self.logging_service.entries().list(body=body).execute(
num_retries=BigQueryTableUsageExtractor.NUM_RETRIES) num_retries=BigQueryTableUsageExtractor.NUM_RETRIES)
while response: while response:
yield response if 'entries' in response:
yield response
try: try:
if 'nextPageToken' in response: if 'nextPageToken' in response:
......
...@@ -125,6 +125,12 @@ FAILURE = {"entries": [ ...@@ -125,6 +125,12 @@ FAILURE = {"entries": [
}, },
}, },
}]} # noqa }]} # noqa
# An empty dict will be ignored, but putting in nextPageToken causes the test
# to loop infinitely, so we need a bogus key/value to ensure that we will try
# to read entries
NO_ENTRIES = { 'key': 'value' } # noqa
KEYFILE_DATA = """ KEYFILE_DATA = """
ewogICJ0eXBlIjogInNlcnZpY2VfYWNjb3VudCIsCiAgInByb2plY3RfaWQiOiAieW91ci1wcm9q ewogICJ0eXBlIjogInNlcnZpY2VfYWNjb3VudCIsCiAgInByb2plY3RfaWQiOiAieW91ci1wcm9q
ZWN0LWhlcmUiLAogICJwcml2YXRlX2tleV9pZCI6ICJiMDQ0N2U1ODEyYTg5ZTAyOTgxYjRkMWE1 ZWN0LWhlcmUiLAogICJwcml2YXRlX2tleV9pZCI6ICJiMDQ0N2U1ODEyYTg5ZTAyOTgxYjRkMWE1
...@@ -217,6 +223,21 @@ class TestBigqueryUsageExtractor(unittest.TestCase): ...@@ -217,6 +223,21 @@ class TestBigqueryUsageExtractor(unittest.TestCase):
self.assertEqual(key.email, 'your-user-here@test.com') self.assertEqual(key.email, 'your-user-here@test.com')
self.assertEqual(value, 1) self.assertEqual(value, 1)
@patch('databuilder.extractor.bigquery_usage_extractor.build')
def test_no_entries(self, mock_build):
config_dict = {
'extractor.bigquery_table_usage.{}'.format(BigQueryTableUsageExtractor.PROJECT_ID_KEY):
'your-project-here',
}
conf = ConfigFactory.from_dict(config_dict)
mock_build.return_value = MockLoggingClient(NO_ENTRIES)
extractor = BigQueryTableUsageExtractor()
extractor.init(Scoped.get_scoped_conf(conf=conf,
scope=extractor.get_scope()))
result = extractor.extract()
self.assertIsNone(result)
@patch('databuilder.extractor.bigquery_usage_extractor.build') @patch('databuilder.extractor.bigquery_usage_extractor.build')
def test_key_path(self, mock_build): def test_key_path(self, mock_build):
""" """
......
Markdown is supported
0% or
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment