recurvedata-lib 0.1.487__py2.py3-none-any.whl

recurvedata/executors/utils.py

@@ -0,0 +1,172 @@
+import asyncio
+import contextlib
+import datetime
+import logging
+import os
+from typing import TYPE_CHECKING, Callable, Optional
+
+import pendulum
+
+from recurvedata.config import EXECUTOR_META_PATH
+from recurvedata.error_codes import ERR
+from recurvedata.exceptions import RecurveException, TimeoutException, WrapRecurveException
+from recurvedata.executors.consts import VAR_CONVERT_STRING_FUNCS, VariableType
+from recurvedata.utils import json_dumps, json_loads
+from recurvedata.utils.timeout import timeout
+
+if TYPE_CHECKING:
+    from recurvedata.executors.schemas import ResponseModel
+
+
+def convert_var_value_from_string(var_type, var_value):
+    """
+    the var value from front-end is in string type，
+    this function will transform var value to corresponding type
+    """
+    if not isinstance(var_value, str):
+        return var_value
+    func = VAR_CONVERT_STRING_FUNCS[var_type]
+    return func(var_value)
+
+
+def get_variable_type_by_value(value):
+    type_mappings = {
+        bool: VariableType.BOOLEAN,
+        int: VariableType.INT,
+        float: VariableType.FLOAT,
+        str: VariableType.STRING,
+        dict: VariableType.STRING,  # 先用 STRING 类型（key 非 str 情况下，JSON 类型会报错）
+        (datetime.datetime, datetime.date, pendulum.DateTime, pendulum.Date): VariableType.DATETIME,
+    }
+    for types, var_type in type_mappings.items():
+        if isinstance(value, types):
+            return var_type
+    return VariableType.STRING
+
+
+def format_meta_file_path(job_id: int, node_key: str, execution_date: datetime.datetime) -> str:
+    sub_path = os.path.join(str(job_id), node_key, execution_date.isoformat())
+    path = os.path.join(EXECUTOR_META_PATH, sub_path)
+    os.makedirs(path, exist_ok=True)
+    return os.path.join(path, "meta.json")
+
+
+def update_meta_file(job_id: int, node_key: str, execution_date: datetime.datetime, meta: dict):
+    if not meta:
+        return
+    file_path = format_meta_file_path(job_id, node_key, execution_date)
+    with open(file_path, "w") as f:
+        f.write(json_dumps(meta))
+
+
+def read_meta_file(
+    job_id: int, node_key: str, execution_date: datetime.datetime, delete_after_read: bool = False
+) -> Optional[dict]:
+    file_path = format_meta_file_path(job_id, node_key, execution_date)
+    if not os.path.exists(file_path):
+        return
+    with open(file_path, "r") as f:
+        meta = json_loads(f.read())
+    if delete_after_read:
+        with contextlib.suppress(OSError, TypeError, ValueError):
+            os.unlink(file_path)
+    return meta
+
+
+def get_airflow_run_id():
+    return os.environ.get("AIRFLOW_CTX_DAG_RUN_ID")
+
+
+def get_airflow_try_number():
+    return os.environ.get("AIRFLOW_CTX_TRY_NUMBER")
+
+
+def get_recurve_node_key():
+    return os.environ.get("RECURVE__NODE_KEY")
+
+
+def run_with_result_handling(func: Callable = None, ttl: int = None, result_filename: str = None, *args, **kwargs):
+    """Run a function with timeout and handle the result.
+
+    Args:
+        func:
+        ttl (int, optional): timeout in seconds.
+        result_filename: the file to dump the result.
+    """
+    from recurvedata.executors.schemas import ResponseError, ResponseModel
+
+    def exec_with_timeout(ttl: int):
+        with timeout(ttl):
+            return func(*args, **kwargs)
+
+    result = ResponseModel(ok=True)
+    try:
+        data = exec_with_timeout(ttl) if ttl else func(*args, **kwargs)
+        result.data = data
+    except Exception as e:
+        result.ok = False
+        if not isinstance(e, RecurveException):
+            e = WrapRecurveException(ERR.UNKNOWN_ERROR, e)
+        result.error = ResponseError.from_recurve_exception(e)
+
+    if result_filename:
+        result.model_dump_json_file(result_filename)
+    else:
+        logging.info(result.model_dump_json(indent=2))
+        return result
+
+
+async def run_with_result_handling_v2(func: Callable = None, ttl: int = None, *args, **kwargs) -> "ResponseModel":
+    """
+    compare with run_with_result_handling,
+    difference is the timeout logic.
+    timeout using signal cannot work on fastapi.
+    parameters:
+        func is a synchronous task.
+    """
+    from recurvedata.executors.schemas import ResponseError, ResponseModel
+
+    result = ResponseModel(ok=True)
+    try:
+        if asyncio.iscoroutinefunction(func):
+            # 如果 func 是一个协程函数，直接 await 它
+            data = await asyncio.wait_for(
+                func(*args, **kwargs),
+                timeout=ttl if ttl else None,
+            )
+        else:
+            # 否则，使用 asyncio.to_thread 运行同步函数
+            data = await asyncio.wait_for(
+                asyncio.to_thread(func, *args, **kwargs),
+                timeout=ttl if ttl else None,
+            )
+        result.data = data
+    except asyncio.TimeoutError:
+        raise TimeoutException(f"Function {func.__name__} timed out after {ttl} seconds")
+    except Exception as e:
+        result.ok = False
+        if not isinstance(e, RecurveException):
+            e = WrapRecurveException(ERR.UNKNOWN_ERROR, e)
+        result.error = ResponseError.from_recurve_exception(e)
+
+    return result
+
+
+def patch_pandas_mysql_connector_cext_missing():
+    """
+    Patch for MySQL Connector/Python C Extension issue.
+
+    When pandas is imported before mysql.connector, the MySQL Connector/Python C Extension
+    may be missing, which can cause connection errors like:
+    '2013: Lost connection to MySQL server during query' when compiling dbt Doris models.
+
+    This function attempts to preemptively import mysql.connector to ensure the C Extension
+    is properly loaded before pandas.
+    """
+    try:
+        # Attempt to import mysql.connector first to ensure C Extension is loaded
+        import mysql.connector  # noqa: F401
+    except ImportError:
+        # Silently continue if mysql.connector is not installed
+        # The error will be handled elsewhere if the connector is actually needed
+        pass

