materials-commons-cli 2.2.3__tar.gz → 2.3.0b2__tar.gz

{materials_commons_cli-2.2.3 → materials_commons_cli-2.3.0b2}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: materials_commons_cli
-Version: 2.2.3
+Version: 2.3.0b2
 Summary: Materials Commons CLI
 Home-page: https://materials-commons.github.io/materials-commons-cli/html/index.html
 Author: Materials Commons development team

materials_commons_cli-2.3.0b2/materials_commons/cli/__init__.py

@@ -0,0 +1 @@
+__version__ = '2.3.0b2'

materials_commons_cli-2.3.0b2/materials_commons/cli/async_reconciler.py

@@ -0,0 +1,104 @@
+import asyncio
+from datetime import timezone
+from pathlib import Path
+from typing import Optional, AsyncIterator
+
+from materials_commons.cli.filedb import FileIndexDB
+from materials_commons.cli.local_project import LocalProject
+from materials_commons.cli.models import FileState, FileDecision, Observation, RemoteFileEntry
+from materials_commons.cli.reconcile3 import SingleFileReconciler, ReconcileMode
+from materials_commons.cli.server import projects
+from materials_commons.cli.walk import ListDirFunc, IgnoreFunc, async_walk, path_to_local_file_entry
+
+
+class AsyncReconciler:
+    def __init__(self,
+                 proj: LocalProject,
+                 db: FileIndexDB,
+                 reconcile_mode: ReconcileMode,
+                 # recompute_checksum: bool = True,
+                 max_concurrent: int = 10):
+        self.proj = proj
+        self.db = db
+        # self.recompute_checksum = recompute_checksum
+        self.max_concurrent = max_concurrent
+        self.reconcile_mode = reconcile_mode
+        self.reconciler = SingleFileReconciler(proj=self.proj, mode=reconcile_mode)
+
+    async def walk(self, path: str | Path, listdir_fn: ListDirFunc, recursive: bool = False,
+                   ignore_fn: Optional[IgnoreFunc] = None) -> AsyncIterator[tuple[Path, dict[str, FileState]]]:
+        sem = asyncio.Semaphore(self.max_concurrent)
+
+        async def single_entry_reconcile(obs: Observation) -> FileDecision:
+            async with sem:
+                return await self.reconciler.reconcile_file(obs)
+
+        async for current_path, observations in async_walk(path, recursive=recursive, listdir_fn=listdir_fn,
+                                                           ignore_fn=ignore_fn):
+            path_states: dict[str, FileState] = {}
+            remote_dir = projects.local_to_remote_project_path(Path(self.proj.local_path), Path(current_path))
+
+            # Create a map of file_records to name
+            file_records = await self.db.get_files_by_dir(remote_dir.as_posix())
+            file_records_map = {file_record.name: file_record for file_record in file_records}
+
+            # Go through all the observations and match to the file records
+            for obs in observations:
+                obs.file_record = file_records_map.get(obs.name, None)
+                path_states[obs.name] = FileState(observation=obs)
+
+            results = await asyncio.gather(
+                *[single_entry_reconcile(state.observation) for state in path_states.values()],
+                return_exceptions=True)
+
+            # At this point path_entries contains entries in one of 4 states:
+            # 1. Both remote and local entries exist
+            # 2. Only remote entry exists
+            # 3. Only local entry exists
+            # 4. Lookup failed, and we got an exception
+
+            for state, result in zip(path_states.values(), results):
+                if isinstance(result, Exception):
+                    state.exception = result
+                else:
+                    state.file_decision = result
+
+            yield current_path, path_states
+
+    async def reconcile_file(self, path: str | Path) -> FileState:
+        remote_project_path = projects.local_to_remote_project_path(Path(self.proj.local_path), Path(path))
+        file_record = await self.db.get_file_by_path(remote_project_path.as_posix())
+        try:
+            mc_remote_file = await asyncio.to_thread(self.proj.remote.get_file_by_path, self.proj.id,
+                                                     remote_project_path.as_posix())
+        except Exception as e:
+            mc_remote_file = None
+        local_file_entry = path_to_local_file_entry(Path(path))
+        if mc_remote_file:
+            remote_file_entry = RemoteFileEntry(
+                path=Path(mc_remote_file.path),
+                name=mc_remote_file.name,
+                kind="file",
+                size=mc_remote_file.size,
+                mtime_ns=int(mc_remote_file.updated_at.replace(tzinfo=timezone.utc).timestamp() * 1_000_000_000),
+                ctime_ns=int(mc_remote_file.created_at.replace(tzinfo=timezone.utc).timestamp() * 1_000_000_000),
+                remote_file_id=getattr(mc_remote_file, "id", None),
+                checksum=getattr(mc_remote_file, "checksum", None),
+                raw=mc_remote_file,
+            )
+        else:
+            remote_file_entry = None
+
+        obs = Observation(
+            remote_entry=remote_file_entry,
+            local_entry=local_file_entry,
+            file_record=file_record,
+            local_path=Path(path),
+            remote_path=remote_project_path,
+        )
+        decision = await self.reconciler.reconcile_file(obs)
+
+        return FileState(
+            observation=obs,
+            file_decision=decision,
+        )

