survey-extraction/etl/utils/sharepoint/sharepoint.py

import aiohttp
import asyncio
from msal import ConfidentialClientApplication
from datetime import datetime, timedelta
from functools import wraps
import logging
from io import BytesIO
from etl.utils.logger import Logger

logger = Logger(name="SharePoint.py", level=logging.DEBUG).get_logger()

async def handle_error(response):
    """
    Handle errors based on HTTP status codes and log detailed information.
    """
    try:
        error_json = await response.json()
    except ValueError:
        error_json = {}

    error_code = error_json.get('code', 'unknownError')
    error_message = error_json.get('message', 'No detailed error message provided.')
    inner_error = error_json.get('innererror', {})
    details = error_json.get('details', [])

    logger.error(f"Error Code: {error_code}")
    logger.error(f"Error Message: {error_message}")
    if inner_error:
        logger.error(f"Inner Error: {inner_error}")
    if details:
        logger.error(f"Error Details: {details}")

    if response.status == 401:
        logger.error("Unauthorized. Token might be invalid.")
    elif response.status == 403:
        logger.error("Forbidden. Access denied to the requested resource.")
    elif response.status == 404:
        logger.error("Not Found. The requested resource doesn’t exist.")
    elif response.status == 429:
        retry_after = int(response.headers.get('Retry-After', 5))  # Default to 5 seconds if not provided
        logger.warning(f"Too Many Requests. Retrying after {retry_after} seconds...")
        await asyncio.sleep(retry_after)
        return 'retry'
    elif response.status in (500, 503):
        retry_after = int(response.headers.get('Retry-After', 5))  # Default to 5 seconds if not provided
        logger.error(f"Server error. Retrying after {retry_after} seconds...")
        await asyncio.sleep(retry_after)
        return 'retry'
    else:
        raise ValueError(f"API request failed with status code {response.status} - {error_message}")


def api_call_decorator(func):
    """
    Handles various aspects of the API call, including refreshing the access token if needed and handling pagination.
    :param func: The function to be decorated.
    :return: The wrapped function.
    """

    @wraps(func)
    async def wrapper(self, *args, **kwargs):
        try:
            # Check and refresh the access token if needed
            if self.is_access_token_expired():
                await self.retrieve_access_token()
                logger.info("Access token refreshed.")

            # Get the HTTP method, URL, and optionally data from the function
            http_method, url, data = await func(self, *args, **kwargs)

            # Initialize the results list and handle pagination if page_size is provided
            results = []
            page_size = kwargs.get('page_size', None)
            response_data = {}

            while url:
                async with aiohttp.ClientSession() as session:
                    async with session.request(http_method, url, headers=self.headers, json=data) as response:
                        # Handle the response
                        if response.status == 200:
                            response_json = await response.json()  # Store the response JSON
                            if page_size:
                                results.extend(response_json.get('value', []))
                                url = response_json.get('@odata.nextLink', None)
                            else:
                                response_data = response_json  # Capture the full response for consistency
                                break
                        else:
                            retry = await handle_error(response)
                            if retry == 'retry':
                                continue

            if page_size:
                response_data = {'value': results}

            return response_data

        except Exception as e:
            logger.exception("An error occurred during the API call.")
            raise e

    return wrapper


