ingestr 0.13.2__py3-none-any.whl → 0.14.104__py3-none-any.whl

ingestr/src/masking.py

@@ -0,0 +1,344 @@
+import hashlib
+import hmac
+import random
+import re
+import string
+import uuid
+from datetime import date, datetime, timedelta
+from typing import Any, Callable, Dict, Optional, Tuple, Union
+
+
+class MaskingEngine:
+    def __init__(self):
+        self.token_cache: Dict[str, Union[str, int]] = {}
+        self.sequential_counter = 0
+
+    def parse_mask_config(self, config: str) -> Tuple[str, str, Optional[str]]:
+        parts = config.split(":")
+        if len(parts) == 2:
+            return parts[0], parts[1], None
+        elif len(parts) == 3:
+            return parts[0], parts[1], parts[2]
+        else:
+            raise ValueError(
+                f"Invalid mask configuration: {config}. Expected format: 'column:algorithm[:param]'"
+            )
+
+    def get_masking_function(
+        self, algorithm: str, param: Optional[str] = None
+    ) -> Callable:
+        algorithm = algorithm.lower()
+
+        # Hash-based masking
+        if algorithm == "hash" or algorithm == "sha256":
+            return self._hash_sha256
+        elif algorithm == "md5":
+            return self._hash_md5
+        elif algorithm == "hmac":
+            return lambda x: self._hash_hmac(x, param or "default-key")
+
+        # Format-preserving masking
+        elif algorithm == "email":
+            return self._mask_email
+        elif algorithm == "phone":
+            return self._mask_phone
+        elif algorithm == "credit_card":
+            return self._mask_credit_card
+        elif algorithm == "ssn":
+            return self._mask_ssn
+
+        # Redaction strategies
+        elif algorithm == "redact":
+            return lambda x: "REDACTED"
+        elif algorithm == "stars":
+            return lambda x: "*" * len(str(x)) if x else ""
+        elif algorithm == "fixed":
+            return lambda x: param or "MASKED"
+        elif algorithm == "random":
+            return self._random_replace
+
+        # Partial masking
+        elif algorithm == "partial":
+            chars = int(param) if param else 2
+            return lambda x: self._partial_mask(x, chars)
+        elif algorithm == "first_letter":
+            return self._first_letter_mask
+
+        # Tokenization
+        elif algorithm == "uuid":
+            return self._tokenize_uuid
+        elif algorithm == "sequential":
+            return self._tokenize_sequential
+
+        # Numeric masking
+        elif algorithm == "round":
+            precision = int(param) if param else 10
+            return lambda x: self._round_number(x, precision)
+        elif algorithm == "range":
+            bucket_size = int(param) if param else 100
+            return lambda x: self._range_mask(x, bucket_size)
+        elif algorithm == "noise":
+            noise_level = float(param) if param else 0.1
+            return lambda x: self._add_noise(x, noise_level)
+
+        # Date masking
+        elif algorithm == "date_shift":
+            max_days = int(param) if param else 30
+            return lambda x: self._date_shift(x, max_days)
+        elif algorithm == "year_only":
+            return self._year_only
+        elif algorithm == "month_year":
+            return self._month_year
+
+        else:
+            raise ValueError(f"Unknown masking algorithm: {algorithm}")
+
+    # Hash functions
+    def _hash_sha256(self, value: Any) -> Optional[str]:
+        if value is None:
+            return None
+        return hashlib.sha256(str(value).encode()).hexdigest()
+
+    def _hash_md5(self, value: Any) -> Optional[str]:
+        if value is None:
+            return None
+        return hashlib.md5(str(value).encode()).hexdigest()
+
+    def _hash_hmac(self, value: Any, key: str) -> Optional[str]:
+        if value is None:
+            return None
+        return hmac.new(key.encode(), str(value).encode(), hashlib.sha256).hexdigest()
+
+    # Format-preserving masks
+    def _mask_email(self, value: Any) -> Any:
+        if value is None or not value:
+            return value
+        email_str = str(value)
+        if "@" not in email_str:
+            return self._partial_mask(email_str, 2)
+
+        local, domain = email_str.split("@", 1)
+        if len(local) <= 2:
+            masked_local = "*" * len(local)
+        else:
+            masked_local = local[0] + "*" * (len(local) - 2) + local[-1]
+        return f"{masked_local}@{domain}"
+
+    def _mask_phone(self, value: Any) -> Any:
+        if value is None or not value:
+            return value
+        phone_str = re.sub(r"\D", "", str(value))
+        if len(phone_str) < 10:
+            return "*" * len(phone_str)
+
+        # Keep country code and area code, mask the rest
+        if len(phone_str) >= 10:
+            return phone_str[:3] + "-***-****"
+        return phone_str
+
+    def _mask_credit_card(self, value: Any) -> Any:
+        if value is None or not value:
+            return value
+        cc_str = re.sub(r"\D", "", str(value))
+        if len(cc_str) < 12:
+            return "*" * len(cc_str)
+        return "*" * (len(cc_str) - 4) + cc_str[-4:]
+
+    def _mask_ssn(self, value: Any) -> Any:
+        if value is None or not value:
+            return value
+        ssn_str = re.sub(r"\D", "", str(value))
+        if len(ssn_str) != 9:
+            return "*" * len(ssn_str)
+        return "***-**-" + ssn_str[-4:]
+
+    # Partial masking
+    def _partial_mask(self, value: Any, chars_to_show: int) -> Any:
+        if value is None or not value:
+            return value
+        val_str = str(value)
+        if len(val_str) <= chars_to_show * 2:
+            return "*" * len(val_str)
+        return (
+            val_str[:chars_to_show]
+            + "*" * (len(val_str) - chars_to_show * 2)
+            + val_str[-chars_to_show:]
+        )
+
+    def _first_letter_mask(self, value: Any) -> Any:
+        if value is None or not value:
+            return value
+        val_str = str(value)
+        if len(val_str) <= 1:
+            return val_str
+        return val_str[0] + "*" * (len(val_str) - 1)
+
+    # Random replacement
+    def _random_replace(self, value: Any) -> Any:
+        if value is None:
+            return value
+
+        if isinstance(value, (int, float)):
+            # Generate random number in similar range
+            if isinstance(value, int):
+                magnitude = len(str(abs(value)))
+                return random.randint(10 ** (magnitude - 1), 10**magnitude - 1)
+            else:
+                return random.uniform(0, abs(value) * 2)
+        elif isinstance(value, str):
+            # Generate random string of same length
+            return "".join(
+                random.choices(string.ascii_letters + string.digits, k=len(value))
+            )
+        else:
+            return str(value)
+
+    # Tokenization
+    def _tokenize_uuid(self, value: Any) -> Optional[str]:
+        if value is None:
+            return None
+        val_str = str(value)
+        if val_str not in self.token_cache:
+            self.token_cache[val_str] = str(uuid.uuid4())
+        return str(self.token_cache[val_str])
+
+    def _tokenize_sequential(self, value: Any) -> Optional[int]:
+        if value is None:
+            return None
+        val_str = str(value)
+        if val_str not in self.token_cache:
+            self.sequential_counter += 1
+            self.token_cache[val_str] = self.sequential_counter
+        return int(self.token_cache[val_str])
+
+    # Numeric masking
+    def _round_number(self, value: Any, precision: int) -> Any:
+        if value is None:
+            return value
+        try:
+            num = float(value)
+            return round(num / precision) * precision
+        except (ValueError, TypeError):
+            return value
+
+    def _range_mask(self, value: Any, bucket_size: int) -> Any:
+        if value is None:
+            return value
+        try:
+            num = float(value)
+            lower = int(num // bucket_size) * bucket_size
+            upper = lower + bucket_size
+            return f"{lower}-{upper}"
+        except (ValueError, TypeError):
+            return value
+
+    def _add_noise(self, value: Any, noise_level: float) -> Any:
+        if value is None:
+            return value
+        try:
+            num = float(value)
+            noise = random.uniform(-noise_level, noise_level) * abs(num)
+            result = num + noise
+            if isinstance(value, int):
+                return int(result)
+            return result
+        except (ValueError, TypeError):
+            return value
+
+    # Date masking
+    def _date_shift(self, value: Any, max_days: int) -> Any:
+        if value is None:
+            return value
+
+        if isinstance(value, (date, datetime)):
+            shift_days = random.randint(-max_days, max_days)
+            return value + timedelta(days=shift_days)
+
+        # Try to parse string dates
+        try:
+            from dateutil import parser  # type: ignore
+
+            dt = parser.parse(str(value))
+            shift_days = random.randint(-max_days, max_days)
+            result = dt + timedelta(days=shift_days)
+            if isinstance(value, str):
+                return result.strftime("%Y-%m-%d")
+            return result
+        except Exception:
+            return value
+
+    def _year_only(self, value: Any) -> Any:
+        if value is None:
+            return value
+
+        if isinstance(value, (date, datetime)):
+            return value.year
+
+        # Try to parse string dates
+        try:
+            from dateutil import parser
+
+            dt = parser.parse(str(value))
+            return dt.year
+        except Exception:
+            return value
+
+    def _month_year(self, value: Any) -> Any:
+        if value is None:
+            return value
+
+        if isinstance(value, (date, datetime)):
+            return f"{value.year}-{value.month:02d}"
+
+        # Try to parse string dates
+        try:
+            from dateutil import parser
+
+            dt = parser.parse(str(value))
+            return f"{dt.year}-{dt.month:02d}"
+        except Exception:
+            return value
+
+
+def create_masking_mapper(mask_configs: list[str]) -> Callable:
+    engine = MaskingEngine()
+
+    # Parse all configurations
+    masks = {}
+    for config in mask_configs:
+        column, algorithm, param = engine.parse_mask_config(config)
+        masks[column] = engine.get_masking_function(algorithm, param)
+
+    def apply_masks(data: Any) -> Any:
+        # Handle PyArrow tables
+        try:
+            import pyarrow as pa  # type: ignore
+
+            if isinstance(data, pa.Table):
+                # Convert to pandas for easier manipulation
+                df = data.to_pandas()
+
+                # Apply masks to each column
+                for column, mask_func in masks.items():
+                    if column in df.columns:
+                        df[column] = df[column].apply(mask_func)
+
+                # Convert back to PyArrow table
+                return pa.Table.from_pandas(df)
+        except ImportError:
+            pass
+
+        # Handle dictionaries (original behavior)
+        if isinstance(data, dict):
+            for column, mask_func in masks.items():
+                if column in data:
+                    try:
+                        data[column] = mask_func(data[column])
+                    except Exception as e:
+                        print(f"Warning: Failed to mask column {column}: {e}")
+            return data
+
+        # Return as-is if not a supported type
+        return data
+
+    return apply_masks

ingestr/src/mixpanel/__init__.py

@@ -0,0 +1,62 @@
+from typing import Iterable
+
+import dlt
+import pendulum
+from dlt.common.typing import TDataItem
+from dlt.sources import DltResource
+
+from .client import MixpanelClient
+
+
+@dlt.source(max_table_nesting=0)
+def mixpanel_source(
+    username: str,
+    password: str,
+    project_id: str,
+    server: str,
+    start_date: pendulum.DateTime,
+    end_date: pendulum.DateTime | None = None,
+) -> Iterable[DltResource]:
+    client = MixpanelClient(username, password, project_id, server)
+
+    @dlt.resource(write_disposition="merge", name="events", primary_key="distinct_id")
+    def events(
+        date=dlt.sources.incremental(
+            "time",
+            initial_value=start_date.int_timestamp,
+            end_value=end_date.int_timestamp if end_date else None,
+            range_end="closed",
+            range_start="closed",
+        ),
+    ) -> Iterable[TDataItem]:
+        if date.end_value is None:
+            end_dt = pendulum.now(tz="UTC")
+        else:
+            end_dt = pendulum.from_timestamp(date.end_value)
+
+        start_dt = pendulum.from_timestamp(date.last_value)
+
+        yield from client.fetch_events(
+            start_dt,
+            end_dt,
+        )
+
+    @dlt.resource(write_disposition="merge", primary_key="distinct_id", name="profiles")
+    def profiles(
+        last_seen=dlt.sources.incremental(
+            "last_seen",
+            initial_value=start_date,
+            end_value=end_date,
+            range_end="closed",
+            range_start="closed",
+        ),
+    ) -> Iterable[TDataItem]:
+        if last_seen.end_value is None:
+            end_dt = pendulum.now(tz="UTC")
+        else:
+            end_dt = last_seen.end_value
+
+        start_dt = last_seen.last_value
+        yield from client.fetch_profiles(start_dt, end_dt)
+
+    return events, profiles

ingestr/src/mixpanel/client.py

@@ -0,0 +1,99 @@
+import json
+from typing import Iterable
+
+import pendulum
+from dlt.sources.helpers.requests import Client
+
+
+class MixpanelClient:
+    def __init__(self, username: str, password: str, project_id: str, server: str):
+        self.username = username
+        self.password = password
+        self.project_id = project_id
+        self.server = server
+        self.session = Client(raise_for_status=False).session
+
+    def fetch_events(
+        self, start_date: pendulum.DateTime, end_date: pendulum.DateTime
+    ) -> Iterable[dict]:
+        if self.server == "us":
+            server = "data"
+        elif self.server == "in":
+            server = "data-in"
+        else:
+            server = "data-eu"
+
+        url = f"https://{server}.mixpanel.com/api/2.0/export/"
+        params = {
+            "project_id": self.project_id,
+            "from_date": start_date.format("YYYY-MM-DD"),
+            "to_date": end_date.format("YYYY-MM-DD"),
+        }
+        headers = {
+            "accept": "text/plain",
+        }
+        from requests.auth import HTTPBasicAuth
+
+        auth = HTTPBasicAuth(self.username, self.password)
+        resp = self.session.get(url, params=params, headers=headers, auth=auth)
+        resp.raise_for_status()
+        for line in resp.iter_lines():
+            if line:
+                data = json.loads(line.decode())
+                if "properties" in data:
+                    for key, value in data["properties"].items():
+                        if key.startswith("$"):
+                            data[key[1:]] = value
+                        else:
+                            data[key] = value
+                    del data["properties"]
+                yield data
+
+    def fetch_profiles(
+        self, start_date: pendulum.DateTime, end_date: pendulum.DateTime
+    ) -> Iterable[dict]:
+        if self.server == "us":
+            server = ""
+        elif self.server == "in":
+            server = "in."
+        else:
+            server = "eu."
+        url = f"https://{server}mixpanel.com/api/query/engage"
+        headers = {
+            "accept": "application/json",
+            "content-type": "application/x-www-form-urlencoded",
+        }
+        from requests.auth import HTTPBasicAuth
+
+        auth = HTTPBasicAuth(self.username, self.password)
+        page = 0
+        session_id = None
+        while True:
+            params = {"project_id": self.project_id, "page": str(page)}
+            if session_id:
+                params["session_id"] = session_id
+            start_str = start_date.format("YYYY-MM-DDTHH:mm:ss")
+            end_str = end_date.format("YYYY-MM-DDTHH:mm:ss")
+            where = f'properties["$last_seen"] >= "{start_str}" and properties["$last_seen"] <= "{end_str}"'
+            params["where"] = where
+            resp = self.session.post(url, params=params, headers=headers, auth=auth)
+
+            resp.raise_for_status()
+            data = resp.json()
+
+            for result in data.get("results", []):
+                for key, value in result["$properties"].items():
+                    if key.startswith("$"):
+                        if key == "$last_seen":
+                            result["last_seen"] = pendulum.parse(value)
+                        else:
+                            result[key[1:]] = value
+                result["distinct_id"] = result["$distinct_id"]
+                del result["$properties"]
+                del result["$distinct_id"]
+                yield result
+            if not data.get("results"):
+                break
+            session_id = data.get("session_id", session_id)
+
+            page += 1