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

recurvedata/core/templating.py

@@ -0,0 +1,206 @@
+import copy
+import datetime
+import inspect
+import re
+import types
+from typing import Any, Callable, TypeVar
+
+import jinja2.nodes
+from jinja2 import Environment, TemplateSyntaxError, meta, pass_context
+from jinja2.runtime import Context
+
+from recurvedata.utils.crontab import previous_schedule
+from recurvedata.utils.registry import jinja2_template_funcs_registry
+
+T = TypeVar("T")
+
+# {% set navigation = [('index.html', 'Index'), ('about.html', 'About')] %} -> navigation
+# {% set key, value = call_something() %} -> key, value
+_jinja2_set_p = re.compile(r"\{%\s*set\s([\w,\s]+?)\s*=.*")
+
+
+def get_template_env() -> Environment:
+    env = Environment(
+        cache_size=0,
+        trim_blocks=True,
+        lstrip_blocks=True,
+    )
+    for func_name, func in jinja2_template_funcs_registry.items():
+        env.globals[func_name] = func
+    return env
+
+
+def extract_vars_from_template_code(template_code: str) -> list[str]:
+    """
+    This function is copied from recurve-server recurve.library.jinja_utils
+    """
+    env = Environment(autoescape=True)
+    ast = env.parse(template_code)
+
+    extracted_var_names = []
+
+    # Helper function to recursively walk through nodes
+    def visit_node(node: jinja2.nodes.Node):
+        # We're looking for Call nodes where the function is 'var'
+        if isinstance(node, jinja2.nodes.Call) and isinstance(node.node, jinja2.nodes.Name) and node.node.name == "var":
+            # The first argument of the Call node is the variable name
+            arguments = [arg.value for arg in node.args if isinstance(arg, jinja2.nodes.Const)]
+            extracted_var_names.append(arguments[0])
+
+        # Recursively visit child nodes
+        for child_node in node.iter_child_nodes():
+            visit_node(child_node)
+
+    # Start the traversal
+    visit_node(ast)
+
+    return extracted_var_names
+
+
+@pass_context
+def var_function(context: Context, name: str, default: Any = None) -> Any:
+    return context.get(name, default)
+
+
+class Renderer(object):
+    def __init__(self):
+        self.env = get_template_env()
+        self.env.globals["var"] = var_function
+
+    @staticmethod
+    def init_context(execution_date: datetime.datetime, schedule_interval: str):
+        yesterday_dttm = execution_date - datetime.timedelta(days=1)
+        tomorrow_dttm = execution_date + datetime.timedelta(days=1)
+        data_interval_start = previous_schedule(schedule_interval, execution_date)
+
+        template_context = {
+            "dt": execution_date.date(),
+            "yesterday": yesterday_dttm,
+            "yesterday_dt": yesterday_dttm.date(),
+            "tomorrow": tomorrow_dttm,
+            "tomorrow_dt": tomorrow_dttm.date(),
+            # "execution_date": execution_date,
+            "logical_date": execution_date,
+            "data_interval_start": data_interval_start,
+            "data_interval_end": execution_date,
+            "data_interval_start_dt": data_interval_start and data_interval_start.date(),
+            "data_interval_end_dt": execution_date.date(),
+        }
+        return template_context
+
+    @staticmethod
+    def get_functions() -> dict[str, Callable]:
+        return dict(jinja2_template_funcs_registry.items())
+
+    def render_template(self, tmpl: T, context: dict) -> T:
+        if isinstance(tmpl, str):
+            result = self.env.from_string(tmpl).render(**context)
+        elif isinstance(tmpl, (tuple, list)):
+            result = [self.render_template(x, context) for x in tmpl]
+        elif isinstance(tmpl, dict):
+            result = {k: self.render_template(v, context) for k, v in tmpl.items()}
+        else:
+            # raise TypeError(f'Type {type(tmpl)} is not supported for templating')
+            result = tmpl
+        return result
+
+    def extract_variables(self, tmpl: str) -> list[str]:
+        ast = self.env.parse(tmpl)
+        var_variables: list[str] = extract_vars_from_template_code(tmpl)
+        variables: set[str] = meta.find_undeclared_variables(ast)
+        variables.update(var_variables)
+
+        # exclude assignments within if blocks.
+        # e.g. for this template, the undefined variables are ['yesterday_ds']
+        # WITHOUT 'dedup_order', which is a local variable defined by set
+        #  {% if yesterday_ds <= '2020-09-25' %}
+        #    {% set dedup_order = "snapshot_time ASC" %}
+        #  {% else %}
+        #    {% set dedup_order = "sell_count DESC NULLS LAST" %}
+        #  {% endif %}
+        assignments: list[str] = _jinja2_set_p.findall(tmpl)
+        for vs in set(assignments):
+            for v in vs.split(","):
+                v = v.strip()
+                if v in variables:
+                    variables.remove(v)
+
+        return sorted(variables)
+
+    def _prepare_jinja_context(
+        self, exist_variables: dict, execution_date: datetime.datetime, schedule_interval: str
+    ) -> dict:
+        template_var_dct = self.init_context(execution_date, schedule_interval)
+        context = copy.copy(template_var_dct)  # shallow copy
+        context.update(exist_variables)  # python code may use exist_variables
+        for func_name, func in self.get_functions().items():
+            context.setdefault(func_name, func)
+
+        return context
+
+    def render_variables(
+        self, variables: dict[str, Any], execution_date: datetime.datetime, schedule_interval: str
+    ) -> dict:
+        """
+        Renders variables in a dictionary using Jinja2 templating.
+
+        This method processes each string value in the input dictionary,
+        rendering it with Jinja2 if it contains template variables. The
+        rendered values are then updated in the original dictionary.
+
+        Args:
+            variables: A dictionary of variables to render.
+            execution_date: a given date that some date function can be used to calculate
+            schedule_interval: A string representing the schedule interval (crontab expression).
+        Returns:
+            A dictionary with the same keys as the input, but with rendered values.
+        """
+        context = self._prepare_jinja_context(variables, execution_date, schedule_interval)
+        update_dct = {}
+        for name, val in variables.items():
+            if not isinstance(val, str):
+                # Process only the Jinja variables within the string.
+                continue
+            try:
+                jinja_variables = self.extract_variables(val)
+            except TemplateSyntaxError:
+                # invalid jinja, leave it unrendered
+                continue
+            if not jinja_variables:
+                continue
+
+            rendered_val = self.render_template(val, context)
+            if rendered_val != val:
+                update_dct[name] = rendered_val
+                context[name] = rendered_val
+        variables.update(update_dct)
+        return variables
+
+    def extract_python_code_variable(
+        self, python_code: str, exist_variables: dict, execution_date: datetime.datetime, schedule_interval: str
+    ) -> dict:
+        result = {}
+        name_space = self._prepare_jinja_context(exist_variables, execution_date, schedule_interval)
+        orig_name_space = copy.copy(name_space)
+        rendered_code = self.render_template(python_code, name_space)
+
+        compiled_code = compile(rendered_code, "", "exec")
+        exec(compiled_code, name_space)
+
+        for key2, value2 in name_space.items():
+            if key2 == "__builtins__":
+                continue
+            if isinstance(value2, types.ModuleType):
+                continue
+            if inspect.isclass(value2):
+                continue
+            if key2 in orig_name_space and orig_name_space[key2] == value2:  # only return defined/updated var
+                continue
+            result[key2] = value2
+        return result
+
+
+if __name__ == "__main__":
+    import doctest
+
+    doctest.testmod()

