1 files changed, 117 insertions, 122 deletions

diff --git a/tap_google_sheets/client.py b/tap_google_sheets/client.py
index 4f38352..0cbad98 100644
--- a/tap_google_sheets/client.py
+++ b/tap_google_sheets/client.py
@@ -1,16 +1,21 @@
 from datetime import datetime, timedelta
 from collections import OrderedDict
 import backoff
-import requests
 import singer
+import logging
+import pickle
+import json
+import os
 from singer import metrics
 from singer import utils
+from google.oauth2 import service_account
+from google_auth_oauthlib.flow import InstalledAppFlow
+from google.auth.transport.requests import Request
+from googleapiclient.errors import HttpError
+import googleapiclient.discovery
 
-BASE_URL = 'https://www.googleapis.com'
-GOOGLE_TOKEN_URI = 'https://oauth2.googleapis.com/token'
 LOGGER = singer.get_logger()
 
-
 class Server5xxError(Exception):
     pass
 
@@ -101,90 +106,85 @@ ERROR_CODE_EXCEPTION_MAPPING = {
     428: GooglePreconditionRequiredError,
     500: GoogleInternalServiceError}
 
-
-def get_exception_for_error_code(error_code):
-    return ERROR_CODE_EXCEPTION_MAPPING.get(error_code, GoogleError)
-
-def raise_for_error(response):
-    try:
-        response.raise_for_status()
-    except (requests.HTTPError, requests.ConnectionError) as error:
-        try:
-            content_length = len(response.content)
-            if content_length == 0:
-                # There is nothing we can do here since Google has neither sent
-                # us a 2xx response nor a response content.
-                return
-            response = response.json()
-            if ('error' in response) or ('errorCode' in response):
-                message = '%s: %s' % (response.get('error', str(error)),
-                                      response.get('message', 'Unknown Error'))
-                error_code = response.get('error', {}).get('code')
-                ex = get_exception_for_error_code(error_code)
-                raise ex(message)
-            raise GoogleError(error)
-        except (ValueError, TypeError):
-            raise GoogleError(error)
-
 class GoogleClient: # pylint: disable=too-many-instance-attributes
-    def __init__(self,
-                 client_id,
-                 client_secret,
-                 refresh_token,
-                 user_agent=None):
-        self.__client_id = client_id
-        self.__client_secret = client_secret
-        self.__refresh_token = refresh_token
-        self.__user_agent = user_agent
-        self.__access_token = None
-        self.__expires = None
-        self.__session = requests.Session()
-        self.base_url = None
-
+    SCOPES = [
+        "https://www.googleapis.com/auth/drive.metadata.readonly",
+        "https://www.googleapis.com/auth/spreadsheets.readonly"
+    ]
+
+    def __init__(self, credentials_file):
+        self.__credentials = self.fetchCredentials(credentials_file)
+        self.__sheets_service = googleapiclient.discovery.build(
+            'sheets',
+            'v4',
+            credentials=self.__credentials,
+            cache_discovery=False
+        )
+        self.__drive_service = googleapiclient.discovery.build(
+            'drive',
+            'v3',
+            credentials=self.__credentials,
+            cache_discovery=False
+        )
+
+    def fetchCredentials(self, credentials_file):
+        LOGGER.debug('authenticate with google')
+        data = None
+
+        # Check a credentials file exist
+        if not os.path.exists(credentials_file):
+            raise Exception("The configured Google credentials file {} doesn't exist".format(credentials_file))
+
+        # Load credentials json file
+        with open(credentials_file) as json_file:
+            data = json.load(json_file)
+
+        if data.get('type', '') == 'service_account':
+            return self.fetchServiceAccountCredentials(credentials_file)
+        elif data.get('installed'):
+            return self.fetchInstalledOAuthCredentials(credentials_file)
+        else:
+            raise Exception("""This Google credentials file is not yet recognize.
+
+            Please use either:
+            - a Service Account (https://github.com/googleapis/google-api-python-client/blob/d0110cf4f7aaa93d6f56fc028cd6a1e3d8dd300a/docs/oauth-server.md)
+            - an installed OAuth client (https://github.com/googleapis/google-api-python-client/blob/d0110cf4f7aaa93d6f56fc028cd6a1e3d8dd300a/docs/oauth-installed.md)"""
+            )
+
+    def fetchServiceAccountCredentials(self, credentials_file):
+        # The service account credentials file can be used for server-to-server applications
+        return service_account.Credentials.from_service_account_file(
+            credentials_file, scopes=GoogleClient.SCOPES)
+
+    def fetchInstalledOAuthCredentials(self, credentials_file):
+        creds = None
+
+        # The file token.pickle stores the user's access and refresh tokens, and is
+        # created automatically when the authorization flow completes for the first
+        # time.
+        if os.path.exists('token.pickle'):
+            with open('token.pickle', 'rb') as token:
+                creds = pickle.load(token)
+
+        # If there are no (valid) credentials available, let the user log in.
+        if not creds or not creds.valid:
+            if creds and creds.expired and creds.refresh_token:
+                creds.refresh(Request())
+            else:
+                flow = InstalledAppFlow.from_client_secrets_file(
+                    credentials_file, GoogleClient.SCOPES)
+                creds = flow.run_local_server(port=0)
+            # Save the credentials for the next run
+            with open('token.pickle', 'wb') as token:
+                pickle.dump(creds, token)
+
+        return creds
 
     def __enter__(self):
