icdev 0.0.3__py3-none-any.whl

tools/databridge/agent/ws_relay.py

@@ -0,0 +1,91 @@
+# CUI // SP-CTI
+"""DataBridge WebSocket Relay Client (D-DB-10).
+
+Outbound-only WebSocket connection from on-prem agent to ICDEV cloud relay.
+No inbound ports required (firewall-friendly).
+"""
+
+import json
+import logging
+import threading
+import time
+from typing import Any, Callable, Dict, Optional
+
+logger = logging.getLogger("databridge.agent.ws_relay")
+
+try:
+    import websocket
+    HAS_WEBSOCKET = True
+except ImportError:
+    websocket = None  # type: ignore[assignment]
+    HAS_WEBSOCKET = False
+
+
+def connect_relay(
+    url: str,
+    agent_id: str,
+    handler: Callable[[Dict[str, Any]], Dict[str, Any]],
+    reconnect_delay: int = 30,
+    max_reconnect_attempts: int = 10,
+) -> Optional[Any]:
+    """Connect to relay server and handle incoming messages.
+
+    Returns the WebSocket app (runs in background thread).
+    """
+    if not HAS_WEBSOCKET:
+        raise ImportError("websocket-client required for relay connection")
+
+    def on_message(ws, message):
+        try:
+            request = json.loads(message)
+            logger.info("Received request: %s", request.get("action"))
+            response = handler(request)
+            response["request_id"] = request.get("request_id", "")
+            response["agent_id"] = agent_id
+            ws.send(json.dumps(response, default=str))
+        except Exception as exc:
+            logger.error("Handler error: %s", exc)
+            ws.send(json.dumps({
+                "status": "error",
+                "message": str(exc),
+                "agent_id": agent_id,
+            }))
+
+    def on_open(ws):
+        logger.info("Connected to relay: %s", url)
+        # Send registration message
+        ws.send(json.dumps({
+            "action": "agent_register",
+            "agent_id": agent_id,
+        }))
+
+    def on_close(ws, close_status_code, close_msg):
+        logger.warning("Relay disconnected (code=%s)", close_status_code)
+
+    def on_error(ws, error):
+        logger.error("WebSocket error: %s", error)
+
+    ws_app = websocket.WebSocketApp(
+        url,
+        on_message=on_message,
+        on_open=on_open,
+        on_close=on_close,
+        on_error=on_error,
+    )
+
+    def _run_with_reconnect():
+        attempts = 0
+        while attempts < max_reconnect_attempts:
+            try:
+                ws_app.run_forever(ping_interval=30, ping_timeout=10)
+            except Exception as exc:
+                logger.error("WebSocket run failed: %s", exc)
+            attempts += 1
+            if attempts < max_reconnect_attempts:
+                delay = min(reconnect_delay * (2 ** (attempts - 1)), 300)
+                logger.info("Reconnecting in %ds (attempt %d/%d)", delay, attempts, max_reconnect_attempts)
+                time.sleep(delay)
+
+    thread = threading.Thread(target=_run_with_reconnect, daemon=True, name="ws-relay")
+    thread.start()
+    return ws_app

tools/databridge/analytics.py