recurvedata/core/tracing.py

@@ -0,0 +1,223 @@
+import asyncio
+import functools
+import inspect
+import json
+import random
+import re
+from typing import Any, Dict, overload
+
+from opentelemetry import context, trace
+from opentelemetry.context import Context
+from opentelemetry.exporter.otlp.proto.http.trace_exporter import OTLPSpanExporter
+from opentelemetry.sdk.resources import SERVICE_NAME, Resource
+from opentelemetry.sdk.trace import Span, TracerProvider
+from opentelemetry.sdk.trace.export import BatchSpanProcessor, SpanExporter
+from opentelemetry.trace.propagation.tracecontext import TraceContextTextMapPropagator
+from pydantic import BaseModel
+
+from .consts import TRACE_KEY, TRACING_CONTEXT_KEY
+
+
+class _OTLPSpanExporter(OTLPSpanExporter):
+    def update_auth_headers(self, headers: Dict[str, str]):
+        if not headers:
+            return
+        self._session.headers.update(headers)
+
+
+class FilteringSpanProcessor(BatchSpanProcessor):
+    """A span processor that filters out unwanted spans before they are created."""
+
+    def __init__(
+        self,
+        span_exporter: SpanExporter,
+        max_queue_size: int = None,
+        schedule_delay_millis: float = None,
+        max_export_batch_size: int = None,
+        export_timeout_millis: float = None,
+    ):
+        super().__init__(
+            span_exporter,
+            schedule_delay_millis=schedule_delay_millis,
+            export_timeout_millis=export_timeout_millis,
+            max_export_batch_size=max_export_batch_size,
+            max_queue_size=max_queue_size,
+        )
+        # Single compiled pattern with OR operator for better performance
+        self._fastapi_pattern = re.compile(r"^(?:handling\s+event|Event\s+.*\s+dispatched$)")
+
+    def _should_filter_span(self, span: Span) -> bool:
+        """Determine if a span should be filtered out using regex pattern."""
+        if not span.name:
+            return False
+        return bool(self._fastapi_pattern.match(span.name))
+
+    def on_end(self, span: Span) -> None:
+        """Called when a span is ended."""
+        if not self._should_filter_span(span):
+            super().on_end(span)
+
+
+class Singleton(type):
+    _instances = {}
+
+    def __call__(cls, *args, **kwargs):
+        if cls not in cls._instances:
+            instance = super().__call__(*args, **kwargs)
+            cls._instances[cls] = instance
+        return cls._instances[cls]
+
+
+class Tracing(metaclass=Singleton):
+    @classmethod
+    def is_instantiated(cls):
+        """
+        Returns the valid singleton instance if it exists, otherwise None.
+        """
+        instance: "Tracing" = cls._instances.get(cls, None)
+        if not instance:
+            return False
+        if not instance.tracer:
+            return False
+        return True
+
+    def __init__(
+        self,
+        endpoint: str = None,
+        service: str = None,
+        export_timeout_millis: int = 1000,
+        schedule_delay_millis: int = 1000,
+    ):
+        if endpoint and service:
+            self.init(endpoint, service, export_timeout_millis, schedule_delay_millis)
+        else:
+            self.exporter = None
+            self.tracer = None
+
+    def init(
+        self,
+        endpoint: str,
+        service: str,
+        export_timeout_millis: int = 1000,
+        schedule_delay_millis: int = 1000,
+    ):
+        self.exporter = _OTLPSpanExporter(endpoint=endpoint)
+        processor = FilteringSpanProcessor(
+            self.exporter, schedule_delay_millis=schedule_delay_millis, export_timeout_millis=export_timeout_millis
+        )
+        resource = Resource.create(attributes={SERVICE_NAME: service})
+        provider = TracerProvider(resource=resource)
+        provider.add_span_processor(processor)
+        trace.set_tracer_provider(provider)
+        self.tracer = trace.get_tracer(service, tracer_provider=provider)
+
+    def set_auth_headers(self, headers: Dict[str, str]):
+        self.exporter.update_auth_headers(headers)
+
+    def dump_context(self) -> Dict[str, str]:
+        carrier = {}
+        TraceContextTextMapPropagator().inject(carrier)
+        return carrier
+
+    def load_context(self, ctx_json: Dict[str, str] | None = None) -> Context:
+        ctx = TraceContextTextMapPropagator().extract(ctx_json or {})
+        context.attach(ctx)
+
+    def send_context(self, payload: dict):
+        payload[TRACE_KEY] = self.dump_context()
+        return payload
+
+    @overload
+    def receive_context(self, payload: dict):
+        ...
+
+    @overload
+    def receive_context(self, payload: str):
+        ...
+
+    @overload
+    def receive_context(self, payload: Any):
+        ...
+
+    def receive_context(self, payload: str | dict | BaseModel | Any):
+        if isinstance(payload, str):
+            try:
+                payload = json.loads(payload)
+            except Exception:
+                return
+        elif isinstance(payload, BaseModel):
+            payload = payload.model_dump(by_alias=True)
+            if not payload.get(TRACING_CONTEXT_KEY, None):
+                return
+            payload = payload.get(TRACING_CONTEXT_KEY)
+            try:
+                payload = json.loads(payload)
+            except Exception:
+                return
+        elif isinstance(payload, dict):
+            pass
+        else:
+            from fastapi import Request
+
+            if isinstance(payload, Request):
+                payload = payload.headers
+            else:
+                return
+
+        if TRACE_KEY in payload:
+            self.load_context(payload.get(TRACE_KEY))
+
+    def create_span(self, name: str = None, sampling_rate: float = 0.0, context_payload_name: str = None):
+        """
+        Decorator to create a span for a function.
+
+        Args:
+            name: The name of the span. If not provided, the function name will be used.
+            sampling_rate: The rate at which to sample the span. Default is 0.0.
+            context_payload_name: The name of the argument that contains the context payload."""
+
+        def decorator(func):
+            def get_all_args_as_dict(func, args, kwargs):
+                """Helper to get all arguments as a dictionary."""
+                sig = inspect.signature(func)
+                bound_args = sig.bind(*args, **kwargs)
+                bound_args.apply_defaults()
+                return dict(bound_args.arguments)
+
+            @functools.wraps(func)
+            async def async_wrapper(*args, **kwargs):
+                all_args = get_all_args_as_dict(func, args, kwargs)
+                if not self.tracer:
+                    return await func(*args, **kwargs)
+                if context_payload_name:
+                    if all_args.get(context_payload_name, None):
+                        self.receive_context(all_args[context_payload_name])
+                span_name = name or f"{func.__module__}.{func.__name__}"
+                if not context.get_current():
+                    if random.random() > sampling_rate:
+                        return await func(*args, **kwargs)
+                with self.tracer.start_as_current_span(span_name):
+                    return await func(*args, **kwargs)
+
+            @functools.wraps(func)
+            def wrapper(*args, **kwargs):
+                all_args = get_all_args_as_dict(func, args, kwargs)
+                if not self.tracer:
+                    return func(*args, **kwargs)
+                if context_payload_name:
+                    if all_args.get(context_payload_name, None):
+                        self.receive_context(all_args[context_payload_name])
+                span_name = name or f"{func.__module__}.{func.__name__}"
+                if not context.get_current():
+                    if random.random() > sampling_rate:
+                        return func(*args, **kwargs)
+                with self.tracer.start_as_current_span(span_name):
+                    return func(*args, **kwargs)
+
+            return async_wrapper if asyncio.iscoroutinefunction(func) else wrapper
+
+        return decorator
+
+    @property
+    def current_span(self) -> Span:
+        return trace.get_current_span()

