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

recurvedata/error_codes.py

@@ -0,0 +1,71 @@
+"""
+copied from recurve-server
+"""
+
+from enum import Enum, EnumMeta
+
+
+class ErrorCodeMeta(EnumMeta):
+    _error_codes = set()
+
+    def __new__(metacls, clsname, bases, classdict):
+        enum_members = {k: v for k, v in classdict.items() if not k.startswith("_")}
+        for name, code in enum_members.items():
+            if type(code) is not tuple:
+                continue
+            for error_code in metacls._error_codes:
+                if code[0] == error_code[1][0]:
+                    raise ValueError(f"Error code {code[0]} in {clsname} already exists globally")
+            metacls._error_codes.add((name, code))
+        return super().__new__(metacls, clsname, bases, classdict)
+
+    @classmethod
+    def error_codes(cls):
+        return sorted(list(cls._error_codes), key=lambda x: x[1][0])
+
+
+class BaseErrorCode(Enum, metaclass=ErrorCodeMeta):
+    @property
+    def code(self):
+        return self.value[0]
+
+    @property
+    def message(self):
+        return self.value[1]
+
+    def to_dict(self):
+        return {"code": self.code, "msg": self.message}
+
+    def exception(self, data: dict | None = None):
+        from recurvedata.exceptions import RecurveException
+
+        return RecurveException(data=data, code=self)
+
+
+class ErrorCode(BaseErrorCode):
+    # 00: General
+    UNAUTHORIZED = ("A0001", "Unauthorized")
+    PERMISSION_DENIED = ("A0003", "Permission Denied")
+    NOT_FOUND = ("A0004", "Not Found")
+    ALREADY_EXISTS = ("A0008", "Already Exists")
+    INVALID_ARGUMENT = ("A0010", "Invalid Argument")
+    MAX_FILE_SIZE_EXCEEDED = ("A0011", "Max File Size Exceeded")
+    INVALID_FILE_FORMAT = ("A0012", "Invalid File Format")
+    INVALID_READONLY_QUERY = ("A0013", "Query Should be Read Only")
+    INTERNAL_SERVER_ERROR = ("B0001", "Internal Server Error")
+    NOT_IMPLEMENTED = ("B0002", "Not Implemented")
+    FAILED_TO_PARSE_QUERY = ("B0003", "Failed to Parse Query")
+
+    UNKNOWN_ERROR = ("D0001", "Unknown Error")
+
+    # agent task
+    API_REQUEST_FAILED = ("A1202", "API request failed")
+    MAX_RETRIES_EXCEEDED = ("A1203", "max retries exceeded")
+    TIMEOUT = ("A1204", "time out")
+
+    # data service
+    DP_FETCH_CONNECTION_FAILED = ("A1401", "DP fetch connection failed")
+    PREVIEW_DATA_FAILED = ("A1402", "Preview data failed")
+
+
+ERR = ErrorCode  # shortcut

recurvedata/exceptions.py

@@ -0,0 +1,72 @@
+from functools import wraps
+
+from recurvedata.error_codes import ERR, BaseErrorCode
+
+
+class RecurveException(Exception):
+    _default_code: ERR = ERR.UNKNOWN_ERROR
+
+    def __init__(self, data: dict | str = None, code: BaseErrorCode = None):
+        self.code = code or self._default_code
+        self.data = data
+
+    def to_dict(self) -> dict:
+        return self.code.to_dict() | {"data": self.data}
+
+
+class InvalidArgument(RecurveException):
+    _default_code: ERR = ERR.INVALID_ARGUMENT
+
+
+class APIError(RecurveException):
+    """Raised when an API request fails."""
+
+    _default_code: ERR = ERR.API_REQUEST_FAILED
+
+
+class UnauthorizedError(RecurveException):
+    """Raised when an unauthorized request is made."""
+
+    _default_code = ERR.UNAUTHORIZED
+
+
+class MaxRetriesExceededException(RecurveException):
+    """Raised when the maximum number of retries is exceeded."""
+
+    _default_code: ERR = ERR.MAX_RETRIES_EXCEEDED
+
+
+class TimeoutException(RecurveException):
+    """Raised when a timeout occurs."""
+
+    _default_code: ERR = ERR.TIMEOUT
+
+
+class WrapRecurveException(RecurveException):
+    """
+    raised in wrap_error function
+    """
+
+    def __init__(self, code: ERR, exception: Exception, data: dict | str = None):
+        super().__init__(data=data, code=code)
+        self.exception: Exception = exception
+
+    def to_dict(self) -> dict:
+        return self.code.to_dict() | {"data": self.data, "exception": str(self.exception)}
+
+
+def wrap_error(err_code: ERR):
+    def decorator(func):
+        @wraps(func)
+        def wrapper(*args, **kwargs):
+            try:
+                return func(*args, **kwargs)
+            except RecurveException:
+                raise
+            except Exception as e:
+                wrapped_error = WrapRecurveException(err_code, e)
+                raise wrapped_error
+
+        return wrapper
+
+    return decorator

