python-jack-knife 0.7.7__py3-none-any.whl → 0.8.0__py3-none-any.whl

pjk/components.py

@@ -2,7 +2,7 @@
 # Copyright 2024 Mike Schultz
 
 from abc import ABC, abstractmethod
-from typing import Any, Optional, List
+from typing import Any, Iterable, Optional, List
 from pjk.usage import Usage, NoBindUsage, ParsedToken
     
 # mixin
@@ -19,9 +19,16 @@ class KeyedSource(ABC):
         """Return the record associated with the given key, or None."""
         pass
 
-    def get_unlookedup_records(self) -> List[Any]:
-        # for outer join
-        pass
+    def lookup_item(self, left_rec):
+        """Return ``(key, record)`` for consumers that need match identity."""
+        record = self.lookup(left_rec)
+        if record is None:
+            return None
+        return id(record), record
+
+    def keyed_items(self) -> Iterable[tuple[Any, dict]]:
+        """Iterate ``(key, record)`` pairs without changing lookup state."""
+        return ((id(record), record) for record in self)
 
     def deep_copy(self):
         return None

pjk/errors.py

@@ -0,0 +1,79 @@
+# SPDX-License-Identifier: Apache-2.0
+# Copyright 2024 Mike Schultz
+
+from __future__ import annotations
+
+from typing import List, Optional
+
+import yaml
+
+CLI_PROBLEM_HEADER = "You've got a problem here."
+
+
+def _quote_token(token: str) -> str:
+    if token.startswith('[') or token.startswith('{'):
+        return '"' + token + '"'
+    return token
+
+
+def format_cli_command(tokens: List[str], token_no: int = 0, marker: str = '^') -> List[str]:
+    """Return command line + caret underline for CLI errors."""
+    if not tokens:
+        return []
+    token_copies = [_quote_token(t) for t in tokens]
+    lines = ['pjk ' + ' '.join(token_copies)]
+    offset = 4 + sum(len(t) + 1 for t in token_copies[:token_no])
+    lines.append(' ' * offset + marker * len(token_copies[token_no]))
+    return lines
+
+
+def find_token_index(tokens: List[str], token: str) -> int:
+    """Locate the failing pipe token in argv tokens."""
+    for i, t in enumerate(tokens):
+        if t == token:
+            return i
+    for i, t in enumerate(tokens):
+        if token in t or t in token:
+            return i
+    return 0
+
+
+class RecordProcessingError(Exception):
+    """Per-record pipeline failure with CLI-friendly formatting."""
+
+    def __init__(
+        self,
+        *,
+        token: str,
+        record: dict,
+        cause: BaseException,
+        message: Optional[str] = None,
+        tokens: Optional[List[str]] = None,
+        token_no: int = 0,
+    ):
+        super().__init__(message or f"Error in {token}")
+        self.token = token
+        self.record = record
+        self.cause = cause
+        self.message = message or f"Error in {token}"
+        self.tokens = tokens
+        self.token_no = token_no
+
+    def __str__(self) -> str:
+        lines = []
+        if self.tokens:
+            lines.extend(format_cli_command(self.tokens, self.token_no))
+        lines.append(CLI_PROBLEM_HEADER)
+        lines.append('')
+        lines.append(str(self.cause))
+        lines.append('')
+        lines.append('On this record:')
+        lines.append(
+            yaml.dump(
+                self.record,
+                default_flow_style=False,
+                sort_keys=False,
+                allow_unicode=True,
+            ).rstrip()
+        )
+        return '\n'.join(lines)

pjk/integrations/rest_pipe.py

