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

recurvedata/utils/normalizer.py

@@ -0,0 +1,102 @@
+class ColumnTypeNormalizer:
+    """
+    A class to normalize database column types to a standard set of types.
+
+    Attributes:
+    - database (str): The name of the database for which normalization is needed.
+    - _normalized_types (dict[str, list[str]]): A dictionary mapping normalized types to their corresponding database-specific types.
+
+    Usage examples:
+    >>> mysql_normalizer = ColumnTypeNormalizer(database='mysql')
+    >>> mysql_normalizer.normalize('varchar')
+    "string"
+    >>> mysql_normalizer.normalize('tinyint(1)')
+    "boolean"
+    >>> postgres_normalizer = ColumnTypeNormalizer(database='postgresql')
+    >>> postgres_normalizer.normalize('int4')
+    "integer"
+    >>> postgres_normalizer.normalize('jsonb')
+    "json"
+    >>> snowflake_normalizer = ColumnTypeNormalizer(database='snowflake')
+    >>> snowflake_normalizer.normalize('timestamp_ntz')
+    "datetime"
+    >>> snowflake_normalizer.normalize('variant')
+    "json"
+    """
+
+    _COMMON_NORMALIZED_TYPES: dict[str, list[str]] = {
+        "integer": [
+            "int",
+            "integer",
+            "smallint",
+            "bigint",
+            "tinyint",
+            "int2",
+            "int4",
+            "int8",
+            "int16",
+            "int32",
+            "int64",
+        ],
+        "float": ["float", "double", "real", "decimal", "numeric"],
+        "string": ["varchar", "char", "text", "character varying", "nvarchar", "nchar", "clob", "string"],
+        "boolean": ["bool", "boolean"],
+        "date": ["date"],
+        "datetime": ["datetime", "timestamp", "timestamp with time zone", "timestamp without time zone"],
+        "time": ["time"],
+        "binary": ["binary", "blob", "varbinary"],
+        "json": ["json", "jsonb", "object", "variant", "array"],
+    }
+
+    def __init__(self, database: str, custom_mappings: dict[str, list[str]] | None = None):
+        """
+        Initializes the normalizer for a specific database with optional custom mappings.
+
+        Args:
+        - database: The name of the database for which normalization is needed.
+        - custom_mappings: A dictionary mapping database-specific types to normalized types.
+        """
+        self.database = database.lower()
+        # Convert common types and their mappings to lowercase
+        self._normalized_types = {
+            k.lower(): [v.lower() for v in values] for k, values in self._COMMON_NORMALIZED_TYPES.items()
+        }
+
+        # Convert custom mappings to lowercase, if provided
+        if custom_mappings:
+            for key, values in custom_mappings.items():
+                key = key.lower()
+                values = [v.lower() for v in values]
+                if key in self._normalized_types:
+                    self._normalized_types[key].extend(values)
+                else:
+                    self._normalized_types[key] = values
+
+    def normalize(self, column_type: str) -> str:
+        """
+        Normalizes a given database column type to a standard type.
+
+        Args:
+        - column_type: The database column type to normalize.
+
+        Returns:
+
+        Usage example:
+        >>> normalizer = ColumnTypeNormalizer('postgresql')
+        >>> normalizer.normalize('int4')
+        "integer"
+        """
+        column_type = column_type.lower()
+
+        for normalized, column_types in self._normalized_types.items():
+            for _type in column_types:
+                if _type.startswith(column_type):
+                    return normalized
+
+        return column_type
+
+
+if __name__ == "__main__":
+    import doctest
+
+    doctest.testmod()

recurvedata/utils/redis_lock.py