@@ -0,0 +1,167 @@
+# CUI // SP-CTI
+"""DataBridge DuckDB Analytics Engine (D-DB-5).
+
+Provides SQL query execution over connected sources using DuckDB.
+Supports cross-source JOINs, column profiling, and in-memory analytics.
+"""
+
+import json
+import logging
+import time
+from typing import Any, Dict, List, Optional
+
+logger = logging.getLogger("databridge.analytics")
+
+try:
+    import duckdb
+    HAS_DUCKDB = True
+except ImportError:
+    duckdb = None  # type: ignore[assignment]
+    HAS_DUCKDB = False
+
+try:
+    import pyarrow as pa
+    HAS_ARROW = True
+except ImportError:
+    pa = None  # type: ignore[assignment]
+    HAS_ARROW = False
+
+
+class AnalyticsEngine:
+    """DuckDB-based analytics engine for DataBridge.
+
+    Executes SQL queries over Arrow tables loaded from connected sources.
+    Zero infrastructure — in-process, air-gap safe.
+    """
+
+    def __init__(self, memory_limit_mb: int = 256, threads: int = 4):
+        if not HAS_DUCKDB:
+            raise ImportError("duckdb required for analytics engine")
+        self._conn = duckdb.connect(":memory:")
+        self._conn.execute(f"SET memory_limit='{memory_limit_mb}MB'")
+        self._conn.execute(f"SET threads={threads}")
+        self._registered_tables: Dict[str, Any] = {}
+
+    def register_table(self, name: str, data: Any) -> Dict[str, Any]:
+        """Register an Arrow Table or list of dicts as a queryable table."""
+        try:
+            if HAS_ARROW and isinstance(data, pa.Table):
+                self._conn.register(name, data)
+            elif HAS_ARROW and isinstance(data, pa.RecordBatch):
+                table = pa.Table.from_batches([data])
+                self._conn.register(name, table)
+            elif isinstance(data, list) and data:
+                if HAS_ARROW:
+                    table = pa.Table.from_pylist(data)
+                    self._conn.register(name, table)
+                else:
+                    # Create from CSV-like inline
+                    return {"status": "error", "message": "Arrow required for list registration"}
+            else:
+                return {"status": "error", "message": "Unsupported data type"}
+
+            self._registered_tables[name] = {"rows": len(data) if hasattr(data, "__len__") else 0}
+            return {"status": "ok", "table": name}
+        except Exception as exc:
+            return {"status": "error", "message": str(exc)}
+
+    def query(self, sql: str, max_rows: int = 100_000) -> Dict[str, Any]:
+        """Execute a SQL query and return results."""
+        start = time.time()
+        try:
+            result = self._conn.execute(sql)
+            columns = [desc[0] for desc in result.description]
+
+            if HAS_ARROW:
+                arrow_table = result.fetch_arrow_table()
+                if arrow_table.num_rows > max_rows:
+                    arrow_table = arrow_table.slice(0, max_rows)
+                data = arrow_table
+                row_count = arrow_table.num_rows
+            else:
+                rows = result.fetchall()
+                if len(rows) > max_rows:
+                    rows = rows[:max_rows]
+                data = [dict(zip(columns, row)) for row in rows]
+                row_count = len(rows)
+
+            duration_ms = int((time.time() - start) * 1000)
+            return {
+                "status": "ok",
+                "data": data,
+                "columns": columns,
+                "row_count": row_count,
+                "duration_ms": duration_ms,
+            }
+        except Exception as exc:
+            return {
+                "status": "error",
+                "message": str(exc),
+                "duration_ms": int((time.time() - start) * 1000),
+            }
+
+    def profile_column(self, table_name: str, column_name: str) -> Dict[str, Any]:
+        """Profile a single column: type, nulls, distinct, min, max, avg, etc."""
+        try:
+            # DuckDB identifiers from config, not user input — nosec B608
+            profile_sql = (  # nosec B608
+                f'SELECT typeof("{column_name}") as data_type, '  # nosec B608
+                f'COUNT(*) as total_rows, COUNT("{column_name}") as non_null, '  # nosec B608
+                f'COUNT(*) - COUNT("{column_name}") as null_count, '  # nosec B608
+                f'COUNT(DISTINCT "{column_name}") as distinct_count, '  # nosec B608
+                f'MIN("{column_name}") as min_value, MAX("{column_name}") as max_value '  # nosec B608
+                f'FROM "{table_name}"'  # nosec B608
+            )
+            stats = self._conn.execute(profile_sql).fetchone()
+
+            result = {
+                "column": column_name,
+                "data_type": stats[0],
+                "total_rows": stats[1],
+                "non_null": stats[2],
+                "null_count": stats[3],
+                "null_pct": round(stats[3] / stats[1] * 100, 1) if stats[1] > 0 else 0,
+                "distinct_count": stats[4],
+                "uniqueness_pct": round(stats[4] / stats[1] * 100, 1) if stats[1] > 0 else 0,
+                "min_value": str(stats[5]),
+                "max_value": str(stats[6]),
+            }
+
+            # Try numeric stats
+            try:
+                num_sql = (  # nosec B608
+                    f'SELECT AVG("{column_name}"), STDDEV("{column_name}"), '  # nosec B608
+                    f'MEDIAN("{column_name}") FROM "{table_name}"'  # nosec B608
+                )
+                num_stats = self._conn.execute(num_sql).fetchone()
+                result["avg"] = float(num_stats[0]) if num_stats[0] is not None else None
+                result["stddev"] = float(num_stats[1]) if num_stats[1] is not None else None
+                result["median"] = float(num_stats[2]) if num_stats[2] is not None else None
+            except Exception:
+                pass
+
+            return {"status": "ok", **result}
+        except Exception as exc:
+            return {"status": "error", "message": str(exc)}
+
+    def profile_table(self, table_name: str) -> Dict[str, Any]:
+        """Profile all columns in a table."""
+        try:
+            col_sql = f"SELECT column_name FROM information_schema.columns WHERE table_name = '{table_name}'"  # nosec B608
+            columns = self._conn.execute(col_sql).fetchall()
+            profiles = []
+            for (col,) in columns:
+                p = self.profile_column(table_name, col)
+                if p.get("status") == "ok":
+                    profiles.append(p)
+            return {"status": "ok", "table": table_name, "columns": profiles}
+        except Exception as exc:
+            return {"status": "error", "message": str(exc)}
+
+    def list_tables(self) -> List[str]:
+        return list(self._registered_tables.keys())
+
+    def close(self) -> None:
+        if self._conn:
+            self._conn.close()
+            self._conn = None

