Source code for galaxy.files.sources.googledrive

try:
    from gdrive_fsspec import GoogleDriveFileSystem
    from google.oauth2.credentials import Credentials
    from googleapiclient.discovery import build
except ImportError:
    GoogleDriveFileSystem = None


from datetime import datetime
from typing import (
    Annotated,
    Optional,
    Union,
)

from fsspec import AbstractFileSystem
from pydantic import (
    AliasChoices,
    Field,
)

from galaxy.files.models import FilesSourceRuntimeContext
from galaxy.util.config_templates import TemplateExpansion
from ._fsspec import (
    CacheOptionsDictType,
    FsspecBaseFileSourceConfiguration,
    FsspecBaseFileSourceTemplateConfiguration,
    FsspecFilesSource,
)

GalaxyGoogleDriveFileSystem: Optional[type[AbstractFileSystem]]

if GoogleDriveFileSystem is not None:

    class _GalaxyGoogleDriveFileSystem(GoogleDriveFileSystem):
        def __init__(self, access_token: str, **kwargs):
            self._galaxy_credentials = Credentials(token=access_token)
            super().__init__(token="galaxy", **kwargs)

        def connect(self, method=None):
            if method == "galaxy":
                srv = build("drive", "v3", credentials=self._galaxy_credentials)
                self.srv = srv
                self.files = srv.files()
            else:
                super().connect(method=method)

        def put_file(self, lpath, rpath, *args, **kwargs):
            parent = self._parent(str(rpath))
            if parent not in self.dircache:
                self.dircache[parent] = []
            return super().put_file(lpath, rpath, *args, **kwargs)

    GalaxyGoogleDriveFileSystem = _GalaxyGoogleDriveFileSystem

else:
    GalaxyGoogleDriveFileSystem = None

AccessTokenField = Field(
    ...,
    validation_alias=AliasChoices("oauth2_access_token", "accessToken", "access_token", "token"),
)


class GoogleDriveFileSourceTemplateConfiguration(FsspecBaseFileSourceTemplateConfiguration):
    access_token: Annotated[Union[str, TemplateExpansion], AccessTokenField]


class GoogleDriveFilesSourceConfiguration(FsspecBaseFileSourceConfiguration):
    access_token: Annotated[str, AccessTokenField]



[docs]
class GoogleDriveFilesSource(
    FsspecFilesSource[GoogleDriveFileSourceTemplateConfiguration, GoogleDriveFilesSourceConfiguration]
):
    plugin_type = "googledrive"
    required_module = GalaxyGoogleDriveFileSystem
    required_package = "gdrive_fsspec"

    template_config_class = GoogleDriveFileSourceTemplateConfiguration
    resolved_config_class = GoogleDriveFilesSourceConfiguration

    def _open_fs(
        self,
        context: FilesSourceRuntimeContext[GoogleDriveFilesSourceConfiguration],
        cache_options: CacheOptionsDictType,
    ):
        if GalaxyGoogleDriveFileSystem is None:
            raise self.required_package_exception
        return GalaxyGoogleDriveFileSystem(access_token=context.config.access_token, **cache_options)

    def _to_filesystem_path(self, path: str, config: GoogleDriveFilesSourceConfiguration) -> str:
        if path in ("", "/"):
            return ""
        return path.lstrip("/")

    def _adapt_entry_path(self, filesystem_path: str, config: GoogleDriveFilesSourceConfiguration) -> str:
        if not filesystem_path or filesystem_path == "/":
            return "/"
        return filesystem_path if filesystem_path.startswith("/") else f"/{filesystem_path}"

    def _extract_timestamp(self, info: dict):
        timestamp = info.get("modifiedTime") or info.get("createdTime") or super()._extract_timestamp(info)
        if isinstance(timestamp, str):
            return datetime.fromisoformat(timestamp.replace("Z", "+00:00"))
        return timestamp



__all__ = ("GoogleDriveFilesSource",)