@@ -0,0 +1,474 @@
+import asyncio
+import threading
+import time
+from typing import Optional
+
+import redis
+from redis.asyncio import Redis as AsyncRedis
+from redis.asyncio.lock import Lock as NativeAsyncRedisLock
+from redis.client import Lock as NativeRedisLock
+from redis.client import Redis
+from redis.exceptions import LockError
+
+from recurvedata.config import REDIS_LOCK_URL
+
+
+class RedisLock:
+    """
+    Redis-based distributed lock implementation using redis.lock
+
+    Usage:
+    ```
+    with RedisLock("my_lock_name", expire=60) as lock:
+        if lock.acquired:
+            # Execute code that needs lock protection
+            pass
+        else:
+            # Failed to acquire lock
+            pass
+    ```
+
+    Or manage manually:
+    ```
+    lock = RedisLock("my_lock_name")
+    if lock.acquire():
+        try:
+            # Execute code that needs lock protection
+            pass
+        finally:
+            lock.release()
+    ```
+
+    You can also extend the lock expiration time:
+    ```
+    lock = RedisLock("my_lock_name", expire=60)
+    if lock.acquire():
+        try:
+            # Start some long operation
+            # ...
+            # Extend the lock if operation takes longer than expected
+            lock.extend(additional_time=60)
+            # Continue operation
+            # ...
+        finally:
+            lock.release()
+    ```
+
+    For long-running operations with unknown duration, you can use auto-extend:
+    ```
+    with RedisLock("my_lock_name", expire=60, auto_extend=True) as lock:
+        if lock.acquired:
+            # Execute long-running code, lock will be automatically extended
+            # until the operation completes or the lock is released
+            pass
+    ```
+    """
+
+    def __init__(
+        self,
+        name: str,
+        expire: int = 60,
+        timeout: int = 0,
+        sleep_interval: float = 0.1,
+        redis_client: Optional[Redis] = None,
+        redis_url: Optional[str] = REDIS_LOCK_URL,
+        auto_extend: bool = False,
+        extend_interval: Optional[int] = None,
+    ):
+        """
+        Initialize Redis lock
+
+        Args:
+            name: Lock name, must be unique
+            expire: Lock expiration time (seconds), prevents deadlock
+            timeout: Lock acquisition timeout (seconds), 0 means try only once
+            sleep_interval: Sleep interval (seconds)
+            redis_client: Optional Redis client, if not provided will create from REDIS_URL
+            auto_extend: Whether to automatically extend the lock while it's held
+            extend_interval: Interval (seconds) to extend the lock, defaults to expire/3
+        """
+        self.name = f"recurve:lock:{name}"
+        self.expire = expire
+        self.timeout = timeout
+        self.sleep_interval = sleep_interval
+        self.acquired = False
+        self.auto_extend = auto_extend
+        self.extend_interval = extend_interval or max(1, int(expire / 3))
+        self._extend_thread = None
+        self._stop_extend = threading.Event()
+        self._lock_token = None  # Store the lock token for cross-thread access
+
+        if redis_client is not None:
+            self.redis = redis_client
+        else:
+            self.redis = redis.from_url(redis_url)
+
+        # Create the native Redis lock
+        self.lock: NativeRedisLock = self.redis.lock(self.name, timeout=self.expire, sleep=self.sleep_interval)
+
+    def _extend_lock_periodically(self):
+        """Background thread that periodically extends the lock"""
+        while not self._stop_extend.is_set():
+            # Sleep for the extend interval
+            for _ in range(int(self.extend_interval / self.sleep_interval)):  # Check stop flag more frequently
+                if self._stop_extend.is_set():
+                    return
+                time.sleep(self.sleep_interval)
+
+            # Extend the lock if we still have it
+            if self.acquired:
+                try:
+                    # Directly extend the lock using Redis commands instead of the lock.extend method
+                    # This avoids the thread-local token issue
+                    success = self._extend_lock_directly()
+                    if not success:
+                        # If extension fails, stop the thread
+                        self._stop_extend.set()
+                except Exception:
+                    # If any exception occurs, stop the thread
+                    self._stop_extend.set()
+
+    def _extend_lock_directly(self) -> bool:
+        """
+        Extend the lock directly using Redis commands
+
+        Returns:
+            bool: Whether successfully extended the lock
+        """
+        if not self.acquired or not self._lock_token:
+            return False
+
+        try:
+            # Use Redis PEXPIRE command to extend the lock if the token matches
+            script = """
+            if redis.call('get', KEYS[1]) == ARGV[1] then
+                return redis.call('pexpire', KEYS[1], ARGV[2])
+            else
+                return 0
+            end
+            """
+            # Convert seconds to milliseconds for pexpire
+            extend_time_ms = int(self.expire * 1000)
+            result = self.redis.eval(script, 1, self.name, self._lock_token, extend_time_ms)
+            return bool(result)
+        except Exception:
+            return False
+
+    def _start_extend_thread(self):
+        """Start the background thread to extend the lock periodically"""
+        if not self._extend_thread:
+            self._stop_extend.clear()
+            self._extend_thread = threading.Thread(
+                target=self._extend_lock_periodically, daemon=True  # Make it a daemon so it doesn't block program exit
+            )
+            self._extend_thread.start()
+
+    def _stop_extend_thread(self):
+        """Stop the background thread that extends the lock"""
+        if self._extend_thread:
+            self._stop_extend.set()
+            if self._extend_thread.is_alive():
+                self._extend_thread.join(timeout=1.0)  # Wait for thread to finish
+            self._extend_thread = None
+
+    def acquire(self) -> bool:
+        """
+        Try to acquire the lock
+
+        Returns:
+            bool: Whether successfully acquired the lock
+        """
+        if self.timeout > 0:
+            # With timeout
+            try:
+                self.acquired = self.lock.acquire(blocking=True, blocking_timeout=self.timeout)
+            except redis.exceptions.LockError:
+                self.acquired = False
+        else:
+            # Without timeout (try once)
+            try:
+                self.acquired = self.lock.acquire(blocking=False)
+            except redis.exceptions.LockError:
+                self.acquired = False
+
+        # Store the lock token for cross-thread access if acquired
+        if self.acquired:
+            # Access the thread-local token from the lock
+            self._lock_token = self.lock.local.token
+
+            # If auto_extend is enabled, start the extend thread
+            if self.auto_extend:
+                self._start_extend_thread()
+
+        return self.acquired
+
+    def release(self) -> bool:
+        """
+        Release the lock
+
+        Returns:
+            bool: Whether successfully released the lock
+        """
+        if not self.acquired:
+            return False
+
+        # Stop the extend thread if it's running
+        self._stop_extend_thread()
+
+        try:
+            self.lock.release()
+            self.acquired = False
+            self._lock_token = None  # Clear the token
+            return True
+        except redis.exceptions.LockError:
+            return False
+
+    def extend(self, additional_time: Optional[int] = None) -> bool:
+        """
+        Extend the lock's expiration time
+
+        Args:
+            additional_time: Additional seconds to extend the lock.
+                             If None, uses the original expire time.
+
+        Returns:
+            bool: Whether successfully extended the lock
+        """
+        if not self.acquired:
+            return False
+
+        try:
+            # If additional_time is not provided, use the original expire time
+            extend_time = additional_time if additional_time is not None else self.expire
+
+            # Try to use the lock's extend method first
+            try:
+                success = self.lock.extend(additional_time=extend_time)
+                return success
+            except AttributeError:
+                # Fallback for older redis-py versions that don't support extend
+                # or if we're in a different thread
+                return self._extend_lock_directly()
+
+        except redis.exceptions.LockError:
+            return False
+
+    def __enter__(self) -> "RedisLock":
+        self.acquire()
+        return self
+
+    def __exit__(self, exc_type, exc_val, exc_tb) -> None:
+        if self.acquired:
+            self.release()
+
+
+class AsyncRedisLock:
+    """
+    Redis-based distributed lock implementation using redis.lock (async version)
+
+    Usage:
+    ```
+    async with AsyncRedisLock("my_lock_name", expire=60) as lock:
+        if lock.acquired:
+            # Execute code that needs lock protection
+            pass
+        else:
+            # Failed to acquire lock
+            pass
+    ```
+
+    Or manage manually:
+    ```
+    lock = AsyncRedisLock("my_lock_name")
+    if await lock.acquire():
+        try:
+            # Execute code that needs lock protection
+            pass
+        finally:
+            await lock.release()
+    ```
+
+    You can also extend the lock expiration time:
+    ```
+    lock = AsyncRedisLock("my_lock_name", expire=60)
+    if await lock.acquire():
+        try:
+            # Start some long operation
+            # ...
+            # Extend the lock if operation takes longer than expected
+            await lock.extend(additional_time=60)
+            # Continue operation
+            # ...
+        finally:
+            await lock.release()
+    ```
+
+    For long-running operations with unknown duration, you can use auto-extend:
+    ```
+    async with AsyncRedisLock("my_lock_name", expire=60, auto_extend=True) as lock:
+        if lock.acquired:
+            # Execute long-running code, lock will be automatically extended
+            # until the operation completes or the lock is released
+            pass
+    ```
+    """
+
+    def __init__(
+        self,
+        name: str,
+        expire: int = 60,
+        timeout: int = 0,
+        sleep_interval: float = 0.1,
+        redis_client: Optional[AsyncRedis] = None,
+        redis_url: Optional[str] = REDIS_LOCK_URL,
+        auto_extend: bool = False,
+        extend_interval: Optional[int] = None,
+    ):
+        """
+        Initialize Redis lock
+
+        Args:
+            name: Lock name, must be unique
+            expire: Lock expiration time (seconds), prevents deadlock
+            timeout: Lock acquisition timeout (seconds), 0 means try only once
+            sleep_interval: Sleep interval (seconds)
+            redis_client: Optional Redis client, if not provided will create from REDIS_URL
+            auto_extend: Whether to automatically extend the lock while it's held
+            extend_interval: Interval (seconds) to extend the lock, defaults to expire/3
+        """
+        self.name = f"recurve:lock:{name}"
+        self.expire = expire
+        self.timeout = timeout
+        self.sleep_interval = sleep_interval
+        self.acquired = False
+        self.auto_extend = auto_extend
+        self.extend_interval = extend_interval or max(1, int(expire / 3))
+        self._extend_task = None
+        self._stop_extend = asyncio.Event()
+        self._lock_token = None
+
+        self.redis = redis_client if redis_client is not None else AsyncRedis.from_url(redis_url)
+        self.lock: NativeAsyncRedisLock = self.redis.lock(self.name, timeout=self.expire, sleep=self.sleep_interval)
+
+    async def _extend_lock_periodically(self):
+        """Background task that periodically extends the lock"""
+        while not self._stop_extend.is_set():
+            try:
+                await asyncio.sleep(self.extend_interval)
+                if self.acquired:
+                    success = await self._extend_lock_directly()
+                    if not success:
+                        self._stop_extend.set()
+            except Exception:
+                self._stop_extend.set()
+
+    async def _extend_lock_directly(self) -> bool:
+        """
+        Extend the lock directly using Redis commands
+
+        Returns:
+            bool: Whether successfully extended the lock
+        """
+        if not self.acquired or not self._lock_token:
+            return False
+
+        try:
+            script = """
+            if redis.call('get', KEYS[1]) == ARGV[1] then
+                return redis.call('pexpire', KEYS[1], ARGV[2])
+            else
+                return 0
+            end
+            """
+            extend_time_ms = int(self.expire * 1000)
+            result = await self.redis.eval(script, 1, self.name, self._lock_token, extend_time_ms)
+            return bool(result)
+        except Exception:
+            return False
+
+    def _start_extend_task(self):
+        """Start the background task to extend the lock periodically"""
+        if not self._extend_task:
+            self._stop_extend.clear()
+            self._extend_task = asyncio.create_task(self._extend_lock_periodically())
+
+    async def _stop_extend_task(self):
+        """Stop the background task that extends the lock"""
+        if self._extend_task:
+            self._stop_extend.set()
+            await self._extend_task
+            self._extend_task = None
+
+    async def acquire(self) -> bool:
+        """
+        Try to acquire the lock
+
+        Returns:
+            bool: Whether successfully acquired the lock
+        """
+        try:
+            if self.timeout > 0:
+                self.acquired = await self.lock.acquire(blocking=True, blocking_timeout=self.timeout)
+            else:
+                self.acquired = await self.lock.acquire(blocking=False)
+
+            if self.acquired:
+                self._lock_token = self.lock.local.token
+                if self.auto_extend:
+                    self._start_extend_task()
+
+            return self.acquired
+        except LockError:
+            self.acquired = False
+            return False
+
+    async def release(self) -> bool:
+        """
+        Release the lock
+
+        Returns:
+            bool: Whether successfully released the lock
+        """
+        if not self.acquired:
+            return False
+
+        await self._stop_extend_task()
+
+        try:
+            await self.lock.release()
+            self.acquired = False
+            self._lock_token = None
+            return True
+        except LockError:
+            return False
+
+    async def extend(self, additional_time: Optional[int] = None) -> bool:
+        """
+        Extend the lock's expiration time
+
+        Args:
+            additional_time: Additional seconds to extend the lock.
+                           If None, uses the original expire time.
+
+        Returns:
+            bool: Whether successfully extended the lock
+        """
+        if not self.acquired:
+            return False
+
+        try:
+            extend_time = additional_time if additional_time is not None else self.expire
+            try:
+                return await self.lock.extend(additional_time=extend_time)
+            except AttributeError:
+                return await self._extend_lock_directly()
+        except LockError:
+            return False
+
+    async def __aenter__(self) -> "AsyncRedisLock":
+        await self.acquire()
+        return self
+
+    async def __aexit__(self, exc_type, exc_val, exc_tb) -> None:
+        if self.acquired:
+            await self.release()