recurvedata/executors/__init__.py

@@ -0,0 +1,4 @@
+from recurvedata.executors.cli import cli
+from recurvedata.executors.debug_executor import DebugExecutor
+from recurvedata.executors.executor import Executor
+from recurvedata.executors.link_executor import LinkExecutor

recurvedata/executors/cli/__init__.py

@@ -0,0 +1,7 @@
+from recurvedata.executors.utils import patch_pandas_mysql_connector_cext_missing
+
+patch_pandas_mysql_connector_cext_missing()
+
+from recurvedata.executors.cli.main import cli  # noqa: E402
+
+__all__ = ["cli"]

recurvedata/executors/cli/connector.py

@@ -0,0 +1,117 @@
+import json
+
+import typer
+
+from recurvedata.executors.service.connector import ConnectionService
+from recurvedata.executors.utils import run_with_result_handling
+from recurvedata.utils import init_logging
+from recurvedata.utils._typer import RecurveTyper
+
+cli = RecurveTyper()
+
+
+@cli.callback()
+def init():
+    init_logging()
+
+
+@cli.command()
+def test_connection(
+    connection_type: str = typer.Option(
+        ..., "--connection-type", help="Type of the connection (e.g., 'mysql', 'postgresql', 'snowflake')"
+    ),
+    config: str = typer.Option(..., "--config", help="JSON string of connection configuration"),
+    result_filename: str = typer.Option(None, "--result-filename", help="Filename to save the json result"),
+    timeout: int = typer.Option(30, "--timeout", help="Timeout for the connection test in seconds"),
+):
+    """
+    Test if a connection is valid
+    """
+    connection_config = json.loads(config)
+    run_with_result_handling(
+        ConnectionService.test_connection, timeout, result_filename, connection_type, connection_config
+    )
+
+
+@cli.command()
+async def list_databases(
+    connection_type: str = typer.Option(
+        ..., "--connection-type", help="Type of the connection (e.g., 'mysql', 'postgresql', 'snowflake')"
+    ),
+    config: str = typer.Option(..., "--config", help="JSON string of connection configuration"),
+    result_filename: str = typer.Option(None, "--result-filename", help="Filename to save the json result"),
+):
+    """
+    List databases for a given connection
+    """
+    connection_config = json.loads(config)
+    run_with_result_handling(
+        ConnectionService.list_databases,
+        result_filename=result_filename,
+        connection_type=connection_type,
+        connection_config=connection_config,
+    )
+
+
+@cli.command()
+async def list_tables(
+    connection_type: str = typer.Option(
+        ..., "--connection-type", help="Type of the connection (e.g., 'mysql', 'postgresql', 'snowflake')"
+    ),
+    config: str = typer.Option(..., "--config", help="JSON string of connection configuration"),
+    database: str = typer.Option(..., "--database", help="Database name"),
+    result_filename: str = typer.Option(None, "--result-filename", help="Filename to save the json result"),
+):
+    """List tables for a given connection and database"""
+    connection_config = json.loads(config)
+    run_with_result_handling(
+        ConnectionService.list_tables,
+        result_filename=result_filename,
+        connection_type=connection_type,
+        connection_config=connection_config,
+        database=database,
+    )
+
+
+@cli.command()
+async def list_columns(
+    connection_type: str = typer.Option(
+        ..., "--connection-type", help="Type of the connection (e.g., 'mysql', 'postgresql', 'snowflake')"
+    ),
+    config: str = typer.Option(..., "--config", help="JSON string of connection configuration"),
+    database: str = typer.Option(..., "--database", help="Database name"),
+    table: str = typer.Option(..., "--table", help="Table name"),
+    result_filename: str = typer.Option(None, "--result-filename", help="Filename to save the json result"),
+):
+    """List columns for a given connection, database and table"""
+    connection_config = json.loads(config)
+    run_with_result_handling(
+        ConnectionService.list_columns,
+        result_filename=result_filename,
+        connection_type=connection_type,
+        connection_config=connection_config,
+        database_name=database,
+        table_name=table,
+    )
+
+
+@cli.command()
+async def list_full_databases(
+    connection_type: str = typer.Option(
+        ..., "--connection-type", help="Type of the connection (e.g., 'mysql', 'postgresql', 'snowflake')"
+    ),
+    config: str = typer.Option(..., "--config", help="JSON string of connection configuration"),
+    result_filename: str = typer.Option(None, "--result-filename", help="Filename to save the json result"),
+):
+    """List full databases for a given connection"""
+    connection_config = json.loads(config)
+    run_with_result_handling(
+        ConnectionService.list_full_databases,
+        result_filename=result_filename,
+        connection_type=connection_type,
+        connection_config=connection_config,
+    )
+
+
+if __name__ == "__main__":
+    cli()

