QuLab 2.10.10__cp313-cp313-win_amd64.whl

qulab/executor/storage.py

@@ -0,0 +1,615 @@
+import hashlib
+import lzma
+import pickle
+import re
+import uuid
+import zipfile
+from dataclasses import dataclass, field
+from datetime import datetime, timedelta
+from functools import lru_cache
+from pathlib import Path
+from typing import Any, Literal
+from urllib.parse import parse_qs
+
+from loguru import logger
+
+try:
+    from paramiko import SSHClient
+    from paramiko.ssh_exception import SSHException
+except:
+    import warnings
+
+    warnings.warn("Can't import paramiko, ssh support will be disabled.")
+
+    class SSHClient:
+
+        def __init__(self):
+            raise ImportError(
+                "Can't import paramiko, ssh support will be disabled.")
+
+        def __enter__(self):
+            return self
+
+        def __exit__(self, exc_type, exc_value, traceback):
+            pass
+
+    class SSHException(Exception):
+        pass
+
+
+from ..cli.config import get_config_value
+
+__current_config_cache = None
+
+
+@dataclass
+class Report():
+    workflow: str = ''
+    in_spec: bool = False
+    bad_data: bool = False
+    fully_calibrated: bool = False
+    calibrated_time: datetime = field(default_factory=datetime.now)
+    checked_time: datetime = field(default_factory=datetime.now)
+    ttl: timedelta = timedelta(days=3650)
+    parameters: dict = field(default_factory=dict)
+    oracle: dict = field(default_factory=dict)
+    other_infomation: dict = field(default_factory=dict)
+    data: Any = field(default_factory=tuple, repr=False)
+    index: int = -1
+    previous_path: Path | None = field(default=None, repr=False)
+    heads: dict[str, Path] = field(default_factory=dict, repr=False)
+    base_path: str | Path | None = field(default=None, repr=False)
+    path: Path | None = field(default=None, repr=False)
+    config_path: Path | None = field(default=None, repr=False)
+    script_path: Path | None = field(default=None, repr=False)
+
+    def __getstate__(self):
+        state = self.__dict__.copy()
+        state.pop('base_path', None)
+        for k in ['path', 'previous_path', 'config_path', 'script_path']:
+            if state[k] is not None:
+                state[k] = str(state[k])
+        return state
+
+    def __setstate__(self, state):
+        for k in ['path', 'previous_path', 'config_path', 'script_path']:
+            if state[k] is not None:
+                state[k] = Path(state[k])
+        self.__dict__.update(state)
+
+    @property
+    def previous(self):
+        if self.previous_path is not None and self.base_path is not None:
+            return load_report(self.previous_path, self.base_path)
+        else:
+            return None
+
+    @property
+    def state(self) -> Literal['OK', 'Bad', 'Outdated']:
+        state = 'Bad'
+        match (self.in_spec, self.bad_data):
+            case (True, False):
+                state = 'OK'
+            case (False, True):
+                state = 'Bad'
+            case (False, False):
+                state = 'Outdated'
+        return state
+
+    @state.setter
+    def state(self, state: Literal['OK', 'Bad', 'Outdated', 'In spec',
+                                   'Out of spec', 'Bad data']):
+        if state not in [
+                'OK', 'Bad', 'Outdated', 'In spec', 'Out of spec', 'Bad data'
+        ]:
+            raise ValueError(
+                f'Invalid state: {state}, state must be one of "OK", "Bad" and "Outdated"'
+            )
+        if state in ['In spec', 'OK']:
+            self.in_spec = True
+            self.bad_data = False
+        elif state in ['Bad data', 'Bad']:
+            self.bad_data = True
+            self.in_spec = False
+        else:
+            self.bad_data = False
+            self.in_spec = False
+
+    @property
+    def config(self):
+        if self.config_path is not None and self.base_path is not None:
+            return load_item(self.config_path, 'items', self.base_path)
+        else:
+            return None
+
+    @property
+    def script(self):
+        if self.script_path is not None and self.base_path is not None:
+            source = load_item(self.script_path, 'items', self.base_path)
+            if isinstance(source, str):
+                return source
+            else:
+                from .template import inject_mapping
+                return inject_mapping(*source)[0]
+        else:
+            return None
+
+    @property
+    def template_source(self):
+        if self.script_path is not None and self.base_path is not None:
+            source = load_item(self.script_path, 'items', self.base_path)
+            return source
+        else:
+            return None
+
+
+def random_path(base: Path) -> Path:
+    while True:
+        s = uuid.uuid4().hex
+        path = Path(s[:2]) / s[2:4] / s[4:6] / s[6:]
+        if not (base / path).exists():
+            return path
+
+
+def find_report(
+    workflow: str, base_path: str | Path = get_config_value("data", Path)
+) -> Report | None:
+    if workflow.startswith("cfg:"):
+        return find_config_key_history(workflow[4:], base_path)
+
+    base_path = Path(base_path)
+    path = get_head(workflow, base_path)
+    if path is None:
+        return None
+    return load_report(path, base_path)
+
+
+def renew_report(workflow: str, report: Report | None, base_path: str | Path):
+    logger.debug(f'Renewing report for "{workflow}"')
+    if report is not None:
+        report.checked_time = datetime.now()
+        return save_report(workflow,
+                           report,
+                           base_path,
+                           overwrite=True,
+                           refresh_heads=True)
+    else:
+        raise ValueError(f"Can't renew report for {workflow}")
+
+
+def revoke_report(workflow: str, report: Report | None, base_path: str | Path):
+    logger.debug(f'Revoking report for "{workflow}"')
+    base_path = Path(base_path)
+    if report is not None:
+        report.in_spec = False
+        report.previous_path = report.path
+        return save_report(workflow,
+                           report,
+                           base_path,
+                           overwrite=False,
+                           refresh_heads=True)
+
+
+def get_report_by_index(
+    index: int, base_path: str | Path = get_config_value("data", Path)
+) -> Report | None:
+    try:
+        path = query_index("report", base_path, index)
+        return load_report(path, base_path)
+    except:
+        raise
+        return None
+
+
+def get_head(workflow: str, base_path: str | Path) -> Path | None:
+    return get_heads(base_path).get(workflow, None)
+
+
+#########################################################################
+##                           Basic Write API                           ##
+#########################################################################
+
+
+def set_head(workflow: str, path: Path, base_path: str | Path):
+    base_path = Path(base_path)
+    base_path.mkdir(parents=True, exist_ok=True)
+    try:
+        with open(base_path / "heads", "rb") as f:
+            heads = pickle.load(f)
+    except:
+        heads = {}
+    heads[workflow] = path
+    with open(base_path / "heads", "wb") as f:
+        pickle.dump(heads, f)
+
+
+def save_report(workflow: str,
+                report: Report,
+                base_path: str | Path,
+                overwrite: bool = False,
+                refresh_heads: bool = True) -> int:
+    if workflow.startswith("cfg:"):
+        return save_config_key_history(workflow[4:], report, base_path)
+
+    logger.debug(
+        f'Saving report for "{workflow}", {report.in_spec=}, {report.bad_data=}, {report.fully_calibrated=}'
+    )
+    base_path = Path(base_path)
+    try:
+        buf = lzma.compress(pickle.dumps(report))
+    except:
+        raise ValueError(f"Can't pickle report for {workflow}")
+    if overwrite:
+        path = report.path
+        if path is None:
+            raise ValueError("Report path is None, can't overwrite.")
+        with open(base_path / 'reports' / path, "rb") as f:
+            index = int.from_bytes(f.read(8), 'big')
+        report.index = index
+    else:
+        path = random_path(base_path / 'reports')
+        (base_path / 'reports' / path).parent.mkdir(parents=True,
+                                                    exist_ok=True)
+        report.path = path
+        report.index = create_index("report",
+                                    base_path,
+                                    context=str(path),
+                                    width=35)
+    with open(base_path / 'reports' / path, "wb") as f:
+        f.write(report.index.to_bytes(8, 'big'))
+        f.write(buf)
+    if refresh_heads:
+        set_head(workflow, path, base_path)
+    return report.index
+
+
+def create_index(name: str,
+                 base_path: str | Path,
+                 context: str,
+                 width: int = -1,
+                 start: int = 0):
+
+    path = Path(base_path) / "index" / name
+    if width == -1:
+        width = len(context)
+    else:
+        width = max(width, len(context))
+
+    if path.with_suffix('.width').exists():
+        width = int(path.with_suffix('.width').read_text())
+        index = int(path.with_suffix('.seq').read_text())
+    else:
+        index = start
+    if width < len(context):
+        raise ValueError(
+            f"Context '{context}' is too long, existing width of '{name}' is {width}."
+        )
+    if not path.with_suffix('.width').exists():
+        path.parent.mkdir(parents=True, exist_ok=True)
+        path.with_suffix('.width').write_text(str(width))
+
+    path.with_suffix('.seq').write_text(str(index + 1))
+
+    with path.with_suffix('.idx').open("a") as f:
+        f.write(f"{context.ljust(width)}\n")
+
+    return index
+
+
+def save_item(item, group, data_path):
+    salt = 0
+    buf = pickle.dumps(item)
+    buf = lzma.compress(buf)
+    h = hashlib.md5(buf)
+
+    while True:
+        h.update(f"{salt}".encode())
+        hashstr = h.hexdigest()
+        item_id = Path(hashstr[:2]) / hashstr[2:4] / hashstr[4:]
+        path = Path(data_path) / group / item_id
+        if not path.exists():
+            path.parent.mkdir(parents=True, exist_ok=True)
+            with open(path, 'wb') as f:
+                f.write(buf)
+            break
+        elif path.read_bytes() == buf:
+            break
+        salt += 1
+    return str(item_id)
+
+
+def append_item_data(data, id, group, base_path):
+    path = Path(base_path) / group / id
+    if not path.exists():
+        raise ValueError(f"Item {id} does not exist.")
+    with open(path, 'ab') as f:
+        f.write(data)
+
+
+def save_config_key_history(key: str, report: Report,
+                            base_path: str | Path) -> int:
+    global __current_config_cache
+    base_path = Path(base_path) / 'state'
+    base_path.mkdir(parents=True, exist_ok=True)
+
+    if __current_config_cache is None:
+        if (base_path / 'parameters.pkl').exists():
+            with open(base_path / 'parameters.pkl', 'rb') as f:
+                __current_config_cache = pickle.load(f)
+        else:
+            __current_config_cache = {}
+
+    __current_config_cache[
+        key] = report.data, report.calibrated_time, report.checked_time
+
+    with open(base_path / 'parameters.pkl', 'wb') as f:
+        pickle.dump(__current_config_cache, f)
+    return 0
+
+
+#########################################################################
+##                            Basic Read API                           ##
+#########################################################################
+
+
+def load_report(path: str | Path, base_path: str | Path) -> Report | None:
+    if isinstance(base_path, str) and base_path.startswith('ssh://'):
+        with SSHClient() as client:
+            cfg = parse_ssh_uri(base_path)
+            remote_base_path = cfg.pop('remote_file_path')
+            client.load_system_host_keys()
+            client.connect(**cfg)
+            report = load_report_from_scp(path, remote_base_path, client)
+            report.base_path = base_path
+            return report
+
+    base_path = Path(base_path)
+    if zipfile.is_zipfile(base_path):
+        return load_report_from_zipfile(path, base_path)
+
+    path = base_path / 'reports' / path
+
+    with open(base_path / 'reports' / path, "rb") as f:
+        index = int.from_bytes(f.read(8), 'big')
+        report = pickle.loads(lzma.decompress(f.read()))
+        report.base_path = base_path
+        report.index = index
+        return report
+
+
+def get_heads(base_path: str | Path) -> Path | None:
+    if isinstance(base_path, str) and base_path.startswith('ssh://'):
+        with SSHClient() as client:
+            cfg = parse_ssh_uri(base_path)
+            remote_base_path = cfg.pop('remote_file_path')
+            client.load_system_host_keys()
+            client.connect(**cfg)
+            return get_heads_from_scp(remote_base_path, client)
+
+    base_path = Path(base_path)
+    if zipfile.is_zipfile(base_path):
+        return get_heads_from_zipfile(base_path)
+    try:
+        with open(base_path / "heads", "rb") as f:
+            heads = pickle.load(f)
+        return heads
+    except:
+        return {}
+
+
+@lru_cache(maxsize=4096)
+def query_index(name: str, base_path: str | Path, index: int):
+    if isinstance(base_path, str) and base_path.startswith('ssh://'):
+        with SSHClient() as client:
+            cfg = parse_ssh_uri(base_path)
+            remote_base_path = cfg.pop('remote_file_path')
+            client.load_system_host_keys()
+            client.connect(**cfg)
+            return query_index_from_scp(name, remote_base_path, client, index)
+
+    base_path = Path(base_path)
+    if zipfile.is_zipfile(base_path):
+        return query_index_from_zipfile(name, base_path, index)
+    path = Path(base_path) / "index" / name
+    width = int(path.with_suffix('.width').read_text())
+
+    with path.with_suffix('.idx').open("r") as f:
+        f.seek(index * (width + 1))
+        context = f.read(width)
+    return context.rstrip()
+
+
+@lru_cache(maxsize=4096)
+def load_item(id, group, base_path):
+    if isinstance(base_path, str) and base_path.startswith('ssh://'):
+        with SSHClient() as client:
+            cfg = parse_ssh_uri(base_path)
+            remote_base_path = cfg.pop('remote_file_path')
+            client.load_system_host_keys()
+            client.connect(**cfg)
+            buf = load_item_buf_from_scp(id, group, remote_base_path, client)
+    else:
+        base_path = Path(base_path)
+        if zipfile.is_zipfile(base_path):
+            buf = load_item_buf_from_zipfile(id, group, base_path)
+        else:
+            path = Path(base_path) / group / id
+            with open(path, 'rb') as f:
+                buf = f.read()
+    item = pickle.loads(lzma.decompress(buf))
+    return item
+
+
+def find_config_key_history(key: str, base_path: str | Path) -> Report | None:
+    global __current_config_cache
+    base_path = Path(base_path) / 'state'
+    if __current_config_cache is None:
+        if (base_path / 'parameters.pkl').exists():
+            with open(base_path / 'parameters.pkl', 'rb') as f:
+                __current_config_cache = pickle.load(f)
+        else:
+            __current_config_cache = {}
+
+    if key in __current_config_cache:
+        value, calibrated_time, checked_time = __current_config_cache.get(
+            key, None)
+        report = Report(
+            workflow=f'cfg:{key}',
+            bad_data=False,
+            in_spec=True,
+            fully_calibrated=True,
+            parameters={key: value},
+            data=value,
+            calibrated_time=calibrated_time,
+            checked_time=checked_time,
+        )
+        return report
+    return None
+
+
+#########################################################################
+##                            Zipfile support                          ##
+#########################################################################
+
+
+def load_report_from_zipfile(path: str | Path,
+                             base_path: str | Path) -> Report | None:
+    path = Path(path)
+    with zipfile.ZipFile(base_path) as zf:
+        path = '/'.join(path.parts)
+        with zf.open(f"{base_path.stem}/reports/{path}") as f:
+            index = int.from_bytes(f.read(8), 'big')
+            report = pickle.loads(lzma.decompress(f.read()))
+            report.base_path = base_path
+            report.index = index
+            return report
+
+
+def get_heads_from_zipfile(base_path: str | Path) -> Path | None:
+    with zipfile.ZipFile(base_path) as zf:
+        with zf.open(f"{base_path.stem}/heads") as f:
+            heads = pickle.load(f)
+    return heads
+
+
+def query_index_from_zipfile(name: str, base_path: str | Path, index: int):
+    with zipfile.ZipFile(base_path) as zf:
+        with zf.open(f"{base_path.stem}/index/{name}.width") as f:
+            width = int(f.read().decode())
+        with zf.open(f"{base_path.stem}/index/{name}.idx") as f:
+            f.seek(index * (width + 1))
+            context = f.read(width).decode()
+    return context.rstrip()
+
+
+def load_item_buf_from_zipfile(id, group, base_path):
+    with zipfile.ZipFile(base_path) as zf:
+        with zf.open(f"{base_path.stem}/{group}/{id}") as f:
+            return f.read()
+
+
+#########################################################################
+##                             SCP support                             ##
+#########################################################################
+
+
+def parse_ssh_uri(uri):
+    """
+    解析 SSH URI 字符串，返回包含连接参数和路径的字典。
+    
+    格式：ssh://[{username}[:{password}]@]{host}[:{port}][?key_filename={key_path}][/{remote_file_path}]
+
+    返回示例：
+    {
+        "username": "user",
+        "password": "pass",
+        "host": "example.com",
+        "port": 22,
+        "key_filename": "/path/to/key",
+        "remote_file_path": "/data/file.txt"
+    }
+    """
+    pattern = re.compile(
+        r"^ssh://"  # 协议头
+        r"(?:([^:@/]+))(?::([^@/]+))?@?"  # 用户名和密码（可选）
+        r"([^:/?#]+)"  # 主机名（必须）
+        r"(?::(\d+))?"  # 端口（可选）
+        r"(/?[^?#]*)?"  # 远程路径（可选）
+        r"(?:\?([^#]+))?"  # 查询参数（如 key_filename）
+        r"$",
+        re.IGNORECASE)
+
+    match = pattern.match(uri)
+    if not match:
+        raise ValueError(f"Invalid SSH URI format: {uri}")
+
+    # 提取分组
+    username, password, host, port, path, query = match.groups()
+
+    # 处理查询参数
+    key_filename = None
+    if query:
+        params = parse_qs(query)
+        key_filename = params.get("key_filename", [None])[0]  # 取第一个值
+
+    # 清理路径开头的斜杠
+    remote_file_path = path
+
+    return {
+        "username": username,
+        "password": password,
+        "hostname": host,
+        "port": int(port) if port else 22,  # 默认端口 22
+        "key_filename": key_filename,
+        "remote_file_path": remote_file_path
+    }
+
+
+def load_report_from_scp(path: str | Path, base_path: Path,
+                         client: SSHClient) -> Report:
+    try:
+        path = Path(path)
+        with client.open_sftp() as sftp:
+            with sftp.open(str(Path(base_path) / 'reports' / path), 'rb') as f:
+                index = int.from_bytes(f.read(8), 'big')
+                report = pickle.loads(lzma.decompress(f.read()))
+                report.index = index
+                return report
+    except SSHException:
+        raise ValueError(f"Can't load report from {path}")
+
+
+def get_heads_from_scp(base_path: Path, client: SSHClient) -> Path | None:
+    try:
+        with client.open_sftp() as sftp:
+            with sftp.open(str(Path(base_path) / 'heads'), 'rb') as f:
+                heads = pickle.load(f)
+        return heads
+    except SSHException:
+        return None
+
+
+def query_index_from_scp(name: str, base_path: Path, client: SSHClient,
+                         index: int):
+    try:
+        with client.open_sftp() as sftp:
+            s = str(Path(base_path) / 'index' / f'{name}.width')
+            with sftp.open(s, 'rb') as f:
+                width = int(f.read().decode())
+            with sftp.open(str(Path(base_path) / 'index' / f'{name}.idx'),
+                           'rb') as f:
+                f.seek(index * (width + 1))
+                context = f.read(width).decode()
+        return context.rstrip()
+    except SSHException:
+        return None
+
+
+def load_item_buf_from_scp(id: str, group: str, base_path: Path,
+                           client: SSHClient):
+    try:
+        with client.open_sftp() as sftp:
+            with sftp.open(str(Path(base_path) / group / str(id)), 'rb') as f:
+                return f.read()
+    except SSHException:
+        return None