recurvedata/filestorage/__init__.py

@@ -0,0 +1,11 @@
+from recurvedata.filestorage._factory import Factory
+from recurvedata.filestorage.interface import AbstractFileStorage, StorageType
+
+factory = Factory[StorageType, AbstractFileStorage](
+    {
+        StorageType.LOCAL: "recurvedata.filestorage.backends.local.LocalStorage",
+        StorageType.OSS: "recurvedata.filestorage.backends.oss.OSSStorage",
+    }
+)
+
+__all__ = ["AbstractFileStorage", "StorageType", "factory"]

recurvedata/filestorage/_factory.py

@@ -0,0 +1,33 @@
+from typing import Any, Generic, TypeVar, cast
+
+from pydantic import ValidationError
+
+from recurvedata.filestorage.interface import AbstractFileStorage, StorageType
+from recurvedata.utils.imports import import_string
+
+_ST = TypeVar("_ST", bound=StorageType)
+_SC = TypeVar("_SC", bound=AbstractFileStorage)
+
+
+class Factory(Generic[_ST, _SC]):
+    def __init__(self, implementations: dict[_ST, str]):
+        self._implementations: dict[_ST, str] = implementations
+
+    def get_supported_backends(self) -> list[_ST]:
+        return list(self._implementations.keys())
+
+    def get_storage_class(self, type_: _ST | str):
+        return cast(type[_SC], import_string(self._implementations[type_]))
+
+    def create(self, type_: _ST | str, options: dict[str, Any]) -> _SC:
+        if type_ not in self._implementations:
+            raise ValueError(f"Unsupported storage backend: {type_}")
+
+        storage_class = self.get_storage_class(type_)
+
+        try:
+            obj = storage_class.from_params(**options)
+        except ValidationError as e:
+            raise ValueError(f"Invalid configuration for {type_}: {e}")
+
+        return obj