recurvedata/executors/cli/dbt.py

@@ -0,0 +1,118 @@
+import logging
+
+from recurvedata.client.server_client import ServerDbtClient
+from recurvedata.core.tracing import Tracing
+from recurvedata.executors.cli import parameters as param
+from recurvedata.executors.utils import run_with_result_handling
+from recurvedata.utils import init_logging
+from recurvedata.utils._typer import RecurveTyper
+from recurvedata.utils.sql import trim_replace_special_character
+
+dbt_tracer = Tracing()
+cli = RecurveTyper()
+logger = logging.getLogger(__name__)
+
+
+@cli.callback()
+def init():
+    from recurvedata.core.tracing import Tracing
+
+    if not Tracing.is_instantiated():
+        from recurvedata.utils.tracing import create_dp_tracer
+
+        create_dp_tracer("recurve-lib-dbt")
+    init_logging()
+
+
+@cli.command()
+@dbt_tracer.create_span(context_payload_name="tracing_context")
+async def compile(
+    project_id: int = param.project_id,
+    alias: str = param.alias,
+    sql: str = param.sql,
+    tracing_context: str = param.tracing_context,
+    force_regenerate_dir: bool = param.force_regenerate_dir,
+    result_filename: str = param.result_filename,
+):
+    try:
+        _compile_using_server(project_id, alias, sql, force_regenerate_dir, result_filename)
+    except Exception as e:
+        logger.exception(f"_compile_using_server failed: {e}, try _compile_using_service")
+        _compile_using_service(project_id, alias, sql, force_regenerate_dir, result_filename)
+
+
+@dbt_tracer.create_span()
+def _compile_using_server(project_id: int, alias: str, sql: str, force_regenerate_dir: bool, result_filename: str):
+    client = ServerDbtClient()
+    res = client.compile(project_id, sql, alias, force_regenerate_dir)
+    if result_filename:
+        res.model_dump_json_file(result_filename)
+
+
+@dbt_tracer.create_span()
+def _compile_using_service(
+    project_id: int, alias: str, sql: str, force_regenerate_dir: bool, result_filename: str = param.result_filename
+):
+    from recurvedata.dbt.service import DbtService
+
+    sql = trim_replace_special_character(sql)
+    service = DbtService(
+        project_id=project_id,
+        project_connection_name=alias,
+        force_regenerate_dir=force_regenerate_dir,
+        need_fetch_variable=True,
+    )
+    run_with_result_handling(service.compile, inline_sql=sql, result_filename=result_filename)
+
+
+@cli.command()
+async def preview(
+    project_id: int = param.project_id,
+    alias: str = param.alias,
+    sql: str = param.sql,
+    limit: int = param.limit,
+    no_data: bool = param.no_data,
+    force_regenerate_dir: bool = param.force_regenerate_dir,
+    result_filename: str = param.result_filename,
+):
+    try:
+        _preview_using_server(project_id, alias, sql, limit, force_regenerate_dir, result_filename)  # todo: add no_data
+    except Exception as e:
+        logger.exception(f"_preview_using_server failed: {e}, try _preview_using_service")
+        _preview_using_service(project_id, alias, sql, limit, force_regenerate_dir, no_data, result_filename)
+
+
+def _preview_using_server(
+    project_id: int, alias: str, sql: str, limit: int, force_regenerate_dir: bool, result_filename: str
+):
+    client = ServerDbtClient()
+    res = client.preview(project_id, sql, alias, limit, force_regenerate_dir)
+    if result_filename:
+        res.model_dump_json_file(result_filename)
+
+
+def _preview_using_service(
+    project_id: int,
+    alias: str,
+    sql: str,
+    limit: int,
+    force_regenerate_dir: bool,
+    no_data: bool,
+    result_filename: str = param.result_filename,
+):
+    from recurvedata.dbt.service import DbtService
+
+    sql = trim_replace_special_character(sql, strip_sufix=True)
+    service = DbtService(
+        project_id=project_id,
+        project_connection_name=alias,
+        force_regenerate_dir=force_regenerate_dir,
+        need_fetch_variable=True,
+    )
+    run_with_result_handling(
+        service.preview, inline_sql=sql, limit=limit, result_filename=result_filename, no_data=no_data
+    )
+
+
+if __name__ == "__main__":
+    cli()