recurvedata/utils/registry.py

@@ -0,0 +1,54 @@
+from typing import Callable, Generic, ItemsView, Iterable, KeysView, TypeVar, Union, ValuesView
+
+_KT = TypeVar("_KT")
+_VT = TypeVar("_VT")
+
+
+class GenericRegistry(Generic[_KT, _VT]):
+    def __init__(self):
+        self._registry: dict[_KT, _VT] = {}
+
+    def add(self, *keys: _KT) -> Callable[[_VT], _VT]:
+        def inner(target: _VT) -> _VT:
+            for k in keys:
+                self._registry[k] = target
+            return target
+
+        return inner
+
+    def get(self, key: _KT, default: _VT = None) -> _VT:
+        return self._registry.get(key, default)
+
+    def __len__(self) -> int:
+        return len(self._registry)
+
+    def keys(self) -> KeysView[_KT]:
+        return self._registry.keys()
+
+    def values(self) -> ValuesView[_VT]:
+        return self._registry.values()
+
+    def items(self) -> ItemsView[_KT, _VT]:
+        return self._registry.items()
+
+    def __getitem__(self, key):
+        return self._registry.get(key)
+
+
+class Registry(GenericRegistry[str, _VT]):
+    def __init__(self, key_callback: Callable[[_VT], Union[str, Iterable[str]]]):
+        self.key_callback = key_callback
+        super().__init__()
+
+    def add(self, target: _VT) -> _VT:
+        keys: Union[str, Iterable[str]] = self.key_callback(target)
+        if isinstance(keys, str):
+            keys = [keys]
+        return super().add(*keys)(target)
+
+    def register(self, target: _VT) -> _VT:  # Compatibility with recurvedata.operator
+        return self.add(target)
+
+
+jinja2_template_funcs_registry = Registry(key_callback=lambda x: x.__name__)
+register_func = jinja2_template_funcs_registry.add