recurvedata/filestorage/backends/fsspec.py

@@ -0,0 +1,45 @@
+from abc import ABC, abstractmethod
+from functools import cached_property
+
+from fsspec.asyn import AsyncFileSystem
+
+from recurvedata.filestorage.interface import AbstractFileStorage
+
+
+class FSSpecAbstractStorage(AbstractFileStorage, ABC):
+    _fs: AsyncFileSystem
+
+    @cached_property
+    def _fs(self) -> AsyncFileSystem:
+        # TODO(liyangliang): we could consider involve a pooling mechanism here
+        # see discussion with ChatGPT https://chat.openai.com/share/972e3bcc-0ebc-43b8-9a49-72f3dd7dc2b6
+        return self.get_fs()
+
+    @abstractmethod
+    def get_fs(self) -> AsyncFileSystem:
+        ...
+
+    def normalize_path(self, path: str) -> str:
+        """Normalize a path to be used with the filesystem."""
+        return path
+
+    async def listdir(self, path: str) -> list[str]:
+        return await self._fs._ls(self.normalize_path(path))
+
+    async def write_bytes(self, path: str, content: bytes):
+        await self._fs._pipe_file(self.normalize_path(path), content)
+
+    async def read_bytes(self, path: str) -> bytes:
+        return await self._fs._cat_file(self.normalize_path(path))
+
+    async def put(self, local_path: str, remote_path: str):
+        await self._fs._put_file(local_path, self.normalize_path(remote_path))
+
+    async def get(self, remote_path: str, local_path: str):
+        await self._fs._get_file(self.normalize_path(remote_path), local_path)
+
+    async def delete(self, path: str):
+        await self._fs._rm(self.normalize_path(path))
+
+    async def exists(self, path: str) -> bool:
+        return await self._fs._exists(self.normalize_path(path))

recurvedata/filestorage/backends/local.py

@@ -0,0 +1,67 @@
+import datetime
+import urllib.parse
+from typing import Any
+
+from jose import JWTError, jwt
+from morefs.asyn_local import AsyncLocalFileSystem
+from pydantic import SecretStr
+
+from recurvedata.filestorage.backends.fsspec import FSSpecAbstractStorage
+from recurvedata.filestorage.interface import StorageConfig, StorageType
+from recurvedata.utils.date_time import utcnow
+
+ALGORITHM = "HS256"
+_DEFAULT_SECRET_KEY = SecretStr("619805f2af666a623f37221ce8dfbec85ce9e83a16b20fe4a424078ed37f2a3a")
+
+
+class LocalStorageConfig(StorageConfig):
+    root_dir: str = "/tmp"
+    auto_mkdir: bool = True
+
+    # below are required for generating presigned url
+    secret_key: SecretStr = _DEFAULT_SECRET_KEY
+    """
+    The secret key to sign and verify the presigned url. Although it's better to
+    use a more secure key, for simplicity, we assign a default value here.
+    To generate a secure key, you can use the following code:
+
+        ```python
+        import secrets
+        key = secrets.token_hex(32)
+        print(key)
+        ```
+    """
+    server_base_url: str = None
+
+
+class LocalStorage(FSSpecAbstractStorage):
+    config_class = LocalStorageConfig
+    config: LocalStorageConfig
+    _fs: AsyncLocalFileSystem
+
+    @classmethod
+    def storage_type(cls) -> StorageType:
+        return StorageType.LOCAL
+
+    def get_fs(self) -> AsyncLocalFileSystem:
+        return AsyncLocalFileSystem(auto_mkdir=self.config.auto_mkdir)
+
+    def normalize_path(self, path: str) -> str:
+        return self.join_path(self.config.root_dir, path)
+
+    def get_presigned_url(self, path: str, expiration: int = 1800, **kwargs) -> str:
+        if not all((self.config.secret_key, self.config.server_base_url)):
+            raise ValueError("secret_key and server_base_url are required to generate presigned url")
+
+        to_encode = {"path": path, "exp": utcnow() + datetime.timedelta(seconds=expiration)}
+        encoded = jwt.encode(to_encode, self.config.secret_key.get_secret_value(), algorithm=ALGORITHM)
+        base_url = str(self.config.server_base_url).rstrip("/")
+        query_string = urllib.parse.urlencode({"token": encoded} | kwargs)
+        return f"{base_url}/{path}?{query_string}"
+
+    def decode_presigned_url(self, token: str) -> dict[str, Any] | None:
+        try:
+            payload = jwt.decode(token, self.config.secret_key.get_secret_value(), algorithms=[ALGORITHM])
+            return payload
+        except JWTError:
+            return None

