make reader compatible with google-cloud-logging v3 releases

gcp_flowlogs_reader/gcp_flowlogs_reader.py

@@ -9,8 +9,15 @@
     NotFound,
     TooManyRequests,
 )
-from google.cloud.logging import Client as LoggingClient
-from google.cloud.logging.entries import StructEntry
+from google.cloud.logging import (
+    Client as LoggingClient,
+    __version__ as gcp_logging_version,
+)
+
+try:
+    from google.cloud.logging import StructEntry
+except ImportError:
+    from google.cloud.logging.entries import StructEntry
 
 try:
     from google.cloud.resource_manager import Client as ResourceManagerClient
@@ -22,6 +29,21 @@
 
 
 def page_helper(logging_client, wait_time=1.0, **kwargs):
+    # handle google-cloud-logging >= 3.0
+    if gcp_logging_version[0] == '3':
+        # the project arg in google-cloud-logging >= 3.0 was changes to resource_names
+        if 'projects' in kwargs:
+            kwargs['resource_names'] = [
+                f'projects/{project}' for project in kwargs['projects']
+            ]
+            del kwargs['projects']
+        # google-cloud-logging >= 3.0 handles paging internally
+        iterator = logging_client.list_entries(**kwargs)
+        for entry in iterator:
+            yield entry
+        return
+
+    # google-cloud-logging < 3.0 requires us to handle paging
     kwargs['page_token'] = None
     while True:
         try:

setup.cfg

@@ -1,6 +1,6 @@
 [metadata]
 name = gcp_flowlogs_reader
-version = 2.0.0
+version = 2.1.0
 license = Apache
 url = https://github.com/obsrvbl-oss/gcp-flowlogs-reader
 description = Reader for Google Cloud VPC Flow Logs
@@ -25,7 +25,7 @@ classifiers =
 packages = find:
 python_requires = >=3.6
 install_requires =
-    google-cloud-logging < 2.0
+    google-cloud-logging < 4.0
     google-cloud-resource-manager
     six
 

tests/test_gcp_flowlogs_reader.py

@@ -6,11 +6,10 @@
 from unittest.mock import MagicMock, patch, call
 from tempfile import NamedTemporaryFile
 
-from gcp_flowlogs_reader.gcp_flowlogs_reader import BASE_LOG_NAME
+from gcp_flowlogs_reader.gcp_flowlogs_reader import BASE_LOG_NAME, page_helper
 from google.api_core.exceptions import GoogleAPIError, PermissionDenied, NotFound
 from google.cloud.logging import Client
-from google.cloud.logging.entries import StructEntry
-from google.cloud.logging.resource import Resource
+from google.cloud.logging import StructEntry, Resource
 from google.oauth2.service_account import Credentials
 
 from gcp_flowlogs_reader.aggregation import aggregated_records
@@ -182,6 +181,11 @@ def __next__(self):
         return ''
 
 
+class MockV3Iterator:
+    def __iter__(self):
+        return iter(SAMPLE_ENTRIES[0])
+
+
 class MockFailedIterator:
     def __init__(self):
         self.pages = self
@@ -207,6 +211,29 @@ def __next__(self):
 class TestClient(Client):
     _credentials = ''
 
+    def list_entries(
+        self,
+        *,
+        projects=None,
+        filter_=None,
+        order_by=None,
+        max_results=None,
+        page_size=None,
+        page_token=None,
+    ):
+        pass
+
+
+class V3PageHelperTests(TestCase):
+    def test_init_outbound(self):
+        MockLoggingClient = MagicMock()
+        MockLoggingClient.return_value.list_entries.return_value = iter(SAMPLE_ENTRIES)
+        iterator = page_helper(logging_client=MockLoggingClient(), projects=['proj1'])
+        self.assertEqual(list(iterator), SAMPLE_ENTRIES)
+        MockLoggingClient.return_value.list_entries.assert_called_with(
+            resource_names=['projects/proj1']
+        )
+
 
 class FlowRecordTests(TestCase):
     def test_init_outbound(self):
@@ -373,6 +400,7 @@ def test_resource_labels(self):
 
 
 @patch(PREFIX('LoggingClient'), autospec=TestClient)
+@patch(PREFIX('gcp_logging_version'), '1.12.2')
 class ReaderTests(TestCase):
     def test_init_with_client(self, MockLoggingClient):
         logging_client = MagicMock(Client)
@@ -713,10 +741,10 @@ def test_custom_key(self):
         )
 
 
+@patch(PREFIX('gcp_logging_version'), '1.12.2')
 class MainCLITests(TestCase):
     def setUp(self):
-        patch_path = PREFIX('LoggingClient')
-        with patch(patch_path, autospec=True) as MockLoggingClient:
+        with patch(PREFIX('LoggingClient'), autospec=TestClient) as MockLoggingClient:
             MockLoggingClient.return_value.project = 'yoyodyne-102010'
             MockLoggingClient.return_value.list_entries.return_value = MockIterator()
             self.reader = Reader()