fix google drive

Author: bary12

app/data_sources/google_drive.py

@@ -2,7 +2,7 @@
 import io
 import logging
 from datetime import datetime
-from typing import Dict
+from typing import Dict, List
 
 from apiclient.discovery import build
 from googleapiclient.http import MediaIoBaseDownload
@@ -44,34 +44,53 @@ def __init__(self, *args, **kwargs):
  self._credentials = ServiceAccountCredentials.from_json_keyfile_dict(self._config, scopes=scopes)
  self._http_auth = self._credentials.authorize(Http())
  self._drive = build('drive', 'v3', http=self._http_auth)
-
- def _should_index_file(self, file):
- mime_types = [
+ 
+ self._supported_mime_types = [
  'application/vnd.google-apps.document',
  'application/vnd.openxmlformats-officedocument.wordprocessingml.document',
  'application/vnd.openxmlformats-officedocument.presentationml.presentation'
  ]
- return file['kind'] == 'drive#file' and file['mimeType'] in mime_types
 
- def _feed_new_documents(self) -> None:
- files = self._drive.files().list(fields='files(kind,id,name,mimeType,owners,webViewLink,modifiedTime,parents)').execute()
- files = files['files']
+ def _should_index_file(self, file):
+ if file['mimeType'] not in self._supported_mime_types:
+ return False
+
+ last_modified = datetime.strptime(file['modifiedTime'], "%Y-%m-%dT%H:%M:%S.%fZ")
+ if last_modified < self._last_index_time:
+ return False
+ 
+ return True
+
+ def _index_files_from_drive(self, drive) -> List[dict]: 
+ is_shared_drive = drive['id'] is not None
+
+ print(f'Indexing drive {drive["name"]}')
+
+ kwargs = {
+ 'corpora': 'drive',
+ 'driveId': drive['id'],
+ 'includeItemsFromAllDrives': True,
+ 'supportsAllDrives': True,
+ } if is_shared_drive else {}
+
+ # Todo: add pagination
+ files = self._drive.files().list(
+ fields='files(kind,id,name,mimeType,lastModifyingUser,webViewLink,modifiedTime,parents)',
+ pageSize=1000,
+ **kwargs
+ ).execute()['files']
+
  files = [file for file in files if self._should_index_file(file)]
+
  documents = []
 
- logging.getLogger().info(f'got {len(files)} documents from google drive.')
+ logging.getLogger().info(f'got {len(files)} documents from drive {drive["name"]}.')
 
  for file in files:
  logging.getLogger().info(f'processing file {file["name"]}')
- last_modified = datetime.strptime(file['modifiedTime'], "%Y-%m-%dT%H:%M:%S.%fZ")
- if last_modified < self._last_index_time:
- continue
 
  file_id = file['id']
- file = self._drive.files().get(fileId=file_id,
- fields='id,name,mimeType,owners,webViewLink,modifiedTime,parents').execute()
  file_to_download = file['name']
-
  if file['mimeType'] == 'application/vnd.google-apps.document':
  content = self._drive.files().export(fileId=file_id, mimeType='text/html').execute().decode('utf-8')
  content = html_to_text(content)
@@ -100,25 +119,38 @@ def _feed_new_documents(self) -> None:
  os.remove(file_to_download)
  except Exception as error:
  logging.exception(f'Error occured parsing file "{file["name"]}" from google drive')
-
+ 
  try:
- parent = self._drive.files().get(fileId=file['parents'][0], fields='name').execute()
+ if is_shared_drive:
+ parent = self._drive.drives().get(driveId=drive['id'], fields='name').execute()
+ else:
+ parent = self._drive.files().get(fileId=file['parents'][0], fields='name').execute()
  parent_name = parent['name']
  except Exception as e:
  logging.exception(f"Error while getting folder name of google docs file {file['name']}")
  parent_name = ''
 
+ last_modified = datetime.strptime(file['modifiedTime'], "%Y-%m-%dT%H:%M:%S.%fZ")
+
  documents.append(BasicDocument(
  id=file_id,
  data_source_id=self._data_source_id,
  type=DocumentType.DOCUMENT,
  title=file['name'],
  content=content,
- author=file['owners'][0]['displayName'],
- author_image_url=file['owners'][0]['photoLink'],
+ author=file['lastModifyingUser']['displayName'],
+ author_image_url=file['lastModifyingUser']['photoLink'],
  location=parent_name,
  url=file['webViewLink'],
  timestamp=last_modified
  ))
 
  IndexingQueue.get().feed(documents)
+
+ def _get_all_drives(self) -> List[dict]:
+ return [{'name': 'My Drive', 'id': None}] \
+ + self._drive.drives().list(fields='drives(id,name)').execute()['drives']
+
+ def _feed_new_documents(self) -> None:
+ for drive in self._get_all_drives():
+ self._index_files_from_drive(drive)