@@ -0,0 +1,330 @@
+# SPDX-License-Identifier: Apache-2.0
+# Copyright 2024 Mike Schultz
+
+"""
+Generic REST integration pipe (1 record in → 1 record out).
+
+Configured via configs.yaml as RestPipe-<instance>. Path placeholders
+``{field}`` are filled from the input record; JSON body is built from the
+record (or a field subset). Response JSON object is yielded as the output
+record.
+"""
+
+from __future__ import annotations
+
+import os
+import re
+import time
+from typing import Any, Dict, List, Optional, Tuple
+from urllib.parse import urljoin
+
+from pjk.common import Integration
+from pjk.components import DeepCopyPipe
+from pjk.errors import RecordProcessingError
+from pjk.progress import papi
+from pjk.usage import ParsedToken, Usage, configs
+
+PATH_PARAM_RE = re.compile(r"\{([A-Za-z_][A-Za-z0-9_]*)\}")
+
+BODY_METHODS = frozenset({"POST", "PUT", "PATCH"})
+
+# Refresh a bit before expires_in to avoid edge-of-expiry 401s.
+_TOKEN_EXPIRY_SKEW_SEC = 30
+
+
+def _format_path(path: str, record: dict) -> str:
+    missing = []
+
+    def repl(match: re.Match) -> str:
+        name = match.group(1)
+        if name not in record or record[name] is None:
+            missing.append(name)
+            return match.group(0)
+        return str(record[name])
+
+    out = PATH_PARAM_RE.sub(repl, path)
+    if missing:
+        raise ValueError(f"missing path field(s): {', '.join(missing)}")
+    return out
+
+
+def _build_url(base_url: str, path: str) -> str:
+    base = base_url if base_url.endswith("/") else base_url + "/"
+    rel = path.lstrip("/")
+    return urljoin(base, rel)
+
+
+class RestPipe(DeepCopyPipe, Integration):
+    @classmethod
+    def usage(cls):
+        u = Usage(
+            name="rest",
+            desc=(
+                "Generic REST integration: one HTTP request per input record.\n"
+                "Yields one output record (JSON response body).\n"
+                "Configure RestPipe-<instance> in configs.yaml with base_url, "
+                "method, and path.\n"
+                "Path placeholders {field} are filled from the input record."
+            ),
+            component_class=cls,
+        )
+        u.def_arg(
+            name="instance",
+            usage=(
+                "configs.yaml entry RestPipe-<instance> must define "
+                "base_url, method, path (and auth as needed)"
+            ),
+        )
+        u.def_config_tuples(
+            [
+                ("base_url", str, None),
+                ("method", str, None),
+                ("path", str, None),
+                ("auth", str, "none", True),  # none|bearer|basic|client_credentials
+                ("token", str, None, True),
+                ("token_env", str, None, True),
+                ("username", str, None, True),
+                ("password", str, None, True),
+                ("auth_url", str, None, True),
+                ("auth_client_id", str, None, True),
+                ("auth_client_secret", str, None, True),
+                ("timeout", float, 30.0, True),
+                ("body", str, "record", True),  # record | none
+            ]
+        )
+        u.def_example(
+            expr_tokens=[
+                "{id: '1', amount: 10}",
+                "rest:demo",
+            ],
+            expect=None,
+        )
+        return u
+
+    def __init__(self, ptok: ParsedToken, usage: Usage, root=None):
+        super().__init__(ptok, usage, root)
+        self.token = ptok.whole_token
+        self.usage_obj = usage
+
+        self.base_url = usage.get_config("base_url")
+        self.method = (usage.get_config("method") or "").upper()
+        self.path = usage.get_config("path")
+        self.auth = (usage.get_config("auth") or "none").lower()
+        self.token_value = usage.get_config("token")
+        self.token_env = usage.get_config("token_env")
+        self.username = usage.get_config("username")
+        self.password = usage.get_config("password")
+        self.auth_url = usage.get_config("auth_url")
+        self.auth_client_id = usage.get_config("auth_client_id")
+        self.auth_client_secret = usage.get_config("auth_client_secret")
+        self.timeout = usage.get_config("timeout") or 30.0
+        self.body_mode = (usage.get_config("body") or "record").lower()
+
+        entry = self._resolved_entry()
+        raw_fields = entry.get("body_fields") if entry else None
+        if isinstance(raw_fields, list):
+            self.body_fields: Optional[List[str]] = [str(x) for x in raw_fields]
+        elif isinstance(raw_fields, str) and raw_fields.strip():
+            self.body_fields = [p.strip() for p in raw_fields.split(",") if p.strip()]
+        else:
+            self.body_fields = None
+
+        raw_headers = entry.get("headers") if entry else None
+        self.extra_headers: Dict[str, str] = (
+            {str(k): str(v) for k, v in raw_headers.items()}
+            if isinstance(raw_headers, dict)
+            else {}
+        )
+
+        if self.method not in {"GET", "POST", "PUT", "PATCH", "DELETE", "HEAD"}:
+            raise ValueError(f"rest: unsupported method {self.method!r}")
+
+        # Cached OAuth2 access token for auth=client_credentials
+        self._access_token: Optional[str] = None
+        self._access_token_expires_at: float = 0.0
+
+        self.recs_in = papi.get_counter(self, "recs_in", display=False)
+        self.recs_out = papi.get_percentage_counter(self, "recs_out", self.recs_in)
+
+    def _resolved_entry(self) -> dict:
+        configs._ensure_loaded()
+        instance = self.usage_obj.get_arg("instance")
+        key = f"{type(self).__name__}-{instance}"
+        return configs._resolve_entry(key) or {}
+
+    def reset(self):
+        pass
+
+    def close(self):
+        pass
+
+    def _fetch_client_credentials_token(self) -> Tuple[str, float]:
+        """
+        OAuth2 client_credentials against auth_url (OIDC/Keycloak-style).
+        Returns (access_token, expires_at_epoch).
+        """
+        import requests
+
+        if not self.auth_url:
+            raise ValueError("rest auth=client_credentials requires config 'auth_url'")
+        if not self.auth_client_id:
+            raise ValueError("rest auth=client_credentials requires config 'auth_client_id'")
+        if not self.auth_client_secret:
+            raise ValueError(
+                "rest auth=client_credentials requires config 'auth_client_secret'"
+            )
+
+        response = requests.post(
+            self.auth_url,
+            data={
+                "grant_type": "client_credentials",
+                "client_id": self.auth_client_id,
+                "client_secret": self.auth_client_secret,
+            },
+            headers={"Content-Type": "application/x-www-form-urlencoded"},
+            timeout=self.timeout,
+        )
+        if response.status_code >= 400:
+            snippet = (response.text or "")[:500]
+            raise RuntimeError(
+                f"HTTP {response.status_code} fetching token from {self.auth_url}: {snippet}"
+            )
+
+        payload = response.json()
+        access_token = payload.get("access_token")
+        if not access_token:
+            raise RuntimeError(
+                f"token response from {self.auth_url} missing access_token"
+            )
+        expires_in = float(payload.get("expires_in") or 300)
+        expires_at = time.time() + max(expires_in - _TOKEN_EXPIRY_SKEW_SEC, 1)
+        return access_token, expires_at
+
+    def _bearer_token(self) -> str:
+        if self.auth == "bearer":
+            token = self.token_value
+            if not token and self.token_env:
+                token = os.environ.get(self.token_env)
+            if not token:
+                raise ValueError(
+                    "rest auth=bearer requires config 'token' or 'token_env' "
+                    f"(env {self.token_env!r} empty/missing)"
+                )
+            return token
+
+        if self.auth == "client_credentials":
+            now = time.time()
+            if self._access_token and now < self._access_token_expires_at:
+                return self._access_token
+            token, expires_at = self._fetch_client_credentials_token()
+            self._access_token = token
+            self._access_token_expires_at = expires_at
+            return token
+
+        raise ValueError(f"rest: no bearer token for auth={self.auth!r}")
+
+    def _auth_headers_and_auth(self):
+        import requests
+
+        headers = dict(self.extra_headers)
+        auth = None
+
+        if self.auth in ("none", "", None):
+            return headers, auth
+
+        if self.auth in ("bearer", "client_credentials"):
+            headers["Authorization"] = f"Bearer {self._bearer_token()}"
+            return headers, auth
+
+        if self.auth == "basic":
+            if not self.username:
+                raise ValueError("rest auth=basic requires config 'username'")
+            auth = requests.auth.HTTPBasicAuth(self.username, self.password or "")
+            return headers, auth
+
+        raise ValueError(
+            f"rest: unsupported auth {self.auth!r} "
+            "(use none|bearer|basic|client_credentials)"
+        )
+
+    def _build_body(self, record: dict) -> Optional[dict]:
+        if self.method not in BODY_METHODS:
+            return None
+        if self.body_mode == "none":
+            return None
+        if self.body_fields is not None:
+            return {k: record[k] for k in self.body_fields if k in record}
+        # body: record
+        return dict(record)
+
+    def _response_record(self, response) -> dict:
+        content_type = response.headers.get("Content-Type", "") or ""
+        text = response.text
+        try:
+            obj = response.json()
+        except Exception:
+            return {
+                "status_code": response.status_code,
+                "content_type": content_type.split(";")[0].strip(),
+                "body": text,
+            }
+
+        if isinstance(obj, dict):
+            return obj
+        if isinstance(obj, list):
+            return {"items": obj}
+        return {"value": obj}
+
+    def __iter__(self):
+        import requests
+
+        for record in self.left:
+            self.recs_in.increment()
+            try:
+                path = _format_path(self.path, record)
+                url = _build_url(self.base_url, path)
+                headers, auth = self._auth_headers_and_auth()
+                body = self._build_body(record)
+
+                kwargs: Dict[str, Any] = {
+                    "method": self.method,
+                    "url": url,
+                    "headers": headers,
+                    "auth": auth,
+                    "timeout": self.timeout,
+                    "allow_redirects": True,
+                }
+                if body is not None:
+                    kwargs["json"] = body
+
+                response = requests.request(**kwargs)
+                # One refresh+retry on 401 when using client_credentials.
+                if (
+                    response.status_code == 401
+                    and self.auth == "client_credentials"
+                ):
+                    self._access_token = None
+                    self._access_token_expires_at = 0.0
+                    headers, auth = self._auth_headers_and_auth()
+                    kwargs["headers"] = headers
+                    kwargs["auth"] = auth
+                    response = requests.request(**kwargs)
+
+                if response.status_code >= 400:
+                    snippet = (response.text or "")[:500]
+                    raise RuntimeError(
+                        f"HTTP {response.status_code} for {self.method} {url}: {snippet}"
+                    )
+
+                out = self._response_record(response)
+            except RecordProcessingError:
+                raise
+            except Exception as e:
+                raise RecordProcessingError(
+                    token=self.token,
+                    record=record if isinstance(record, dict) else {"value": record},
+                    cause=e,
+                ) from e
+
+            self.recs_out.increment()
+            yield out

