<reponame>WangCHEN9/azure_client_hydra<gh_stars>0
import os
from configparser import ConfigParser
from datetime import datetime
from typing import List
from pathlib import Path
import inspect

from loguru import logger
from azure.storage.blob import BlobServiceClient, BlobClient, ContainerClient
from azure.cosmosdb.table.tableservice import TableService
from azure.cosmosdb.table.models import Entity

from utils.help_functions import get_config
from .tableItem import tableItem


class azureClient:
    """this class manage the connection to cloud using the secret keys"""

    def __init__(self) -> None:
        """init function, reading config with secret keys"""
        self.cfg = get_config()
        self._read_config_and_set_up()
        self._set_up_azure_connection()

    def _set_up_proxy(self, proxy: str) -> None:
        if self.cfg.proxy.use_proxy:
            os.environ["HTTP_PROXY"] = proxy
            os.environ["HTTPS_PROXY"] = proxy

    def _read_config_and_set_up(self) -> None:
        """read config"""

        AZURE_STORAGE_CONNECTION_STRING = self.cfg.azure.storage_connection_string
        AZURE_STORAGE_KEY = self.cfg.azure.storage_key
        PROXY = f"http://{self.cfg.proxy.user_account}:{self.cfg.proxy.user_password}@{self.cfg.proxy.proxy_address}"

        os.environ["AZURE_STORAGE_CONNECTION_STRING"] = AZURE_STORAGE_CONNECTION_STRING
        os.environ["AZURE_STORAGE_KEY"] = AZURE_STORAGE_KEY
        self._set_up_proxy(PROXY)

    def _set_up_azure_connection(self) -> None:
        """[set up azure connection, consider we have a blob service and a table service]"""

        connect_str = os.getenv("AZURE_STORAGE_CONNECTION_STRING")
        azure_storage_key = os.getenv("AZURE_STORAGE_KEY")

        self.blob_service_client = BlobServiceClient.from_connection_string(connect_str)
        self.table_service = TableService(
            account_name=self.cfg.azure.account_name,
            account_key=azure_storage_key,
        )
        self.table_name = self.cfg.azure.table_name

    def _creat_or_get_container_client(self, container_name: str) -> ContainerClient:
        """[try to create the container, if failed , means it exist, then get_container_client]

        Args:
            container_name ([str]): [name of contrainer]
        """
        try:
            container_client: ContainerClient = (
                self.blob_service_client.create_container(container_name)
            )
        except:
            container_client: ContainerClient = (
                self.blob_service_client.get_container_client(container_name)
            )
        return container_client

    def upload_file(
        self,
        container_name: str,
        file_path: Path,
        metadata=None,
        delete_after_upload: bool = False,
    ) -> None:
        """upload files to cloud

        Args:
            container_name (str): [name of blob container]
            file_path (Path): [upload file path]
            metadata ([type], optional): [metadata to upload]. Defaults to None.
            delete_after_upload (bool, optional): [if delete local file after upload]. Defaults to False.
        """
        if file_path.exists():
            container_client = self._creat_or_get_container_client(container_name)
            with file_path.open(mode="rb") as data:
                blob_client = container_client.upload_blob(
                    name=os.path.basename(file_path),
                    data=data,
                    metadata=metadata,
                    overwrite=True,
                )
            if delete_after_upload and blob_client.exists():
                try:
                    file_path.unlink()
                    logger.info(f"{str(file_path)} deleted after upload to azure")
                except Exception as err:
                    logger.critical(f"{err}")
        else:
            logger.warning(f"{str(file_path)} is not found..")

    def download_blob(self, blob_client: BlobClient, download_file_path: Path) -> None:
        """download blob from cloud

        Args:
            blob_client ([BlobClient]): [blob client]
            download_file_path (Path): [path for save download file]
        """
        logger.debug(f"Downloading blob to \n\t {download_file_path}")
        with open(download_file_path, "wb") as download_file:
            download_file.write(blob_client.download_blob().readall())

    def create_table(self, table_name) -> None:
        """create a new table"""
        self.table_service.create_table(table_name)

    def add_entity_to_table(self, item: tableItem) -> None:
        """add new entity to the table

        Args:
            item (tableItem): [item data class]
        """
        data_to_add = item.unpack_dict_and_prepare_format_for_upload()
        logger.debug(f"uploading data to {self.table_name}: {data_to_add}")
        if not self.table_service.exists(self.table_name):
            logger.info(
                f"table is not existing yet -> creating table: {self.table_name}"
            )
            self.create_table(self.table_name)
        self.table_service.insert_or_replace_entity(self.table_name, data_to_add)

    def download_blob_with_title(
        self, title: str, container_client: ContainerClient, local_folder: Path
    ):
        blob_name = f"{title}.jpg"
        blob_client = container_client.get_blob_client(blob=blob_name)
        self.download_blob(blob_client, download_file_path=local_folder / blob_name)