recurvedata/executors/cli/main.py

@@ -0,0 +1,82 @@
+import typer
+
+from recurvedata.consts import ScheduleType
+from recurvedata.executors.cli.connector import cli as connector_cli
+from recurvedata.executors.cli.dbt import cli as dbt_cli
+from recurvedata.executors.debug_executor import DebugExecutor
+from recurvedata.executors.executor import Executor
+from recurvedata.utils import init_logging
+from recurvedata.utils._typer import RecurveTyper
+from recurvedata.utils.date_time import astimezone
+
+cli = RecurveTyper()
+cli.add_typer(connector_cli, name="connector")
+cli.add_typer(dbt_cli, name="dbt")
+
+
+@cli.callback()
+def init():
+    from recurvedata.core.tracing import Tracing
+
+    if not Tracing.is_instantiated():
+        from recurvedata.utils.tracing import create_dp_tracer
+
+        create_dp_tracer("recurve-lib")
+    init_logging()
+
+
+@cli.command()
+def execute(
+    dag_slug: str = typer.Option(..., "--dag_slug", help="Unique identifier for the DAG"),
+    node_slug: str = typer.Option(..., "--node_slug", help="Unique identifier for the node"),
+    execution_date: str = typer.Option(..., "--execution_date", help="Date/time of execution"),
+    # RecurveLink settings
+    link_workflow_id: int = typer.Option(None, "--link_workflow_id", help="ID of linked workflow"),
+    link_node_id: int = typer.Option(None, "--link_node_id", help="ID of linked node"),
+    is_link_workflow: bool = typer.Option(False, "--is_link_workflow", help="Whether this is a linked workflow"),
+):
+    """
+    Execute a specific node in a DAG at the given execution date.
+
+    Optionally configure workflow linking settings.
+    """
+    # Initialize and run the executor
+    executor = Executor(dag_slug=dag_slug, node_slug=node_slug, execution_date=execution_date)
+    executor.set_link_settings(
+        link_workflow_id=link_workflow_id, link_node_id=link_node_id, is_link_workflow=is_link_workflow
+    )
+    executor.run()
+
+
+@cli.command()
+def debug(
+    workflow_id: int = typer.Option(..., "--workflow_id", help="ID of the workflow to debug"),
+    node_key: str = typer.Option(..., "--node_key", help="Key identifier of the node to debug"),
+    schedule_type: ScheduleType = typer.Option(..., "--schedule_type", help="Type of schedule"),
+    schedule_interval: str = typer.Option(..., "--schedule_interval", help="Schedule interval specification"),
+    execution_date: str = typer.Option(..., "--execution_date", help="Execution timestamp"),
+    timezone: str = typer.Option(..., "--timezone", help="Timezone for execution"),
+    celery_task_id: str = typer.Option(..., "--celery_task_id", help="Celery task ID for tracking"),
+):
+    """
+    Debug a workflow node by executing it in isolation.
+
+    This command allows debugging a specific node from a workflow by running it independently.
+    The execution context (schedule, timing, etc) can be controlled via the parameters.
+    """
+    # execution_date = ensure_datetime(execution_date).replace(tzinfo=ensure_tz(timezone))
+    execution_date = astimezone(execution_date, timezone)
+    executor = DebugExecutor(
+        workflow_id=workflow_id,
+        node_key=node_key,
+        schedule_type=schedule_type,
+        schedule_interval=schedule_interval,
+        execution_date=execution_date,
+        timezone=timezone,
+        celery_task_id=celery_task_id,
+    )
+    executor.run()
+
+
+if __name__ == "__main__":
+    cli()

