<commit_before>from apiclient.discovery import build
from apiclient.http import MediaIoBaseUpload
from dateutil.parser import parse
from django.conf import settings
from django.core.files import File
from django.core.files.storage import Storage
import httplib2
from io import BytesIO
import mimetypes 
from oauth2client.client import SignedJwtAssertionCredentials
import os.path
import requests

class GoogleDriveStorage(Storage):
    """
    Storage class for Django that interacts with Google Drive as persistent storage.
    This class uses a system account for Google API that create an application drive 
    (the drive is not owned by any Google User, but it is owned by the application declared on 
    Google API console).
    """
    
    _GOOGLE_DRIVE_FOLDER_MIMETYPE_ = "application/vnd.google-apps.folder"

    def __init__(self, service_email = None, private_key_file = None):
        self._drive_service = None
        try:
            service_email = service_email if service_email is not None else settings.GOOGLE_DRIVE_STORAGE["service_account"]["email"]
            private_key_file = private_key_file if private_key_file is not None else settings.GOOGLE_DRIVE_STORAGE["service_account"]["private_key_file_path"]
            key = None
            
            # Creating a Google Drive Service API using a system account (without OAuth)
            # See https://developers.google.com/drive/web/service-accounts#console_name_project_service_accounts for more info
            private_key_abs_path = "{0}{1}{2}".format(settings.BASE_DIR, os.path.sep, private_key_file)
            if not os.path.exists(private_key_abs_path):
                raise ValueError("Unable to find provided key file")
            else:
                with file(private_key_abs_path, 'rb') as f:
                    key = f.read()
                credentials = SignedJwtAssertionCredentials(
                    service_email,
                    key,
                    scope = "https://www.googleapis.com/auth/drive"
                )
                http = httplib2.Http()
                http = credentials.authorize(http)
                
                self._drive_service = build('drive', 'v2', http=http)
        except KeyError:
            raise ValueError("You must configure properly your settings file. Check Google Drive Storage docs for more info")
        
    def _split_path(self, p):
        """
        Split a complete path in a list of strings
        
        :param p: Path to be splitted
        :type p: string
        :returns: list - List of strings that composes the path
        """
        p = p[1:] if p[0] == '/' else p
        a,b = os.path.split(p)
        return (self._split_path(a) if len(a) and len(b) else []) + [b]
        
    def _get_or_create_folder(self, path, parent_id = None):
        """
        Create a folder on Google Drive. 
        It creates folders recursively.
        If the folder already exists, it retrieves only the unique identifier.
        
        :param path: Path that had to be created
        :type path: string
        :param parent_id: Unique identifier for its parent (folder)
        :type parent_id: string
        :returns: dict
        """
        folder_data = self._check_file_exists(path, parent_id)
        if folder_data is None:
            # Folder does not exists, have to create
            split_path = self._split_path(path)
            current_folder_data = None
            for p in split_path:
                meta_data = {
                    'title': p,
                    'mimeType': self._GOOGLE_DRIVE_FOLDER_MIMETYPE_
                }
                if current_folder_data is not None:
                    meta_data['parents'] = [{'id': current_folder_data['id']}]
                else:
                    # This is the first iteration loop so we have to set the parent_id
                    # obtained by the user, if available
                    if parent_id is not None:
                        meta_data['parents'] = [{'id': parent_id}]
                current_folder_data = self._drive_service.files().insert(body=meta_data).execute()
            return current_folder_data
        else:
            return folder_data
    
    def _check_file_exists(self, filename, parent_id = None):
        """
        Check if a file with specific parameters exists in Google Drive.
        
        :param filename: File or folder to search
        :type filename: string
        :param parent_id: Unique identifier for its parent (folder)
        :type parent_id: string
        :param mime_type: Mime Type for the file to search
        :type mime_type: string
        :returns: dict containing file / folder data if exists or None if does not exists
        """
        split_filename = self._split_path(filename)
        if len(split_filename) > 1:
            # This is an absolute path with folder inside
            # First check if the first element exists as a folder
            # If so call the method recursively with next portion of path
            # Otherwise the path does not exists hence the file does not exists
            q = u"mimeType = '{0}' and title = '{1}'".format(self._GOOGLE_DRIVE_FOLDER_MIMETYPE_, unicode(split_filename[0]))
            if parent_id is not None:
                q = u"{0} and '{1}' in parents".format(q, parent_id)
            max_results = 1000 # Max value admitted from google drive
            folders = self._drive_service.files().list(q=q, maxResults=max_results).execute()
            for folder in folders["items"]:
                if folder["title"] == unicode(split_filename[0]):
                    # Assuming every folder has a single parent
                    return self._check_file_exists(os.path.sep.join(split_filename[1:]), folder["id"])
            return None
        else:
            # This is a file, checking if exists
            q = u"title = '{0}'".format(split_filename[0])
            if parent_id is not None:
                q = u"{0} and '{1}' in parents".format(q, parent_id)
            max_results = 1000 # Max value admitted from google drive
            file_list = self._drive_service.files().list(q=q, maxResults=max_results).execute()
            if len(file_list["items"]) == 0:
                q = u"" if parent_id is None else u"'{0}' in parents".format(parent_id)
                file_list = self._drive_service.files().list(q=q, maxResults=max_results).execute()
                for element in file_list["items"]:
                    if unicode(split_filename[0]) in element["title"]:
                        return element
                return None
            else:
                return file_list["items"][0]
        
    # Methods that had to be implemented
    # to create a valid storage for Django
    
    def _open(self, name, mode='rb'):
        file_data = self._check_file_exists(name)
        r = requests.get(file_data['webContentLink'])
        return File(BytesIO(r.content), name) 
        
     
    def _save(self, name, content):
        folder_path = os.path.sep.join(self._split_path(name)[:-1])
        folder_data = self._get_or_create_folder(folder_path)
        parent_id = None if folder_data is None else folder_data['id']
        # Now we had created (or obtained) folder on GDrive
        # Upload the file
        fd = BytesIO(content.file.read())
        mime_type = mimetypes.guess_type(name)
        media_body = MediaIoBaseUpload(fd, mime_type, resumable=True)
        body = {
            'title': name,
            'mimeType': mime_type
        }
        # Set the parent folder.<commit_msg>Fix a problem when Python could not determine file mime type by extension. In this case the mime type will be generic 'application/octet-stream'.<commit_after>from apiclient.discovery import build