class SharePointClient:
    access_token = None
    access_token_request_timestamp = None
    access_token_expiry = None
    headers = None

    TIMESTAMP_FORMAT = "%Y-%m-%dT%H:%M:%S.%fZ"

    def __init__(self, tenant_id, client_id, client_secret, site_id, access_token=None,
                 access_token_expiration_details=None):
        """
        Initializes the SharePointClient with necessary credentials and site information.
        :param tenant_id: The tenant ID.
        :param client_id: The client ID.
        :param client_secret: The client secret.
        :param site_id: The site ID.
        :param access_token: The access token (optional)
        :param access_token_expiration_details: The access token expiration details (optional)
        """
        self.tenant_id = tenant_id
        self.client_id = client_id
        self.client_secret = client_secret

        if access_token:
            if not access_token_expiration_details:
                raise ValueError("Access token expiration details must be provided.")
            self.access_token = access_token
            self.set_access_token_expiration_details(access_token_expiration_details)
            self.headers = {
                'Authorization': f"Bearer {self.access_token['access_token']}"
            }
        else:
            # Avoid asyncio.run() here since it's called in a running event loop
            asyncio.ensure_future(self.retrieve_access_token())

        # Retrieve static identifiers
        self.site_id = site_id
        asyncio.ensure_future(self.initialize_document_drive())

    async def initialize_document_drive(self):
        self.document_drive = await self.get_documents_drive()
        self.document_drive_id = self.document_drive["id"]

    def get_token_expiration_details(self):
        """
        Returns the access token expiration details. Converts the datetime objects to strings for serialization.
        :return:
        """
        return {
            'access_token_request_timestamp': datetime.strftime(
                self.access_token_request_timestamp, self.TIMESTAMP_FORMAT
            ),
            'access_token_expiry': datetime.strftime(self.access_token_expiry, self.TIMESTAMP_FORMAT)
        }

    def set_access_token_expiration_details(self, access_token_expiration_details):
        """
        Sets the access token expiration details from a serialized dictionary.
        :param access_token_expiration_details: The serialized access token expiration details.
        :return:
        """
        self.access_token_request_timestamp = datetime.strptime(
            access_token_expiration_details['access_token_request_timestamp'], self.TIMESTAMP_FORMAT
        )
        self.access_token_expiry = datetime.strptime(
            access_token_expiration_details['access_token_expiry'], self.TIMESTAMP_FORMAT
        )

    def is_access_token_expired(self):
        """
        Checks if the access token has expired. If it has, a new access token is retrieved.
        :return: True if expired, False otherwise.
        """
        return datetime.now() >= self.access_token_expiry

    async def retrieve_access_token(self, refresh=False):
        """
        Implements authentication using MSAL.
        :param refresh: If True, force a refresh of the access token.
        :return: None
        """
        app = ConfidentialClientApplication(
            self.client_id,
            authority=f"https://login.microsoftonline.com/{self.tenant_id}",
            client_credential=self.client_secret
        )

        scope = ["https://graph.microsoft.com/.default"]

        access_token_request_timestamp = datetime.now()

        if refresh:
            logger.info("Forcing refresh of access token.")
            token = app.acquire_token_for_client(scopes=scope)
        else:
            # Check if a token is already cached
            token = app.acquire_token_silent(scope, account=None)

            if not token:
                token = app.acquire_token_for_client(scopes=scope)

        if "access_token" not in token:
            logger.error("Authentication failed.")
            raise ValueError("Authentication failed")

        access_token_expiry = access_token_request_timestamp + timedelta(
            seconds=token['expires_in'] - 20
        )

        self.access_token = token
        self.access_token_request_timestamp = access_token_request_timestamp
        self.access_token_expiry = access_token_expiry
        self.headers = {
            'Authorization': f"Bearer {self.access_token['access_token']}"
        }

        logger.info("Access token retrieved successfully.")

    @api_call_decorator
    async def get_documents_drive(self):
        """
        Get the document drive of the SharePoint site.
        :return: Tuple containing HTTP method, URL, and None for data.
        """
        url = f"https://graph.microsoft.com/v1.0/sites/{self.site_id}/drive"
        logger.info(f"Getting document drive from URL: {url}")
        return 'GET', url, None

    @api_call_decorator
    async def list_folder_contents(self, folder_path: str, page_size: int = 100):
        """
        This function will list the contents of a folder in SharePoint.
        :param folder_path: The path of the folder.
        :param page_size: The number of items per page...
        """
        url = f"https://graph.microsoft.com/v1.0/sites/{self.site_id}/drive/root:/{folder_path}:/children"
        logger.info(f"Listing folder contents for {folder_path}")
        return 'GET', url, None