materials_commons_cli-2.3.0b2/materials_commons/cli/config.py

@@ -0,0 +1,224 @@
+import getpass
+import json
+import os
+import re
+from dataclasses import asdict, dataclass, field
+from os.path import join
+from typing import Any, Dict, List, Optional
+
+import requests
+from materials_commons.api.client import Client
+
+
+def default_config_dir() -> str:
+    user = getpass.getuser()
+    return join(os.path.expanduser("~" + user), ".materialscommons")
+
+
+def resolve_config_file(
+    config_dir_path: Optional[str],
+    config_file_name: str,
+) -> str:
+    if not config_dir_path:
+        config_dir_path = default_config_dir()
+    return join(config_dir_path, config_file_name)
+
+
+def default_data() -> Dict[str, Any]:
+    return {
+        "apikey": None,
+        "mcurl": None,
+        "email": None,
+        "remotes": [],
+        "globus": {},
+        "client_uuid": None,
+    }
+
+
+def load_json(path: str) -> Dict[str, Any]:
+    if os.path.exists(path):
+        with open(path, "r") as f:
+            return json.load(f)
+    return default_data()
+
+
+def apply_environment(data: Dict[str, Any]) -> Dict[str, Any]:
+    data = dict(data)
+
+    env_apikey = os.environ.get("MC_API_KEY")
+    env_mcurl = os.environ.get("MC_API_URL")
+    env_email = os.environ.get("MC_API_EMAIL")
+
+    if env_apikey:
+        data["apikey"] = env_apikey
+    if env_mcurl:
+        data["mcurl"] = env_mcurl
+    if env_email:
+        data["email"] = env_email
+
+    return data
+
+
+def apply_overrides(
+    data: Dict[str, Any],
+    override_config: Optional[Dict[str, Any]],
+) -> Dict[str, Any]:
+    data = dict(data)
+    if override_config:
+        data.update(override_config)
+    return data
+
+
+def normalize_default_remote(data: Dict[str, Any]) -> Dict[str, Any]:
+    data = dict(data)
+
+    if data.get("mcurl") and data.get("apikey") and data.get("email"):
+        data["default_remote"] = {
+            "mcurl": data.get("mcurl"),
+            "email": data.get("email"),
+            "mcapikey": data.get("apikey"),
+        }
+    elif "default_remote" not in data:
+        data["default_remote"] = {
+            "mcurl": data.get("mcurl"),
+            "email": "__default__",
+            "mcapikey": data.get("apikey"),
+        }
+
+    return data
+
+
+@dataclass(eq=True)
+class RemoteConfig:
+    mcurl: Optional[str] = None
+    email: Optional[str] = None
+    mcapikey: Optional[str] = None
+
+    def get_params(self) -> Dict[str, Optional[str]]:
+        return {"apikey": self.mcapikey}
+
+    def make_client(self) -> Client:
+        return Client(self.mcapikey, self.mcurl)
+
+
+@dataclass
+class GlobusConfig:
+    transfer_rt: Optional[str] = None
+    endpoint_id: Optional[str] = None
+
+
+@dataclass
+class Config:
+    """Configuration variables."""
+
+    default_remote: RemoteConfig = field(default_factory=RemoteConfig)
+    remotes: List[RemoteConfig] = field(default_factory=list)
+    globus: GlobusConfig = field(default_factory=GlobusConfig)
+    developer_mode: bool = False
+    REST_logging: bool = False
+    client_uuid: Optional[str] = None
+    config_file: Optional[str] = None
+
+    @classmethod
+    def load(
+        cls,
+        config_dir_path: Optional[str] = None,
+        config_file_name: str = "config.json",
+        override_config: Optional[Dict[str, Any]] = None,
+    ) -> "Config":
+        config_file = resolve_config_file(config_dir_path, config_file_name)
+
+        data = load_json(config_file)
+        data = apply_environment(data)
+        data = apply_overrides(data, override_config)
+        data = normalize_default_remote(data)
+
+        return cls(
+            default_remote=RemoteConfig(**data.get("default_remote", {})),
+            remotes=[RemoteConfig(**kwargs) for kwargs in data.get("remotes", [])],
+            globus=GlobusConfig(**data.get("globus", {})),
+            developer_mode=data.get("developer_mode", False),
+            REST_logging=data.get("REST_logging", False),
+            client_uuid=data.get("client_uuid"),
+            config_file=config_file,
+        )
+
+    def to_dict(self) -> Dict[str, Any]:
+        return {
+            "default_remote": asdict(self.default_remote),
+            "remotes": [asdict(value) for value in self.remotes],
+            "globus": asdict(self.globus),
+            "developer_mode": self.developer_mode,
+            "REST_logging": self.REST_logging,
+            "client_uuid": self.client_uuid or "",
+        }
+
+    def save(self) -> None:
+        if not self.config_file:
+            raise ValueError("Cannot save Config without config_file")
+
+        config_dir_path = os.path.dirname(self.config_file)
+        if config_dir_path and not os.path.exists(config_dir_path):
+            os.makedirs(config_dir_path, exist_ok=True)
+
+        with open(self.config_file, "w") as f:
+            json.dump(self.to_dict(), f, indent=2)
+
+        os.chmod(self.config_file, 0o600)
+
+    def find_remote(
+        self,
+        email: Optional[str],
+        mcurl: Optional[str],
+    ) -> Optional[RemoteConfig]:
+        for remote in self.remotes:
+            if remote.email == email and remote.mcurl == mcurl:
+                return remote
+        return None
+
+
+def get_remote_config_and_login_if_necessary(
+    email: Optional[str] = None,
+    mcurl: Optional[str] = None,
+) -> RemoteConfig:
+    """Prompt for login if remote is not stored in Config."""
+    config = Config.load()
+    existing_remote = config.find_remote(email=email, mcurl=mcurl)
+
+    if existing_remote is not None:
+        return existing_remote
+
+    remote_config = RemoteConfig(mcurl=mcurl, email=email)
+
+    while True:
+        try:
+            print("Login to:", email, mcurl)
+            password = getpass.getpass(prompt="password: ")
+            remote_config.mcapikey = Client.get_apikey(email, password, mcurl)
+            break
+        except requests.exceptions.HTTPError as e:
+            print(str(e))
+            if not re.search("Bad Request for url", str(e)):
+                raise
+            print("Wrong password for " + email + " at " + mcurl)
+        except requests.exceptions.ConnectionError as e:
+            print("Could not connect to " + mcurl)
+            raise e
+
+    config.remotes.append(remote_config)
+    config.save()
+
+    print()
+    print("Added APIKey for", email, "at", mcurl, "to", config.config_file)
+    print()
+
+    return remote_config
+
+
+def make_client_and_login_if_necessary(
+    email: Optional[str] = None,
+    mcurl: Optional[str] = None,
+) -> Client:
+    """Make Client, prompting for login if remote is not stored in Config."""
+    remote_config = get_remote_config_and_login_if_necessary(mcurl=mcurl, email=email)
+    return remote_config.make_client()