from apiclient.http import MediaIoBaseUpload
from dateutil.parser import parse
from django.conf import settings
from django.core.files import File
from django.core.files.storage import Storage
import httplib2
from io import BytesIO
import mimetypes 
from oauth2client.client import SignedJwtAssertionCredentials
import os.path
import requests

class GoogleDriveStorage(Storage):
    """
    Storage class for Django that interacts with Google Drive as persistent storage.
    This class uses a system account for Google API that create an application drive 
    (the drive is not owned by any Google User, but it is owned by the application declared on 
    Google API console).
    """
    
    _UNKNOWN_MIMETYPE_ = "application/octet-stream"
    _GOOGLE_DRIVE_FOLDER_MIMETYPE_ = "application/vnd.google-apps.folder"

    def __init__(self, service_email = None, private_key_file = None):
        self._drive_service = None
        try:
            service_email = service_email if service_email is not None else settings.GOOGLE_DRIVE_STORAGE["service_account"]["email"]
            private_key_file = private_key_file if private_key_file is not None else settings.GOOGLE_DRIVE_STORAGE["service_account"]["private_key_file_path"]
            key = None
            
            # Creating a Google Drive Service API using a system account (without OAuth)
            # See https://developers.google.com/drive/web/service-accounts#console_name_project_service_accounts for more info
            private_key_abs_path = "{0}{1}{2}".format(settings.BASE_DIR, os.path.sep, private_key_file)
            if not os.path.exists(private_key_abs_path):
                raise ValueError("Unable to find provided key file")
            else:
                with file(private_key_abs_path, 'rb') as f:
                    key = f.read()
                credentials = SignedJwtAssertionCredentials(
                    service_email,
                    key,
                    scope = "https://www.googleapis.com/auth/drive"
                )
                http = httplib2.Http()
                http = credentials.authorize(http)
                
                self._drive_service = build('drive', 'v2', http=http)
        except KeyError:
            raise ValueError("You must configure properly your settings file. Check Google Drive Storage docs for more info")
        
    def _split_path(self, p):
        """
        Split a complete path in a list of strings
        
        :param p: Path to be splitted
        :type p: string
        :returns: list - List of strings that composes the path
        """
        p = p[1:] if p[0] == '/' else p
        a,b = os.path.split(p)
        return (self._split_path(a) if len(a) and len(b) else []) + [b]
        
    def _get_or_create_folder(self, path, parent_id = None):
        """
        Create a folder on Google Drive. 
        It creates folders recursively.
        If the folder already exists, it retrieves only the unique identifier.
        
        :param path: Path that had to be created
        :type path: string
        :param parent_id: Unique identifier for its parent (folder)
        :type parent_id: string
        :returns: dict
        """
        folder_data = self._check_file_exists(path, parent_id)
        if folder_data is None:
            # Folder does not exists, have to create
            split_path = self._split_path(path)
            current_folder_data = None
            for p in split_path:
                meta_data = {
                    'title': p,
                    'mimeType': self._GOOGLE_DRIVE_FOLDER_MIMETYPE_
                }
                if current_folder_data is not None:
                    meta_data['parents'] = [{'id': current_folder_data['id']}]
                else:
                    # This is the first iteration loop so we have to set the parent_id
                    # obtained by the user, if available
                    if parent_id is not None:
                        meta_data['parents'] = [{'id': parent_id}]
                current_folder_data = self._drive_service.files().insert(body=meta_data).execute()
            return current_folder_data
        else:
            return folder_data
    
    def _check_file_exists(self, filename, parent_id = None):
        """
        Check if a file with specific parameters exists in Google Drive.
        
        :param filename: File or folder to search
        :type filename: string
        :param parent_id: Unique identifier for its parent (folder)
        :type parent_id: string
        :param mime_type: Mime Type for the file to search
        :type mime_type: string
        :returns: dict containing file / folder data if exists or None if does not exists
        """
        split_filename = self._split_path(filename)
        if len(split_filename) > 1:
            # This is an absolute path with folder inside
            # First check if the first element exists as a folder
            # If so call the method recursively with next portion of path
            # Otherwise the path does not exists hence the file does not exists
            q = u"mimeType = '{0}' and title = '{1}'".format(self._GOOGLE_DRIVE_FOLDER_MIMETYPE_, unicode(split_filename[0]))
            if parent_id is not None:
                q = u"{0} and '{1}' in parents".format(q, parent_id)
            max_results = 1000 # Max value admitted from google drive
            folders = self._drive_service.files().list(q=q, maxResults=max_results).execute()
            for folder in folders["items"]:
                if folder["title"] == unicode(split_filename[0]):
                    # Assuming every folder has a single parent
                    return self._check_file_exists(os.path.sep.join(split_filename[1:]), folder["id"])
            return None
        else:
            # This is a file, checking if exists
            q = u"title = '{0}'".format(split_filename[0])
            if parent_id is not None:
                q = u"{0} and '{1}' in parents".format(q, parent_id)
            max_results = 1000 # Max value admitted from google drive
            file_list = self._drive_service.files().list(q=q, maxResults=max_results).execute()
            if len(file_list["items"]) == 0:
                q = u"" if parent_id is None else u"'{0}' in parents".format(parent_id)
                file_list = self._drive_service.files().list(q=q, maxResults=max_results).execute()
                for element in file_list["items"]:
                    if unicode(split_filename[0]) in element["title"]:
                        return element
                return None
            else:
                return file_list["items"][0]
        
    # Methods that had to be implemented
    # to create a valid storage for Django
    
    def _open(self, name, mode='rb'):
        file_data = self._check_file_exists(name)
        r = requests.get(file_data['webContentLink'])
        return File(BytesIO(r.content), name) 
        
     
    def _save(self, name, content):
        folder_path = os.path.sep.join(self._split_path(name)[:-1])
        folder_data = self._get_or_create_folder(folder_path)
        parent_id = None if folder_data is None else folder_data['id']
        # Now we had created (or obtained) folder on GDrive
        # Upload the file
        fd = BytesIO(content.file.read())
        mime_type = mimetypes.guess_type(name)
        if mime_type is None:
            mime_type = self._UNKNOWN_MIMETYPE_
        media_body = MediaIoBaseUpload(fd, mime_type, resumable=True)
        body = {
            'title': name,
            'mimeType': mime_type
        }
        # Set the parent folder.
{'fraction_non_alphanumeric': 0.08114035087719298, 'fraction_numerical': 0.0043859649122807015, 'mean_word_length': 2.1022408963585435, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 5, 'https://': 4, 'lorem ipsum': 0, 'www.': 2, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '3352534', 'n_tokens_mistral': 4249, 'n_tokens_neox': 4115, 'n_words': 1327}