recurvedata/core/transformer.py

@@ -0,0 +1,186 @@
+import json
+import struct
+import zlib
+from typing import Literal
+
+from recurvedata.pigeon.schema import Schema
+from recurvedata.utils.crypto_util import CryptoUtil
+
+
+class Transformer:
+    @property
+    def input_schema(self):
+        """Returns the schema of input data"""
+        return getattr(self, "_input_schema", None)
+
+    @input_schema.setter
+    def input_schema(self, schema):
+        """Should be called by the handler"""
+        assert isinstance(schema, Schema)
+        setattr(self, "_input_schema", schema)
+
+    @property
+    def output_schema(self):
+        """Subclasses that change the rows schema should provide the output schema.
+
+        These operations will change the output schema:
+        - Add or remove fields
+        - Change the name of fields
+        - Change the type of fields
+
+        An example of valid schema:
+
+        from recurvedata.pigeon.schema import Schema, Field, types
+
+        Schema([
+            Field(name='id', type=types.INT32),
+            Field(name='name', type=types.STRING, size=64),
+            Field(name='snapshot_time', type=types.DATETIME),
+            Field(name='is_active', type=types.BOOLEAN)
+        ])
+
+        Allowed types:
+
+        - INT8 = 'INT8'  # 1-byte (8-bit) signed integers
+        - INT16 = 'INT16'  # 2-byte (16-bit) signed integers
+        - INT32 = 'INT32'  # 4-byte (32-bit) signed integers
+        - INT64 = 'INT64'  # 8-byte (64-bit) signed integers
+        - FLOAT32 = 'FLOAT32'  # 4-byte (32-bit) single-precision floating
+        - FLOAT64 = 'FLOAT64'  # 8-byte (64-bit) double-precision floating
+        - BOOLEAN = 'BOOLEAN'
+        - DATETIME = 'DATETIME'
+        - DATE = 'DATE'
+        - STRING = 'STRING'
+        """
+        return None
+
+    def transform(self, row: dict, *args, **kwargs) -> dict | list[dict] | None:
+        """This is the method called by Handler.
+
+        It internally calls `transform_impl` to do the real transform logic.
+        Subclasses should implement `transform_impl` but not this method.
+
+        :param row: a Row (namedtuple) object contains a row record fetched from database
+        :type row: collection.namedtuple
+        :returns: returns one (tuple) or multiple (list of tuple) rows
+        """
+        return self.transform_impl(row, *args, **kwargs)
+
+    def transform_impl(self, row: dict, *args, **kwargs) -> dict | list[dict] | None:
+        """subclass should override this method to implement the custom transform operations"""
+        return row
+
+    @staticmethod
+    def convert_json_to_hive_map(data: str | dict) -> str:
+        from recurvedata.connectors.connectors.hive import (  # lazy import
+            HIVE_MAP_ITEM_DELIMITER,
+            HIVE_MAP_KV_DELIMITER,
+            HIVE_NULL,
+        )
+
+        if not data:
+            return HIVE_NULL
+
+        if isinstance(data, str):
+            d = json.loads(data)
+        else:
+            d = data
+
+        items = []
+        for key, value in d.items():
+            key = str(key).strip()
+            value = str(value).strip()
+            item = f"{key}{HIVE_MAP_KV_DELIMITER}{value}"
+            items.append(item)
+        return HIVE_MAP_ITEM_DELIMITER.join(items)
+
+    @staticmethod
+    def convert_json_to_hive_array(data: str | list) -> str:
+        from recurvedata.connectors.connectors.hive import HIVE_ARRAY_DELIMITER, HIVE_NULL
+
+        if not data:
+            return HIVE_NULL
+
+        if isinstance(data, str):
+            items = json.loads(data)
+        else:
+            items = data
+
+        return HIVE_ARRAY_DELIMITER.join(items)
+
+    @staticmethod
+    def mysql_uncompress(value: bytes, return_str=False) -> str | bytes:
+        """An Python implementation of UNCOMPRESS function of MySQL.
+
+        Used to decompress result of COMPRESS function.
+
+        https://dev.mysql.com/doc/refman/5.7/en/encryption-functions.html#function_compress
+
+        :param value: the compressed data in bytes
+        :type value: bytes
+        :param return_str: the return value should be unicode
+        :type return_str: bool
+        :rtype: bytes | str
+        """
+
+        # Empty strings are stored as empty strings.
+        # Nonempty strings are stored as a 4-byte length of the uncompressed string
+        if not value or len(value) < 4:
+            return value
+
+        rv = zlib.decompress(value[4:])
+
+        if return_str:
+            rv = rv.decode()
+        return rv
+
+    @staticmethod
+    def mysql_compress(value: str) -> bytes | None:
+        if value is None:
+            return None
+        if value == "":
+            return b""
+        size = struct.pack("I", len(value))
+        data = zlib.compress(value.encode())
+        return size + data
+
+    @staticmethod
+    def json_loads(*args, **kwargs):
+        return json.loads(*args, **kwargs)
+
+    @staticmethod
+    def json_dumps(*args, **kwargs) -> str:
+        return json.dumps(*args, **kwargs)
+
+    def aes_encrypt(
+        self, key_name: str, data: str | bytes, mode: Literal["ECB", "CBC"] = "ECB", iv: str | bytes = None
+    ) -> str:
+        return CryptoUtil.base64_encode(CryptoUtil.aes_encrypt(key_name, data, mode, iv))
+
+    def aes_decrypt(self, key_name: str, data: bytes | str) -> str:
+        if isinstance(data, str):
+            data = CryptoUtil.base64_decode(data)
+        return CryptoUtil.aes_decrypt(key_name, data)
+
+    def rsa_encrypt(self, key_name: str, data: str | bytes) -> str:
+        return CryptoUtil.base64_encode(CryptoUtil.rsa_encrypt(key_name, data))
+
+    def rsa_decrypt(self, key_name: str, data: bytes | str) -> str:
+        if isinstance(data, str):
+            data = CryptoUtil.base64_decode(data)
+        return CryptoUtil.rsa_decrypt(key_name, data)
+
+    def base64_encode(self, data: str | bytes) -> str:
+        return CryptoUtil.base64_encode(data)
+
+    def base64_decode(self, data: str | bytes) -> str:
+        return CryptoUtil.base64_decode(data)
+
+    def md5(self, data: str | bytes) -> str:
+        return CryptoUtil.md5(data)
+
+    def sha1(self, data: str | bytes) -> str:
+        return CryptoUtil.sha1(data)
+
+    def sha256(self, data: str | bytes) -> str:
+        return CryptoUtil.sha256(data)