recurvedata/executors/cli/parameters.py

@@ -0,0 +1,18 @@
+from typer import Option
+
+connection_type = Option(
+    ..., "--connection-type", help="Type of the connection (e.g., 'mysql', 'postgresql', 'snowflake')"
+)
+connection_config = Option(..., "--config", help="JSON string of connection configuration")
+result_filename = Option(None, "--result-filename", help="Filename to save the json result")
+timeout = Option(30, "--timeout", help="Timeout for the connection test in seconds")
+database = Option(..., "--database", help="Database name")
+table = Option(..., "--table", help="Table name")
+project_id = Option(..., "--project-id", help="Project ID")
+connection_name = Option(..., "--connection-name", help="Connection Name")
+alias = Option(..., "--alias", help="Alias")  # project connection name
+limit = Option(100, "--limit")
+sql = Option(..., "--sql")
+tracing_context = Option(None, "--tracing-context", help="Tracing context")
+force_regenerate_dir = Option(False, "--force-regenerate-dir", help="Whether to force regenerate dbt project")
+no_data = Option(False, "--no-data", help="Whether to include data in preview result")

recurvedata/executors/client.py

@@ -0,0 +1,190 @@
+import datetime
+from typing import Any, Optional
+
+from recurvedata.client import Client
+from recurvedata.config import RECURVE_EXECUTOR_PYENV_NAME
+from recurvedata.consts import ETLExecutionStatus, Operator
+from recurvedata.executors.schemas import (
+    ConnectionItem,
+    DebugEnd,
+    DebugLogRecord,
+    DebugStart,
+    JobNodeItem,
+    TaskInstanceEnd,
+    TaskInstanceStart,
+    TaskInstanceStartResponse,
+    TaskLogRecord,
+    WorkflowNodeItem,
+)
+from recurvedata.utils import get_env_id
+
+
+class ExecutorClient(Client):
+    def get_node(self, job_id: int, node_id: int) -> JobNodeItem:
+        params = {
+            "env_id": get_env_id(),
+            "job_id": job_id,
+        }
+
+        return self.request(
+            "GET", path=f"/api/executor/node/{node_id}", response_model_class=JobNodeItem, params=params
+        )
+
+    def get_connection(self, project_id: int, connection_name: str) -> ConnectionItem:
+        params = {
+            "env_id": get_env_id(),
+            "name": connection_name,
+            "project_id": project_id,
+        }
+        return self.request("GET", path="/api/executor/connection", response_model_class=ConnectionItem, params=params)
+
+    def task_instance_start(
+        self,
+        job_id: int,
+        node_id: int,
+        operator: Operator,
+        task: str,
+        execution_date: datetime.datetime,
+        rendered_config: dict,
+        start_time: datetime.datetime,
+        hostname: Optional[str],
+        pid: Optional[int],
+    ):
+        payload = TaskInstanceStart(
+            job_id=job_id,
+            node_id=node_id,
+            task=task,
+            operator=operator,
+            execution_date=execution_date.isoformat(),
+            rendered_config=rendered_config,
+            start_time=start_time.isoformat(),
+            hostname=hostname,
+            pid=pid,
+        ).model_dump()
+        params = {
+            "env_id": get_env_id(),
+        }
+        return self.request(
+            "POST",
+            path="/api/executor/start",
+            response_model_class=TaskInstanceStartResponse,
+            params=params,
+            json=payload,
+        )
+
+    def task_instance_end(
+        self,
+        meta: Any,
+        traceback: Any,
+        status: ETLExecutionStatus,
+        end_time: datetime.datetime,
+        execution_date: datetime.datetime,
+        job_id: int,
+        node_id: int,
+        run_id: str,
+        current_retry_number: Optional[int],
+        max_retry_number: Optional[int],
+        link_workflow_id: Optional[int],
+        link_node_id: Optional[int],
+        data_interval_end: Optional[str],
+        **kwargs,
+    ):
+        payload = TaskInstanceEnd(
+            job_id=job_id,
+            node_id=node_id,
+            run_id=run_id,
+            end_time=end_time.isoformat(),
+            execution_date=execution_date.isoformat(),
+            meta=meta,
+            traceback=traceback,
+            status=status,
+            current_retry_number=current_retry_number,
+            max_retry_number=max_retry_number,
+            link_workflow_id=link_workflow_id,
+            link_node_id=link_node_id,
+            data_interval_end=data_interval_end,
+        ).model_dump()
+        params = {
+            "env_id": get_env_id(),
+        }
+        return self.request(
+            "POST",
+            path="/api/executor/end",
+            params=params,
+            json=payload,
+            timeout=10,  # todo: backend process time is slow, dispatch still wait
+        )
+
+    def get_workflow_node(self, workflow_id: int, node_id: int) -> WorkflowNodeItem:
+        params = {
+            "env_id": get_env_id(),
+            "workflow_id": workflow_id,
+        }
+        return self.request(
+            "GET", path=f"/api/executor/workflow_node/{node_id}", response_model_class=WorkflowNodeItem, params=params
+        )
+
+    def send_back_logs(self, record: TaskLogRecord):
+        params = {
+            "env_id": get_env_id(),
+        }
+        return self.request("POST", path="/api/executor/logs", params=params, json={"records": [record.model_dump()]})
+
+    def send_back_debug_logs(self, record: DebugLogRecord):
+        params = {
+            "env_id": get_env_id(),
+        }
+        return self.request(
+            "POST", path="/api/executor/debug_logs", params=params, json={"records": [record.model_dump()]}
+        )
+
+    def get_debug_node(self, workflow_id: int, node_key: str) -> WorkflowNodeItem:
+        params = {
+            "env_id": get_env_id(),
+            "node_key": node_key,
+        }
+        return self.request(
+            "GET", path=f"/api/executor/debug_node/{workflow_id}", response_model_class=WorkflowNodeItem, params=params
+        )
+
+    def debug_start(self, workflow_id: int, node_key: str, celery_task_id: str):
+        payload = DebugStart(workflow_id=workflow_id, node_key=node_key, celery_task_id=celery_task_id).model_dump()
+        params = {
+            "env_id": get_env_id(),
+        }
+        return self.request(
+            "POST",
+            path="/api/executor/debug_start",
+            params=params,
+            json=payload,
+        )
+
+    def debug_end(self, workflow_id: int, node_key: str, celery_task_id: str, is_success: bool):
+        payload = DebugEnd(
+            workflow_id=workflow_id, node_key=node_key, celery_task_id=celery_task_id, is_success=is_success
+        ).model_dump()
+        params = {
+            "env_id": get_env_id(),
+        }
+        return self.request(
+            "POST",
+            path="/api/executor/debug_end",
+            params=params,
+            json=payload,
+            timeout=10,
+        )
+
+    def get_py_conn_configs(
+        self,
+        conn_type: str = "python",
+        pyenv_name: str = RECURVE_EXECUTOR_PYENV_NAME,
+        project_conn_name: str = "",
+        project_id: int = 0,
+    ) -> dict:
+        params = {
+            "conn_type": conn_type,
+            "pyenv_name": pyenv_name,
+            "project_conn_name": project_conn_name,
+            "project_id": project_id,
+        }
+        return self.request("GET", path="/api/executor/python-conn-configs", params=params)