tools/databridge/arrow_pipeline.py

@@ -0,0 +1,327 @@
+# CUI // SP-CTI
+"""DataBridge Arrow Pipeline — read → transform → write (D-DB-3, D-DB-15).
+
+All data flows through Apache Arrow RecordBatches. Transform steps are
+deterministic Arrow compute operations (air-gap safe, reproducible).
+"""
+
+import json
+import logging
+import time
+from dataclasses import dataclass, field
+from typing import Any, Callable, Dict, List, Optional
+
+logger = logging.getLogger("databridge.arrow_pipeline")
+
+try:
+    import pyarrow as pa
+    import pyarrow.compute as pc
+
+    HAS_ARROW = True
+except ImportError:
+    pa = None  # type: ignore[assignment]
+    pc = None  # type: ignore[assignment]
+    HAS_ARROW = False
+
+
+# ---------------------------------------------------------------------------
+# Transform step
+# ---------------------------------------------------------------------------
+
+@dataclass
+class TransformStep:
+    """A single deterministic transform in the pipeline."""
+
+    name: str
+    operation: str  # rename, cast, filter, select, add_column, drop_column, ...
+    params: Dict[str, Any] = field(default_factory=dict)
+
+
+# ---------------------------------------------------------------------------
+# Pipeline
+# ---------------------------------------------------------------------------
+
+class ArrowPipeline:
+    """Executes a DAG of Arrow transforms on RecordBatches (D-DB-15).
+
+    Usage:
+        pipeline = ArrowPipeline()
+        pipeline.add_step(TransformStep("select_cols", "select", {"columns": ["a", "b"]}))
+        pipeline.add_step(TransformStep("cast_a", "cast", {"column": "a", "target_type": "int64"}))
+        result = pipeline.execute(record_batch)
+    """
+
+    def __init__(self):
+        self.steps: List[TransformStep] = []
+        self._operations: Dict[str, Callable] = {
+            "select": self._op_select,
+            "drop": self._op_drop,
+            "rename": self._op_rename,
+            "cast": self._op_cast,
+            "filter": self._op_filter,
+            "add_column": self._op_add_column,
+            "sort": self._op_sort,
+            "limit": self._op_limit,
+            "deduplicate": self._op_deduplicate,
+            "fill_null": self._op_fill_null,
+            "upper": self._op_upper,
+            "lower": self._op_lower,
+            "trim": self._op_trim,
+        }
+
+    def add_step(self, step: TransformStep) -> "ArrowPipeline":
+        self.steps.append(step)
+        return self
+
+    def clear(self) -> "ArrowPipeline":
+        self.steps.clear()
+        return self
+
+    def execute(self, data: Any) -> Dict[str, Any]:
+        """Execute all steps on an Arrow Table/RecordBatch.
+
+        Returns: {status, data, steps_executed, duration_ms, errors}
+        """
+        if not HAS_ARROW:
+            return {"status": "error", "errors": ["pyarrow not installed"]}
+
+        # Normalize to Table for easier manipulation
+        if isinstance(data, pa.RecordBatch):
+            table = pa.Table.from_batches([data])
+        elif isinstance(data, pa.Table):
+            table = data
+        else:
+            return {"status": "error", "errors": [f"Unsupported data type: {type(data)}"]}
+
+        start = time.time()
+        errors: List[str] = []
+        steps_executed = 0
+
+        for step in self.steps:
+            op_fn = self._operations.get(step.operation)
+            if not op_fn:
+                errors.append(f"Unknown operation: {step.operation}")
+                continue
+            try:
+                table = op_fn(table, step.params)
+                steps_executed += 1
+            except Exception as exc:
+                errors.append(f"Step '{step.name}' ({step.operation}): {exc}")
+                logger.error("Pipeline step '%s' failed: %s", step.name, exc)
+
+        duration_ms = int((time.time() - start) * 1000)
+
+        return {
+            "status": "ok" if not errors else "partial",
+            "data": table,
+            "row_count": table.num_rows,
+            "column_count": table.num_columns,
+            "steps_executed": steps_executed,
+            "steps_total": len(self.steps),
+            "duration_ms": duration_ms,
+            "errors": errors,
+        }
+
+    # -------------------------------------------------------------------
+    # Operations
+    # -------------------------------------------------------------------
+
+    def _op_select(self, table: Any, params: Dict) -> Any:
+        """Select specific columns."""
+        columns = params.get("columns", [])
+        return table.select(columns)
+
+    def _op_drop(self, table: Any, params: Dict) -> Any:
+        """Drop specific columns."""
+        columns = params.get("columns", [])
+        return table.drop(columns)
+
+    def _op_rename(self, table: Any, params: Dict) -> Any:
+        """Rename columns: {old_name: new_name}."""
+        mapping = params.get("mapping", {})
+        names = list(table.column_names)
+        new_names = [mapping.get(n, n) for n in names]
+        return table.rename_columns(new_names)
+
+    def _op_cast(self, table: Any, params: Dict) -> Any:
+        """Cast column to target type."""
+        column = params["column"]
+        target = params["target_type"]
+        type_map = {
+            "int8": pa.int8(), "int16": pa.int16(), "int32": pa.int32(),
+            "int64": pa.int64(), "float32": pa.float32(), "float64": pa.float64(),
+            "utf8": pa.utf8(), "bool": pa.bool_(), "date32": pa.date32(),
+            "timestamp": pa.timestamp("us"),
+        }
+        arrow_type = type_map.get(target)
+        if not arrow_type:
+            raise ValueError(f"Unknown target type: {target}")
+        idx = table.column_names.index(column)
+        casted = pc.cast(table.column(column), arrow_type)
+        return table.set_column(idx, column, casted)
+
+    def _op_filter(self, table: Any, params: Dict) -> Any:
+        """Filter rows where column matches condition."""
+        column = params["column"]
+        op = params.get("op", "==")
+        value = params["value"]
+        col = table.column(column)
+        ops = {
+            "==": pc.equal, "!=": pc.not_equal,
+            ">": pc.greater, ">=": pc.greater_equal,
+            "<": pc.less, "<=": pc.less_equal,
+        }
+        fn = ops.get(op)
+        if not fn:
+            raise ValueError(f"Unknown filter op: {op}")
+        mask = fn(col, value)
+        return table.filter(mask)
+
+    def _op_add_column(self, table: Any, params: Dict) -> Any:
+        """Add a column with a constant value."""
+        name = params["name"]
+        value = params["value"]
+        dtype = params.get("type", "utf8")
+        arr = pa.array([value] * table.num_rows, type=pa.utf8() if dtype == "utf8" else pa.int64())
+        return table.append_column(name, arr)
+
+    def _op_sort(self, table: Any, params: Dict) -> Any:
+        """Sort by column(s)."""
+        by = params.get("by", [])
+        order = params.get("order", "ascending")
+        sort_keys = [(col, order) for col in by]
+        indices = pc.sort_indices(table, sort_keys=sort_keys)
+        return table.take(indices)
+
+    def _op_limit(self, table: Any, params: Dict) -> Any:
+        """Limit to N rows."""
+        n = params.get("n", 100)
+        return table.slice(0, n)
+
+    def _op_deduplicate(self, table: Any, params: Dict) -> Any:
+        """Remove duplicate rows based on columns."""
+        columns = params.get("columns", table.column_names)
+        # Group by columns, take first row per group
+        seen = set()
+        indices = []
+        for i in range(table.num_rows):
+            key = tuple(table.column(c)[i].as_py() for c in columns)
+            if key not in seen:
+                seen.add(key)
+                indices.append(i)
+        return table.take(indices)
+
+    def _op_fill_null(self, table: Any, params: Dict) -> Any:
+        """Replace nulls in column with value."""
+        column = params["column"]
+        value = params["value"]
+        idx = table.column_names.index(column)
+        filled = pc.if_else(pc.is_null(table.column(column)), value, table.column(column))
+        return table.set_column(idx, column, filled)
+
+    def _op_upper(self, table: Any, params: Dict) -> Any:
+        """Uppercase a string column."""
+        column = params["column"]
+        idx = table.column_names.index(column)
+        return table.set_column(idx, column, pc.utf8_upper(table.column(column)))
+
+    def _op_lower(self, table: Any, params: Dict) -> Any:
+        """Lowercase a string column."""
+        column = params["column"]
+        idx = table.column_names.index(column)
+        return table.set_column(idx, column, pc.utf8_lower(table.column(column)))
+
+    def _op_trim(self, table: Any, params: Dict) -> Any:
+        """Trim whitespace from string column."""
+        column = params["column"]
+        idx = table.column_names.index(column)
+        return table.set_column(idx, column, pc.utf8_trim(table.column(column), " "))
+
+
+# ---------------------------------------------------------------------------
+# Convenience functions
+# ---------------------------------------------------------------------------
+
+def convert_format(
+    data: Any,
+    output_format: str,
+    output_path: Optional[str] = None,
+) -> Dict[str, Any]:
+    """Convert Arrow data to another format.
+
+    Supported: parquet, csv, json, ipc
+    """
+    if not HAS_ARROW:
+        return {"status": "error", "message": "pyarrow not installed"}
+
+    if isinstance(data, pa.RecordBatch):
+        table = pa.Table.from_batches([data])
+    elif isinstance(data, pa.Table):
+        table = data
+    else:
+        return {"status": "error", "message": f"Unsupported data type: {type(data)}"}
+
+    start = time.time()
+
+    if output_format == "parquet":
+        import pyarrow.parquet as pq
+        if output_path:
+            pq.write_table(table, output_path)
+        else:
+            buf = pa.BufferOutputStream()
+            pq.write_table(table, buf)
+            return {
+                "status": "ok",
+                "format": "parquet",
+                "bytes": buf.getvalue().size,
+                "rows": table.num_rows,
+                "duration_ms": int((time.time() - start) * 1000),
+            }
+    elif output_format == "csv":
+        import pyarrow.csv as pcsv
+        if output_path:
+            pcsv.write_csv(table, output_path)
+        else:
+            buf = pa.BufferOutputStream()
+            pcsv.write_csv(table, buf)
+            return {
+                "status": "ok",
+                "format": "csv",
+                "bytes": buf.getvalue().size,
+                "rows": table.num_rows,
+                "duration_ms": int((time.time() - start) * 1000),
+            }
+    elif output_format == "json":
+        rows = table.to_pydict()
+        json_str = json.dumps(rows, default=str)
+        if output_path:
+            with open(output_path, "w") as f:
+                f.write(json_str)
+        return {
+            "status": "ok",
+            "format": "json",
+            "bytes": len(json_str),
+            "rows": table.num_rows,
+            "duration_ms": int((time.time() - start) * 1000),
+        }
+    elif output_format == "ipc":
+        if output_path:
+            writer = pa.ipc.new_file(output_path, table.schema)
+            writer.write_table(table)
+            writer.close()
+        return {
+            "status": "ok",
+            "format": "ipc",
+            "rows": table.num_rows,
+            "duration_ms": int((time.time() - start) * 1000),
+        }
+    else:
+        return {"status": "error", "message": f"Unsupported format: {output_format}"}
+
+    return {
+        "status": "ok",
+        "format": output_format,
+        "output_path": output_path,
+        "rows": table.num_rows,
+        "duration_ms": int((time.time() - start) * 1000),
+    }