recurvedata/filestorage/backends/oss.py

@@ -0,0 +1,56 @@
+from functools import cached_property
+
+import oss2
+from ossfs.async_oss import AioOSSFileSystem
+from pydantic import ConfigDict
+
+from recurvedata.filestorage.backends.fsspec import FSSpecAbstractStorage
+from recurvedata.filestorage.interface import StorageConfig, StorageType
+
+
+class OSSStorageConfig(StorageConfig):
+    endpoint: str
+    access_key_id: str
+    access_key_secret: str
+    bucket_name: str
+    directory: str = ""
+    security_token: str | None = None  # for sts token
+
+    model_config = ConfigDict(extra="allow")
+
+
+class OSSStorage(FSSpecAbstractStorage):
+    config_class = OSSStorageConfig
+    config: OSSStorageConfig
+    _fs: AioOSSFileSystem
+
+    @classmethod
+    def storage_type(cls) -> StorageType:
+        return StorageType.OSS
+
+    def get_fs(self) -> AioOSSFileSystem:
+        return AioOSSFileSystem(
+            endpoint=self.config.endpoint,
+            key=self.config.access_key_id,
+            secret=self.config.access_key_secret,
+            token=self.config.security_token,
+        )
+
+    def normalize_path(self, path: str) -> str:
+        return self.join_path(self.config.bucket_name, self.config.directory, path)
+
+    @cached_property
+    def public_bucket(self) -> oss2.Bucket:
+        """The public Bucket object, used for generating public download urls."""
+        if "-internal" not in self.config.endpoint:
+            public_endpoint = self.config.endpoint
+        else:
+            public_endpoint = self.config.endpoint.replace("-internal", "")
+        auth = oss2.Auth(self.config.access_key_id, self.config.access_key_secret)
+        return oss2.Bucket(auth, public_endpoint, self.config.bucket_name)
+
+    def get_presigned_url(self, path: str, expiration: int = 1800, **kwargs) -> str:
+        headers = {"content-disposition": "attachment"}
+        return self.public_bucket.sign_url(
+            "GET", self.join_path(self.config.directory, path), expiration, headers=headers, slash_safe=True
+        )

recurvedata/filestorage/interface.py

