[xy] Improve google sheets source. (mage-ai#4561)

Author: wangxiaoyou1993

mage_integrations/mage_integrations/connections/google_sheets/__init__.py

@@ -1,9 +1,13 @@
+from typing import Dict
+
+import backoff
+import requests
 from google.oauth2 import service_account
 from googleapiclient.discovery import build
+
 from mage_integrations.connections.base import Connection
 from mage_integrations.connections.utils.google import CredentialsInfoType
 from mage_integrations.utils.dictionary import merge_dict
-from typing import Dict
 
 
 class GoogleSheets(Connection):
@@ -22,11 +26,11 @@ def __init__(
 
     def connect(self):
         """Create a connection to the Google Search Console API and return service object.
-        
+
         Returns:
             service (object): Google Search Console service object.
         """
-        
+
         scope = ['https://www.googleapis.com/auth/spreadsheets.readonly']
         if self.credentials_info is None:
             credentials = service_account.Credentials.from_service_account_file(
@@ -41,12 +45,20 @@ def connect(self):
             credentials=credentials
         )
 
+    @backoff.on_exception(
+        backoff.expo,
+        requests.exceptions.HTTPError,
+        max_tries=5,
+        factor=2,
+    )
     def load(
         self,
         spreadsheet_id: str,
         value_range: str = None,
-        opts: Dict = dict(),
+        opts: Dict = None,
     ):
+        if opts is None:
+            opts = dict()
         service = self.connect()
         kwargs = dict(
             spreadsheetId=spreadsheet_id,
@@ -60,6 +72,12 @@ def load(
 
         return response.get('values', [])
 
+    @backoff.on_exception(
+        backoff.expo,
+        requests.exceptions.HTTPError,
+        max_tries=5,
+        factor=2,
+    )
     def get_spreadsheet_metadata(
         self,
         spreadsheet_id: str,

mage_integrations/mage_integrations/sources/base.py

@@ -1,36 +1,34 @@
+import inspect
+import json
+import os
+import sys
+import traceback
 from datetime import datetime
+from os.path import isfile
+from typing import Dict, Generator, List
+
+import dateutil.parser
+import pandas as pd
+import simplejson
+import singer
+from singer import utils
+from singer.schema import Schema
+
 from mage_integrations.sources.catalog import Catalog, CatalogEntry
 from mage_integrations.sources.constants import (
     REPLICATION_METHOD_FULL_TABLE,
     REPLICATION_METHOD_INCREMENTAL,
     REPLICATION_METHOD_LOG_BASED,
 )
 from mage_integrations.sources.messages import write_records, write_schema, write_state
-from mage_integrations.sources.utils import (
-    get_standard_metadata,
-    parse_args,
-)
+from mage_integrations.sources.utils import get_standard_metadata, parse_args
 from mage_integrations.utils.array import find_index
 from mage_integrations.utils.dictionary import extract, group_by, merge_dict
 from mage_integrations.utils.files import get_abs_path
 from mage_integrations.utils.logger import Logger
 from mage_integrations.utils.logger.constants import TYPE_SAMPLE_DATA
 from mage_integrations.utils.schema_helpers import extract_selected_columns
 
-from os.path import isfile
-from singer import utils
-from singer.schema import Schema
-from typing import Dict, Generator, List
-import dateutil.parser
-import inspect
-import json
-import os
-import pandas as pd
-import simplejson
-import singer
-import sys
-import traceback
-
 LOGGER = singer.get_logger()
 
 
@@ -47,7 +45,7 @@ def __init__(
         load_sample_data: bool = False,
         log_to_stdout: bool = False,
         logger=LOGGER,
-        query: Dict = {},
+        query: Dict = None,
         schemas_folder: str = 'schemas',
         selected_streams: List[str] = None,
         settings: Dict = None,
@@ -56,6 +54,8 @@ def __init__(
         test_connection: bool = False,
         verbose: int = 1,
     ):
+        if query is None:
+            query = {}
         args = parse_args([])
         if args:
             if args.catalog:
@@ -231,7 +231,8 @@ def process(self) -> None:
             elif self.count_records_mode:
                 arr = []
                 selected_streams_arr = self.catalog.get_selected_streams(self.state or {}) or []
-                streams = [stream for stream in selected_streams_arr if stream.tap_stream_id in self.selected_streams]
+                streams = [stream for stream in selected_streams_arr
+                           if stream.tap_stream_id in self.selected_streams]
                 for stream in streams:
                     tap_stream_id = stream.tap_stream_id
                     count = self.count_records(
@@ -443,9 +444,10 @@ def sync_stream(self, stream, properties: Dict = None) -> int:
 
                 write_state(state)
 
-        self.logger.info(f'Load data for stream {tap_stream_id} completed.', tags=merge_dict(tags, dict(
-            records=record_count,
-        )))
+        self.logger.info(
+            f'Load data for stream {tap_stream_id} completed.',
+            tags=merge_dict(tags, dict(records=record_count)),
+        )
 
         return record_count
 
@@ -491,7 +493,7 @@ def write_records(self, stream, rows: List[Dict], properties: Dict = None) -> Li
                 if self.is_sorted:
                     state = {}
 
-                    for idx, col in enumerate(bookmark_properties):
+                    for _, col in enumerate(bookmark_properties):
                         singer.write_bookmark(
                             state,
                             tap_stream_id,
@@ -534,9 +536,10 @@ def sync(self, catalog: Catalog, properties: Dict = None) -> None:
             self.process_stream(stream, properties)
             record_count = self.sync_stream(stream, properties)
 
-            self.logger.info(f'Sync for stream {tap_stream_id} completed.', tags=merge_dict(tags, dict(
-                records=record_count,
-            )))
+            self.logger.info(
+                f'Sync for stream {tap_stream_id} completed.',
+                tags=merge_dict(tags, dict(records=record_count)),
+            )
 
         self.logger.info('Sync completed.')
 
@@ -566,7 +569,8 @@ def build_catalog_entry(
             replication_method=replication_method or self.get_forced_replication_method(stream_id),
             schema=schema.to_dict(),
             stream_id=stream_id,
-            valid_replication_keys=bookmark_properties or self.get_valid_replication_keys(stream_id),
+            valid_replication_keys=bookmark_properties or
+            self.get_valid_replication_keys(stream_id),
         )
 
         idx = find_index(lambda x: len(x['breadcrumb']) == 0, metadata)
@@ -582,11 +586,13 @@ def build_catalog_entry(
                 database=None,
                 disable_column_type_check=None,
                 is_view=None,
-                key_properties=key_properties or [],  # User customizes this after creating catalog from discover.
+                # User customizes this after creating catalog from discover.
+                key_properties=key_properties or [],
                 metadata=metadata,
                 partition_keys=None,
                 replication_key=replication_key or '',
-                replication_method=replication_method or self.get_forced_replication_method(stream_id),
+                replication_method=replication_method or
+                self.get_forced_replication_method(stream_id),
                 row_count=None,
                 schema=schema,
                 stream=stream_id,
@@ -611,7 +617,7 @@ def count_records(
         self,
         stream,
         bookmarks: Dict = None,
-        query: Dict = {},
+        query: Dict = None,
         **kwargs,
     ) -> int:
         return 0
@@ -620,7 +626,7 @@ def load_data(
         self,
         stream,
         bookmarks: Dict = None,
-        query: Dict = {},
+        query: Dict = None,
         sample_data: bool = False,
         start_date: datetime = None,
         **kwargs,