pjk/main.py

@@ -8,6 +8,7 @@ import shlex
 from typing import List, Dict
 from pjk.parser import ExpressionParser
 from pjk.usage import UsageError
+from pjk.errors import RecordProcessingError, find_token_index
 from pjk.log import init as init_logging
 import traceback
 import concurrent.futures
@@ -37,6 +38,14 @@ def execute_threaded(sinks, stop_progress=None):
         executor.shutdown(wait=False, cancel_futures=True)
         os._exit(130)
 
+    except RecordProcessingError:
+        if stop_progress:
+            try: stop_progress()
+            except Exception: pass
+        for f in futures:
+            f.cancel()
+        executor.shutdown(wait=False, cancel_futures=True)
+        raise
     except Exception as e:
         if stop_progress:
             try: stop_progress()
@@ -116,6 +125,7 @@ def execute_tokens(tokens: List[str]):
     parser = ExpressionParser(registry)
 
     display = None
+    display_cleared = False
     try:
         sink = parser.parse(tokens)
         write_history(sys.argv[1:]) # now that it's parsed sucessfully
@@ -133,17 +143,28 @@ def execute_tokens(tokens: List[str]):
 
         if len(sinks) > 1:
             # pass a stopper so we halt the UI before tracebacks / shutdown
-            execute_threaded(sinks, stop_progress=(display.stop if display else None))
+            execute_threaded(
+                sinks,
+                stop_progress=(lambda: display.stop(clear=True) if display else None),
+            )
         else:
             sink.drain()
 
     except UsageError as e:
         print(e, file=sys.stderr)
         sys.exit(2)
+    except RecordProcessingError as e:
+        if display:
+            display.stop(clear=True, timeout=0.5)
+            display_cleared = True
+        e.tokens = tokens
+        e.token_no = find_token_index(tokens, e.token)
+        print(e, file=sys.stderr)
+        sys.exit(2)
     except KeyboardInterrupt:
         pass
     finally:
-        if display:
+        if display and not display_cleared:
             # short join so Ctrl-C is immediate
             try: display.stop(timeout=0.1)
             except Exception: pass

pjk/man_page.py

@@ -104,8 +104,9 @@ def display_configs():
             print('   ', highlight(name, 'bold', name))
             print('=========================================')
 
-            if 'password' in body_dict:
-                body_dict['password'] = '*************'
+            for secret_key in ('password', 'token', 'auth_client_secret'):
+                if secret_key in body_dict:
+                    body_dict[secret_key] = '*************'
 
             try:
                 yaml.dump(