materials_commons_cli-2.3.0b2/materials_commons/cli/downloader.py

@@ -0,0 +1,116 @@
+import logging
+from pathlib import Path
+
+from igittigitt import IgnoreParser, igittigitt
+
+from materials_commons.cli.async_reconciler import AsyncReconciler
+from materials_commons.cli.filedb import FileIndexDB
+from materials_commons.cli.local_project import LocalProject
+from materials_commons.cli.requests import DownloadRequest
+from materials_commons.cli.server.service_container import ServiceContainer
+from materials_commons.cli.server.service_runtime import ServiceRuntime
+from materials_commons.cli.walk import make_merged_listdir_func
+
+logger = logging.getLogger(__name__)
+
+
+class Downloader:
+    def __init__(self, proj: LocalProject, db: FileIndexDB, ignore_parser: IgnoreParser, container: ServiceContainer,
+                 service_runtime: ServiceRuntime):
+        self.proj = proj
+        self.db = db
+        self.container = container
+        self.service_runtime = service_runtime
+        self.ignore_parser = ignore_parser
+        self.async_reconciler = AsyncReconciler(db=db, proj=proj, reconcile_mode="download")
+        self.listdir_fn = make_merged_listdir_func(proj)
+        self.transfer_ids = []
+
+    @classmethod
+    async def init(cls, proj: LocalProject):
+        db = await proj.get_filedb()
+
+        # Setup ignore parser
+        ignore_parser = igittigitt.IgnoreParser()
+        ignore_parser.parse_rule_files(base_dir=proj.local_path, filename=".mcignore", add_default_patterns=False)
+
+        # Start services
+        container = ServiceContainer.create()
+        service_runtime = ServiceRuntime(container)
+        await service_runtime.start(db_manager=True, file_download_manager=True)
+
+        return cls(proj=proj, db=db, ignore_parser=ignore_parser, container=container, service_runtime=service_runtime)
+
+    async def run(self, paths: list[str | Path], recursive: bool = False, force: bool = False):
+        try:
+            for path in paths:
+                p = Path(path)
+                if p.is_dir():
+                    # Walk and download local dir
+                    await self._download_dir(p, recursive=recursive)
+                elif p.is_file():
+                    # Download local file
+                    await self._download_file(p)
+                else:
+                    # If we are here, then the path entry is a remote path. We look up the path to
+                    # determine the type of download to perform.
+                    remote_path = self.proj.to_remote_path(p)
+                    try:
+                        f = self.proj.remote.get_file_by_path(self.proj.id, remote_path.as_posix())
+                    except Exception as e:
+                        # logger.error(f"Error getting file by path {remote_path}: {e}")
+                        print(f"No such remote file {p}")
+                        continue
+
+                    if f is None:
+                        continue
+
+                    if f.mime_type == "directory":
+                        # Remote entry is a directory
+                        await self._download_dir(p, recursive=recursive, force=force)
+                    else:
+                        # Remote entry is a file
+                        await self._download_file(p, force=force)
+
+        except Exception as e:
+            logger.error(f"Unexpected error during download: {e}")
+            return False
+
+        finally:
+            await self.container.file_download_manager.wait_all()
+            await self.service_runtime.stop(drain=True)
+            await self.db.close()
+
+    async def _download_file(self, path: Path, force: bool = False) -> None:
+        file_state = await self.async_reconciler.reconcile_file(path)
+        if file_state.exception:
+            logger.error(f"Error encountered while processing {path}: {file_state.exception}")
+            print(f"Error encountered while processing {path}: {file_state.exception}")
+            return
+
+        if file_state.file_decision.action == "download" or force:
+            download_request = DownloadRequest(observation=file_state.observation,
+                                               updated_record=file_state.file_decision.updated_record,
+                                               project=self.proj)
+            print(f"Downloading {path}...")
+            transfer_id = await self.container.file_download_manager.download_file(download_request)
+            self.transfer_ids.append(transfer_id)
+        else:
+            print(f"Skipping download of {path} because {file_state.file_decision.reason}")
+
+    async def _download_dir(self, path: Path, recursive: bool = False, force: bool = False):
+        async for current_path, path_entries in self.async_reconciler.walk(path=path, listdir_fn=self.listdir_fn,
+                                                                           recursive=recursive, ignore_fn=None):
+            for entry_name in sorted(path_entries):
+                file_state = path_entries[entry_name]
+                if file_state.exception:
+                    logger.error(f"Error encountered while processing {entry_name}: {file_state.exception}")
+                    print(f"Error encountered while processing {entry_name}: {file_state.exception}")
+                    continue
+                if file_state.file_decision.action == "download" or force:
+                    download_request = DownloadRequest(observation=file_state.observation,
+                                                       updated_record=file_state.file_decision.updated_record,
+                                                       project=self.proj)
+                    print(f"Downloading {entry_name} to {current_path / entry_name}...")
+                    transfer_id = await self.container.file_download_manager.download_file(download_request)
+                    self.transfer_ids.append(transfer_id)