@@ -0,0 +1,84 @@
+import os
+from abc import ABC, abstractmethod
+from enum import Enum
+from typing import ClassVar, Self
+
+from pydantic import BaseModel, ConfigDict
+
+
+class StorageType(str, Enum):
+    LOCAL = "local"
+    S3 = "s3"
+    OSS = "oss"
+
+    def __str__(self) -> str:
+        return str.__str__(self)
+
+
+class StorageConfig(BaseModel):
+    """Base class for storage configurations."""
+
+    model_config = ConfigDict(extra="forbid")
+
+
+class AbstractFileStorage(ABC):
+    config_class: ClassVar[type[StorageConfig]] = StorageConfig
+    config: StorageConfig
+
+    def __init__(self, config: StorageConfig):
+        self.config = config
+
+    @classmethod
+    def from_params(cls, **kwargs) -> Self:
+        return cls(cls.config_class.model_validate(kwargs))
+
+    @classmethod
+    @abstractmethod
+    def storage_type(cls) -> StorageType:
+        """Return the storage type."""
+        ...
+
+    @abstractmethod
+    async def listdir(self, path: str) -> list[str]:
+        """List the contents of a directory."""
+        ...
+
+    @abstractmethod
+    async def write_bytes(self, path: str, content: bytes):
+        """Write bytes content to a destination path."""
+        ...
+
+    @abstractmethod
+    async def read_bytes(self, path: str) -> bytes:
+        """Read bytes content from a source path."""
+        ...
+
+    @abstractmethod
+    async def put(self, local_path: str, remote_path: str):
+        """Upload a file from a source path to a destination path."""
+        ...
+
+    @abstractmethod
+    async def get(self, remote_path: str, local_path: str):
+        """Download a file from a source path to a destination path."""
+        ...
+
+    @abstractmethod
+    async def delete(self, path: str):
+        """Delete a file from the given path."""
+        ...
+
+    @abstractmethod
+    async def exists(self, path: str) -> bool:
+        """Check if a file exists at the given path."""
+        ...
+
+    @abstractmethod
+    def get_presigned_url(self, path: str, expiration: int = 1800, **kwargs) -> str:
+        """Generate a presigned URL for a file for temporary access."""
+        ...
+
+    @staticmethod
+    def join_path(*parts: str) -> str:
+        """Join path parts together."""
+        return os.path.normpath(os.path.join(*parts))

recurvedata/operators/__init__.py

@@ -0,0 +1,10 @@
+from recurvedata.operators.context import context  # noqa
+from recurvedata.operators.dbt_operator import DBTOperator
+from recurvedata.operators.link_operator import LinkOperator
+from recurvedata.operators.notify_operator import NotifyOperator
+from recurvedata.operators.operator import get_operator_class  # noqa
+from recurvedata.operators.python_operator import PythonOperator
+from recurvedata.operators.sensor_operator import SensorOperator
+from recurvedata.operators.sql_operator import SQLOperator
+from recurvedata.operators.transfer_operator import TransferOperator
+from recurvedata.operators.web_init import init_operator_web

recurvedata/operators/base.py

@@ -0,0 +1,28 @@
+from typing import Any
+
+import jsonschema
+
+from recurvedata.core.translation import convert_lazy_string
+
+
+class Configurable(object):
+    enabled = True
+
+    @classmethod
+    def config_schema(cls) -> dict:
+        return {}
+
+    @classmethod
+    def name(cls) -> str:
+        return cls.__name__
+
+    @classmethod
+    def to_dict(cls) -> dict:
+        return {"name": cls.name(), "config_schema": cls.config_schema()}
+
+    @classmethod
+    def validate(cls, configuration: dict[str, Any]) -> dict[str, Any]:
+        schema = cls.config_schema()
+        schema = convert_lazy_string(schema)
+        jsonschema.validate(configuration, schema)
+        return configuration

recurvedata/operators/config.py

@@ -0,0 +1,21 @@
+import os
+
+from recurvedata.operators.utils.fs import get_exist_path
+
+
+class Config(object):  # todo: move to somewhere else
+    DATA_ROOT = "/opt/recurve/worker_data"
+
+    PYENV_BASE_BIN_PATH = os.path.join(
+        get_exist_path([os.environ.get("RECURVE__WORKER__PYENV__BASE"), "~/.pyenv"]) or "/opt/pyenv",
+        "versions/{pyenv}/bin",
+    )
+    PYENV_PYTHON_PATH = os.path.join(PYENV_BASE_BIN_PATH, "python")
+
+    RECURVE_EXECUTOR_PYENV_BIN_PATH = PYENV_BASE_BIN_PATH.format(pyenv="recurve_executor")
+
+    REDIS_URL = "redis://localhost:6381/13"
+
+
+# CONF 是一个全局对象，用于获取配置项。
+CONF = Config()