recurvedata/core/translation.py

@@ -0,0 +1,91 @@
+from contextvars import ContextVar
+from gettext import NullTranslations
+from typing import Any, Union
+
+_current_translations: ContextVar[NullTranslations] = ContextVar("_current_translations", default=None)
+
+
+class Translator:
+    _translations_cache: dict[str, NullTranslations] = {}
+    _translations_dir: str = "locale"
+
+    @classmethod
+    def set_translations_dir(cls, translations_dir: str):
+        cls._translations_dir = translations_dir
+
+    @classmethod
+    def _get_translations(cls, locale: str) -> NullTranslations:
+        """Get translations for the given locale."""
+        from babel.support import Translations
+
+        if locale not in cls._translations_cache:
+            try:
+                translations = Translations.load(dirname=cls._translations_dir, locales=[locale])
+                cls._translations_cache[locale] = translations
+            except FileNotFoundError:
+                return None
+        return cls._translations_cache[locale]
+
+    @classmethod
+    def set_locale(cls, locale: str):
+        """Set locale for the current context."""
+        translations = cls._get_translations(locale)
+        if translations:
+            translations.install()
+            _current_translations.set(translations)
+
+    @classmethod
+    def gettext(cls, message: str) -> str:
+        """Translate a message based on the current locale."""
+        translations = _current_translations.get()
+        if translations:
+            return translations.gettext(message)
+        return message
+
+
+class LazyString:
+    def __init__(self, message: str):
+        self.message = message
+
+    def __str__(self) -> str:
+        return Translator.gettext(self.message)
+
+    def __hash__(self):
+        return hash(str(self))
+
+    def __eq__(self, other: Union["LazyString", str]):
+        if isinstance(other, LazyString):
+            return str(self) == str(other)
+        return str(self) == other
+
+    def format(self, *args, **kwargs):
+        return str(self).format(*args, **kwargs)
+
+
+def convert_lazy_string(v: Any) -> Any:
+    """Recursively convert LazyString instances to str in nested data structures.
+
+    This function traverses dictionaries, lists and other data structures recursively,
+    converting any LazyString instances to regular strings. This is needed because
+    JsonSchema and Pydantic (see Pydantic issue #8439) do not support LazyString.
+
+    Args:
+        v: The value to convert. Can be a LazyString, dict, list, or any other type.
+
+    Returns:
+        The input value with all LazyString instances converted to str. The structure
+        of the input (dict/list/etc) is preserved.
+    """
+    if isinstance(v, dict):
+        return {k: convert_lazy_string(v) for k, v in v.items()}
+    if isinstance(v, list):
+        return [convert_lazy_string(v) for v in v]
+    if isinstance(v, LazyString):
+        return str(v)
+    return v
+
+
+gettext = Translator.gettext
+lazy_gettext = LazyString
+_ = Translator.gettext
+_l = LazyString