-        self.get_access_token()
         return self
 
     def __exit__(self, exception_type, exception_value, traceback):
-        self.__session.close()
-
-
-    @backoff.on_exception(backoff.expo,
-                          Server5xxError,
-                          max_tries=5,
-                          factor=2)
-    def get_access_token(self):
-        # The refresh_token never expires and may be used many times to generate each access_token
-        # Since the refresh_token does not expire, it is not included in get access_token response
-        if self.__access_token is not None and self.__expires > datetime.utcnow():
-            return
-
-        headers = {}
-        if self.__user_agent:
-            headers['User-Agent'] = self.__user_agent
-
-        response = self.__session.post(
-            url=GOOGLE_TOKEN_URI,
-            headers=headers,
-            data={
-                'grant_type': 'refresh_token',
-                'client_id': self.__client_id,
-                'client_secret': self.__client_secret,
-                'refresh_token': self.__refresh_token,
-            })
-
-        if response.status_code >= 500:
-            raise Server5xxError()
-
-        if response.status_code != 200:
-            raise_for_error(response)
-
-        data = response.json()
-        self.__access_token = data['access_token']
-        self.__expires = datetime.utcnow() + timedelta(seconds=data['expires_in'])
-        LOGGER.info('Authorized, token expires = {}'.format(self.__expires))
-
+        LOGGER.debug('exiting google client')
 
     # Rate Limit: https://developers.google.com/sheets/api/limits
     #   100 request per 100 seconds per User
@@ -193,53 +193,48 @@ class GoogleClient: # pylint: disable=too-many-instance-attributes
                           max_tries=7,
                           factor=3)
     @utils.ratelimit(100, 100)
-    def request(self, method, path=None, url=None, api=None, **kwargs):
-        self.get_access_token()
-        self.base_url = 'https://sheets.googleapis.com/v4'
-        if api == 'files':
-            self.base_url = 'https://www.googleapis.com/drive/v3'
-
-        if not url and path:
-            url = '{}/{}'.format(self.base_url, path)
-
-        # endpoint = stream_name (from sync.py API call)
-        if 'endpoint' in kwargs:
-            endpoint = kwargs['endpoint']
-            del kwargs['endpoint']
+    def request(self, endpoint=None, params={}, **kwargs):
+        formatted_params = {}
+        for (key, value) in params.items():
+            # API parameters interpolation
+            # will raise a KeyError in case a necessary argument is missing
+            formatted_params[key] = value.format(**kwargs)
+
+        # Call the correct Google API depending on the stream name
+        if endpoint == 'spreadsheet_metadata' or endpoint == 'sheet_metadata':
+            # https://developers.google.com/sheets/api/reference/rest/v4/spreadsheets/get
+            request = self.__sheets_service.spreadsheets().get(**formatted_params)
+        elif endpoint == 'sheets_loaded':
+            # https://developers.google.com/sheets/api/reference/rest/v4/spreadsheets.values/get
+            request = self.__sheets_service.spreadsheets().values().get(**formatted_params)
+        elif endpoint == 'file_metadata':
+            # https://developers.google.com/drive/api/v3/reference/files/get
+            request = self.__drive_service.files().get(**formatted_params)
         else:
-            endpoint = None
-        LOGGER.info('{} URL = {}'.format(endpoint, url))
-
-        if 'headers' not in kwargs:
-            kwargs['headers'] = {}
-        kwargs['headers']['Authorization'] = 'Bearer {}'.format(self.__access_token)
+            raise Exception('{} not implemented yet!'.format(endpoint))
 
-        if self.__user_agent:
-            kwargs['headers']['User-Agent'] = self.__user_agent
+        with metrics.http_request_timer(endpoint) as timer:
+            error = None
+            status_code = 400
 
-        if method == 'POST':
-            kwargs['headers']['Content-Type'] = 'application/json'
+            try:
+                response = request.execute()
+                status_code = 200
+            except HttpError as e:
+                status_code = e.resp.status or status_code
+                error = e
 
-        with metrics.http_request_timer(endpoint) as timer:
-            response = self.__session.request(method, url, **kwargs)
-            timer.tags[metrics.Tag.http_status_code] = response.status_code
+            timer.tags[metrics.Tag.http_status_code] = status_code
 
-        if response.status_code >= 500:
+        if status_code >= 500:
             raise Server5xxError()
 
-        #Use retry functionality in backoff to wait and retry if
-        #response code equals 429 because rate limit has been exceeded
-        if response.status_code == 429:
+        # Use retry functionality in backoff to wait and retry if
+        # response code equals 429 because rate limit has been exceeded
+        if status_code == 429:
             raise Server429Error()
 
-        if response.status_code != 200:
-            raise_for_error(response)
-
-        # Ensure keys and rows are ordered as received from API
-        return response.json(object_pairs_hook=OrderedDict)
-
-    def get(self, path, api, **kwargs):
-        return self.request(method='GET', path=path, api=api, **kwargs)
+        if status_code != 200:
+            raise error
 
-    def post(self, path, api, **kwargs):
-        return self.request(method='POST', path=path, api=api, **kwargs)
+        return response