recurvedata/utils/shell.py

@@ -0,0 +1,15 @@
+import logging
+import subprocess
+from typing import Optional
+
+_logger = logging.getLogger(__name__)
+
+
+def run(cmd: str, logger: Optional[logging.Logger] = _logger) -> None:
+    logger.debug("Running command: %s", cmd)
+    subprocess.check_call(cmd, shell=True)
+
+
+def run_output(cmd: str, logger: Optional[logging.Logger] = _logger) -> str:
+    logger.debug("Running command: %s", cmd)
+    return subprocess.check_output(cmd, shell=True).decode("utf-8").strip()

recurvedata/utils/singleton.py

@@ -0,0 +1,33 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+from typing import Generic, TypeVar
+
+T = TypeVar("T")
+
+
+class Singleton(type, Generic[T]):
+    """Metaclass that allows to implement singleton pattern."""
+
+    _instances: dict[Singleton[T], T] = {}
+
+    def __call__(cls: Singleton[T], *args, **kwargs) -> T:
+        if cls not in cls._instances:
+            cls._instances[cls] = super().__call__(*args, **kwargs)
+        return cls._instances[cls]

recurvedata/utils/sql.py

@@ -0,0 +1,6 @@
+def trim_replace_special_character(sql: str, strip_sufix: bool = False) -> str:
+    sql = sql.replace("\\n", "\n")  # todo: may cause error if \\n is in `like '%\\n%'` format
+    # process \\n in sql
+    if strip_sufix:
+        sql = sql.strip(";")
+    return sql