python-jack-knife 0.7.0__py3-none-any.whl → 0.7.4__py3-none-any.whl

pjk/history.py

@@ -30,6 +30,9 @@ def read_history(log_path: str) -> List[int]:
                 if not line:
                     continue
 
+                if 'pjk ' in line: # legacy
+                    line = line.split('pjk ', 1)[1]
+
                 # Expected format: <command_string>
                 line = line.strip()
                         

pjk/integrations/postgres_pipe.py

@@ -6,6 +6,7 @@
 import base64
 import datetime as _dt
 import uuid
+import time
 from decimal import Decimal
 from typing import Any, Dict, Optional
 
@@ -13,33 +14,55 @@ from pjk.usage import ParsedToken, Usage
 from pjk.common import Integration
 from pjk.pipes.query_pipe import QueryPipe
 
+MAX_RETRIES = 3
+BASE_DELAY = 0.1  # seconds
 
 class DBClient:
-    """Simple shared-connection wrapper for pg8000."""
-    _connection = None
-    
-    def __init__(self, host: str, username: str, password: Optional[str],
-                 db_name: str, port: int = 5432, ssl: bool = False):
-        import pg8000 # lazy import
-        if DBClient._connection is None:
-            try:
-                kwargs = dict(user=username, password=password, host=host, database=db_name, port=port)
-                if ssl:
-                    import ssl as _ssl
-                    kwargs["ssl_context"] = _ssl.create_default_context()
-                DBClient._connection = pg8000.connect(**kwargs)
-                DBClient._connection.autocommit = True
-            except Exception as e:
-                print("Failed to connect to DB")
-                raise e
-        self.conn = DBClient._connection
+    """Per-instance pg8000 connection wrapper. No shared state."""
+
+    def __init__(
+        self,
+        host: str,
+        username: str,
+        password: Optional[str],
+        db_name: str,
+        port: int = 5432,
+        ssl: bool = False,
+    ):
+        import pg8000  # lazy import
+
+        kwargs = dict(
+            user=username,
+            password=password,
+            host=host,
+            database=db_name,
+            port=port,
+        )
+        if ssl:
+            import ssl as _ssl
+
+            kwargs["ssl_context"] = _ssl.create_default_context()
+
+        try:
+            self.conn = pg8000.connect(**kwargs)
+            self.conn.autocommit = True
+        except Exception as e:
+            print("Failed to connect to DB")
+            raise e
 
     def close(self):
-        if self.conn is not None:
-            try:
-                self.conn.close()
-            finally:
-                DBClient._connection = None
+        if getattr(self, "conn", None) is None:
+            return
+
+        import pg8000  # lazy
+
+        try:
+            self.conn.close()
+        except pg8000.exceptions.InterfaceError:
+            # Already closed / broken; ignore.
+            pass
+        finally:
+            self.conn = None
 
 
 def _iso_dt(x: _dt.datetime) -> str:
@@ -90,43 +113,67 @@ def _row_to_dict(cursor, row) -> Dict[str, Any]:
     return {col: normalize(val) for col, val in zip(cols, row)}
 
 
-class PostgresPipe(QueryPipe,Integration):
-    name = 'postgres'
+class PostgresPipe(QueryPipe, Integration):
+    name = "postgres"
     desc = "Postgres query pipe; executes SQL over input record['query']."
-    arg0 = ("instance", 'instance of database.')
+    arg0 = ("instance", "instance of database.")
     examples = [
-        ['myquery.sql', 'postgres:mydb', '-'],
-        ["{'query': 'SELECT * from MY_TABLE;'}", 'postgres:mydb', '-'],
-        ["{'query': 'SELECT * FROM pg_catalog.pg_tables;'}", 'postgres:mydb'],
-        ["{'query': 'SELECT stored_procedure(%s, ...), batch_params:{...}"]
+        ["myquery.sql", "postgres:mydb", "-"],
+        ["{'query': 'SELECT * from MY_TABLE;'}", "postgres:mydb", "-"],
+        ["{'query': 'SELECT * FROM pg_catalog.pg_tables;'}", "postgres:mydb"],
+        ["{'query': 'SELECT procedure_batch(%s, ...), batch_params:{...}"],
+        ["{'query': 'SELECT procedure_jsonb(%s, ...), json_params:json_string"],
     ]
 
     # name, type, default
     config_tuples = [
-            ('db_name', str, None),
-            ('host', str, None),
-            ('user', str, None),
-            ('password', str, None),
-            ('port', int, 5432),
-            ('ssl', bool, False)
+        ("db_name", str, None),
+        ("host", str, None),
+        ("user", str, None),
+        ("password", str, None),
+        ("port", int, 5432),
+        ("ssl", bool, False),
     ]
-    
-    def __init__(self, ptok: ParsedToken, u: Usage):
-        super().__init__(ptok, u)
 
-        self.db_name = u.get_config('db_name')
+    def __init__(self, ptok: ParsedToken, u: Usage, root=None):
+        super().__init__(ptok, u, root=root)
+
+        self.db_name = u.get_config("db_name")
         self.db_host = u.get_config("host")
         self.db_user = u.get_config("user")
         self.db_pass = u.get_config("password")
         self.db_port = u.get_config("port")
-        self.db_ssl  = u.get_config("ssl")
+        self.db_ssl = u.get_config("ssl")
 
-        self.params_field = "params"  # optional: list/tuple (positional) or dict (named)
+        # Standard params field: single-exec params (list/tuple/dict/single value)
+        self.params_field = "params"
+
+        # Legacy batch path: list[tuple|list|dict] → executemany
+        self.batch_field = "batch_params"
+
+        # Explicit JSON payload field (no query sniffing).
+        # If present, this value is passed to cur.execute(query, json_params).
+        self.json_params_field = "json_params"
+
+        # One DB client (and thus one connection) per PostgresPipe instance.
+        # Under your invariant (one thread per pipe), this is thread-safe.
+        self.client = DBClient(
+            host=self.db_host,
+            username=self.db_user,
+            password=self.db_pass,
+            db_name=self.db_name,
+            port=self.db_port,
+            ssl=self.db_ssl,
+        )
 
     def reset(self):
         # stateless across reset
         pass
 
+    def close(self):
+        if self.client is not None:
+            self.client.close()
+
     def _make_header(self, cur, query: str, params=None) -> Dict[str, Any]:
         """
         Inspect the cursor and build a full header record.
@@ -136,7 +183,7 @@ class PostgresPipe(QueryPipe,Integration):
             "db": self.db_name,
             "dbhost": self.db_host,
         }
-        if params:
+        if params is not None:
             h["params"] = params
 
         if cur.description:
@@ -155,64 +202,67 @@ class PostgresPipe(QueryPipe,Integration):
         return h
 
     def execute_query_returning_S_xO_iterable(self, record):
-        client = DBClient(
-            host=self.db_host,
-            username=self.db_user,
-            password=self.db_pass,
-            db_name=self.db_name,
-            port=self.db_port,
-            ssl=self.db_ssl,
-        )
+        query = record.get(self.query_field)
+        if not query:
+            record["_error"] = "missing query"
+            yield record
+            return
+
+        # Priority: json_params > batch_params > params
+        json_params = record.get(self.json_params_field, None)
+        batch = record.get(self.batch_field, None)
+        params = record.get(self.params_field, None)
+
+        cur = self.client.conn.cursor()
         try:
-            query = record.get(self.query_field)
-            if not query:
-                record['_error'] = 'missing query'
-                yield record
-                return
-
-            params = record.get(self.params_field)          # single-exec params
-            batch  = record.get("batch_params", None)       # list[tuple|dict] for batching
-
-            cur = client.conn.cursor()
-            try:
-                did_executemany = False
-
-                # ---------- execute ----------
-                if batch is not None:
-                    # Handle batch sizes explicitly to preserve single-SELECT streaming semantics
-                    if len(batch) == 0:
-                        # No-op batch; execute a lightweight statement so we can still emit a header
-                        cur.execute("SELECT 1")
-                        header_params = {"batch_size": 0}
-                    elif len(batch) == 1:
-                        cur.execute(query, batch[0])
-                        header_params = batch[0]
-                    else:
-                        cur.executemany(query, batch)
-                        did_executemany = True
-                        header_params = {"batch_size": len(batch)}
+            did_executemany = False
+            header_params = None
+
+            # ---------- execute ----------
+            if json_params is not None:
+                # Explicit JSON payload; caller controls shape.
+                # We don't inspect query or payload.
+                if isinstance(json_params, (list, tuple, dict)):
+                    cur.execute(query, json_params)
+                else:
+                    cur.execute(query, (json_params,))
+                header_params = {self.json_params_field: json_params}
+
+            elif batch is not None:
+                # Legacy executemany path; no magic.
+                if len(batch) == 0:
+                    cur.execute("SELECT 1")
+                    header_params = {"batch_size": 0}
+                elif len(batch) == 1:
+                    cur.execute(query, batch[0])
+                    header_params = {"batch_size": 1, "params": batch[0]}
+                else:
+                    cur.executemany(query, batch)
+                    did_executemany = True
+                    header_params = {"batch_size": len(batch)}
+
+            else:
+                # Single-statement path.
+                if params is None:
+                    cur.execute(query)
+                    header_params = None
                 else:
-                    if params is None:
-                        cur.execute(query)
+                    if isinstance(params, (list, tuple, dict)):
+                        cur.execute(query, params)
                     else:
-                        if isinstance(params, (list, tuple, dict)):
-                            cur.execute(query, params)
-                        else:
-                            cur.execute(query, (params,))
+                        cur.execute(query, (params,))
                     header_params = params
 
-                # ---------- header ----------
-                yield self._make_header(cur, query, header_params)
+            # ---------- header ----------
+            yield self._make_header(cur, query, header_params)
 
-                # ---------- stream rows (only meaningful for single execute that returns rows) ----------
-                # Note: executemany() typically doesn't expose per-execution result sets.
-                if not did_executemany and cur.description:
-                    cols = [d[0] for d in cur.description]
-                    if not (len(cols) == 1 and cols[0] == "ingest_event"):
-                        for row in cur:
-                            yield _row_to_dict(cur, row)
+            # ---------- stream rows (only meaningful for single execute that returns rows) ----------
+            if not did_executemany and cur.description:
+                cols = [d[0] for d in cur.description]
+                if not (len(cols) == 1 and cols[0] == "ingest_event"):
+                    for row in cur:
+                        yield _row_to_dict(cur, row)
 
-            finally:
-                cur.close()
         finally:
-            client.close()
+            cur.close()
+            # connection stays open for this pipe; closed in .close()

pjk/pipes/factory.py

@@ -38,7 +38,7 @@ COMPONENTS = {
         'reduce': ReducePipe,        
         'sort': SortPipe,
         'where': WherePipe,
-        'sel': SelectFields,
+        'select': SelectFields,
         'sample': SamplePipe,
         'explode': DenormPipe,
         'postgres': PostgresPipe,

pjk/pipes/query_pipe.py

@@ -35,8 +35,8 @@ class QueryPipe(Pipe):
         return u
 
 
-    def __init__(self, ptok: ParsedToken, usage: Usage):
-        super().__init__(ptok, usage)
+    def __init__(self, ptok: ParsedToken, usage: Usage, root = None):
+        super().__init__(ptok, usage, root=root)
         self.output_shape = usage.get_param('shape')
         self.count = usage.get_param('count')
         self.query_field = 'query' # for all subclasses

pjk/pipes/select.py

@@ -10,12 +10,12 @@ class SelectFields(DeepCopyPipe):
     @classmethod
     def usage(cls):
         usage = Usage(
-            name='sel',
+            name='select',
             desc='Select specific fields from each record.',
             component_class=cls
         )
         usage.def_arg(name='fields', usage='Comma-separated list of fields to retain')
-        usage.def_example(expr_tokens=["{id:1, dir:'up', color:'blue'}", 'sel:id,color'], expect="id: 1, color:'blue'")
+        usage.def_example(expr_tokens=["{id:1, dir:'up', color:'blue'}", 'select:id,color'], expect="id: 1, color:'blue'")
         return usage
 
     def __init__(self, ptok: ParsedToken, usage: Usage):

pjk/sinks/graph_bar_line.py

@@ -20,10 +20,6 @@ from typing import Any, Dict, Iterable, List, Optional, Sequence
 from datetime import date, datetime
 from collections import defaultdict
 
-import numpy as np
-import pandas as pd
-
-
 # ----------------------------- Public Params -----------------------------
 @dataclass
 class GraphParams:
@@ -48,6 +44,8 @@ class TimeDetector:
 
     @staticmethod
     def is_time(xs: pd.Series) -> bool:
+        import numpy as np # lazy
+        import pandas as pd # lazy
         # Already datetime dtype?
         if pd.api.types.is_datetime64_any_dtype(xs):
             return True
@@ -74,6 +72,7 @@ class TimeDetector:
 
     @staticmethod
     def parse_times(series: pd.Series) -> pd.Series:
+        import pandas as pd # lazy
         numeric = pd.to_numeric(series, errors="coerce")
         parsed = None
         if numeric.notna().mean() >= 0.9:
@@ -92,6 +91,8 @@ class MultiYAdapter:
     """Builds wide dataframe: columns = ['x'] + y_fields; sums duplicates of x."""
     @staticmethod
     def to_df(records: Iterable[Dict[str, Any]], x_field: str, y_fields: Sequence[str]) -> pd.DataFrame:
+        import pandas as pd # lazy
+        import numpy as np # lazy
         rows: List[Dict[str, Any]] = []
         for r in records:
             if x_field not in r:
@@ -120,7 +121,8 @@ class MultiYAdapter:
 class SingleYWithSetsAdapter:
     """Legacy: single y_field + optional per-row set_name to create series."""
     @staticmethod
-    def to_df(records: Iterable[Dict[str, Any]], x_field: str, y_field: str) -> pd.DataFrame:
+    def to_df(records: Iterable[Dict[str, Any]], x_field: str, y_field: str):
+        import pandas as pd # lazy
         triplets = []  # (x, y, set_name)
         for r in records:
             if x_field in r and y_field in r:
@@ -141,8 +143,10 @@ class GraphPlotter:
         self.y_fields = list(dict.fromkeys(self.pms.y_fields))  # dedupe, preserve order
 
     def plot(self, chart_type: str = "line"):
-        import matplotlib.pyplot as plt
-        import matplotlib.dates as mdates
+        import matplotlib.pyplot as plt # lazy
+        import matplotlib.dates as mdates # lazy
+        import pandas as pd # lazy
+        import numpy as np # lazy
 
         fig = plt.figure()
         ax = plt.gca()
@@ -258,7 +262,7 @@ class GraphPlotter:
     # ---------- Formatting helpers ----------
     @staticmethod
     def _format_time_axis(ax, df: pd.DataFrame) -> None:
-        import matplotlib.dates as mdates
+        import matplotlib.dates as mdates # lazy
         fig = ax.get_figure()
         ts = df["ts"]
         if ts.empty:
@@ -289,6 +293,7 @@ class GraphPlotter:
 
     def _bars_time(self, ax, df: pd.DataFrame, y_cols: Sequence[str]) -> None:
         # Grouped bars at each timestamp using index positions
+        import numpy as np # lazy
         x_vals = df["ts"].to_numpy(); idx = np.arange(len(x_vals))
         n = len(y_cols); width = 0.8 / max(n, 1)
         for i, y in enumerate(y_cols):
@@ -298,6 +303,7 @@ class GraphPlotter:
         ax.set_xticks(idx, [pd.to_datetime(t).strftime("%Y-%m-%d %H:%M") for t in x_vals], rotation=45)
 
     def _bars_categorical(self, ax, df: pd.DataFrame, y_cols: Sequence[str]) -> None:
+        import numpy as np # lazy
         seen = set(); ordered_x: List[Any] = []
         for x in df["x"].tolist():
             if x not in seen:
@@ -310,6 +316,7 @@ class GraphPlotter:
         ax.set_xticks(idx, ordered_x, rotation=45)
 
     def _lines_categorical(self, ax, df: pd.DataFrame, y_cols: Sequence[str]) -> None:
+        import numpy as np # lazy
         seen = set(); ordered_x: List[Any] = []
         for x in df["x"].tolist():
             if x not in seen:
@@ -322,7 +329,7 @@ class GraphPlotter:
 
     # ---------- Misc ----------
     def _apply_args_dict(self) -> None:
-        import matplotlib.pyplot as plt
+        import matplotlib.pyplot as plt # lazy
         for name, val in getattr(self.pms, "args_dict", {}).items():
             fn = getattr(plt, name, None)
             if callable(fn):
@@ -345,7 +352,7 @@ def graph_bar_line(obj, type):
     Returns (fig, ax) for optional downstream tweaks (safe to ignore).
     """
     # Lazy import (ensures MPL backend)
-    import matplotlib.pyplot as plt  # noqa: F401
+    import matplotlib.pyplot as plt  # noqa: F401 # lazy
 
     # Normalize y_fields from string or list
     raw_y = obj.y_field if isinstance(obj.y_field, str) else str(obj.y_field)

pjk/sources/factory.py

@@ -13,9 +13,11 @@ from pjk.sources.inline_source import InlineSource
 from pjk.sources.user_source_factory import UserSourceFactory
 from pjk.sources.parquet_source import ParquetSource
 from pjk.sources.format_source import FormatSource
+from pjk.sources.s3_select_source import S3SelectSource
 
 COMPONENTS = {
         'inline': InlineSource,
+        's3s': S3SelectSource,
         'json': JsonSource,
         'jsonl': JsonSource,
         'csv': CSVSource,
@@ -35,6 +37,11 @@ class SourceFactory(ComponentFactory):
     def create(self, token: str) -> Source:
         token = token.strip()
 
+        # s3s is a pseudo source only in the above list to provide easy man page
+        # it's instantiated by the parser when <file>.s3s, so disallow standard search for it.
+        if token == 's3s':
+            return None
+
         if InlineSource.is_inline(token):
             return InlineSource(token)
         
@@ -44,6 +51,12 @@ class SourceFactory(ComponentFactory):
             source = UserSourceFactory.create(ptok)
             if source:
                 return source
+            
+        # s3 select file
+        if ptok.pre_colon.endswith('.s3s'):
+            source = S3SelectSource(ptok, None)
+            if source:
+                return source
 
         source_cls = self.get_component_class(ptok.pre_colon)
         if source_cls and not issubclass(source_cls, FormatSource):

pjk/sources/npy_source.py

@@ -4,18 +4,15 @@
 import json
 from typing import Iterator, Dict, Any
 
-import numpy as np
-from pjk.usage import NoBindUsage
-from pjk.components import Source
 from pjk.sources.lazy_file import LazyFile
 from pjk.sources.format_source import FormatSource
 from pjk.log import logger
 
-
 class NpySource(FormatSource):
     extension = 'npy'
 
     def __init__(self, lazy_file: LazyFile):
+        super().__init__(root=None)
         self.lazy_file = lazy_file
         self.num_vecs = 0
 
@@ -32,9 +29,11 @@ class NpySource(FormatSource):
 
         try:
             # Use mmap to avoid loading entire array in RAM at once.
+            import numpy as np #lazy import
             arr = np.load(path, mmap_mode="r", allow_pickle=False)
         except Exception as e:
             logger.error(f"Failed to load .npy file at {path}: {e}")
+            raise Exception(f"Failed to load .npy file at {path}: {e}")
             return
 
         if arr.size == 0:

pjk/sources/s3_select_source.py

@@ -0,0 +1,373 @@
+# SPDX-License-Identifier: Apache-2.0
+from __future__ import annotations
+
+import json
+import re
+from threading import Lock
+from typing import Any, Dict, Iterator, List, Optional, Tuple
+
+import yaml
+
+from pjk.components import Source
+from pjk.usage import ParsedToken, Usage
+
+
+# ============================================================
+#  Per-object S3 Select reader
+# ============================================================
+
+class S3SelectObjectSource(Source):
+    """
+    Runs S3 Select on a single S3 key and streams JSON rows.
+    """
+
+    def __init__(
+        self,
+        s3_client,
+        bucket: str,
+        key: str,
+        query: str,
+        input_format: str,
+        is_gz: bool,
+    ):
+        super().__init__(root=None)
+        self._s3 = s3_client
+        self._bucket = bucket
+        self._key = key
+        self._query = query
+        self._input_format = input_format
+        self._is_gz = is_gz
+
+    def _build_input_serialization(self) -> Dict[str, Any]:
+        fmt = self._input_format.lower()
+
+        if fmt == "json":
+            base = {"JSON": {"Type": "LINES"}}
+        elif fmt == "csv":
+            base = {"CSV": {"FileHeaderInfo": "USE", "FieldDelimiter": ","}}
+        elif fmt == "tsv":
+            base = {"CSV": {"FileHeaderInfo": "USE", "FieldDelimiter": "\t"}}
+        elif fmt == "parquet":
+            base = {"Parquet": {}}
+        else:
+            raise ValueError(f"Unsupported input format for S3 Select: {self._input_format}")
+
+        if self._is_gz:
+            base["CompressionType"] = "GZIP"
+
+        return base
+
+    def __iter__(self) -> Iterator[Dict[str, Any]]:
+        input_ser = self._build_input_serialization()
+        output_ser = {"JSON": {}}  # JSON objects per record
+
+        resp = self._s3.select_object_content(
+            Bucket=self._bucket,
+            Key=self._key,
+            ExpressionType="SQL",
+            Expression=self._query,
+            InputSerialization=input_ser,
+            OutputSerialization=output_ser,
+        )
+
+        decoder = json.JSONDecoder()
+        buffer = ""
+
+        for event in resp["Payload"]:
+            if "Records" not in event:
+                continue
+
+            chunk = event["Records"]["Payload"].decode("utf-8")
+            buffer += chunk
+
+            # peel off as many complete JSON objects as we can
+            while True:
+                stripped = buffer.lstrip()
+                if not stripped:
+                    buffer = ""
+                    break
+
+                try:
+                    obj, end = decoder.raw_decode(stripped)
+                except json.JSONDecodeError:
+                    # incomplete JSON; wait for more data
+                    break
+
+                yield obj
+                buffer = stripped[end:]
+
+
+# ============================================================
+#  Shared S3 Select state (prefix iteration + threading)
+# ============================================================
+
+class _SharedS3SelectState:
+    """
+    Shared, thread-safe lazy iterator over S3 objects for S3 Select queries.
+
+    Driven by a .s3s YAML config that specifies:
+      - s3_bucket
+      - prefix
+      - optional sub_keys: [ "01", "02", ... ]
+      - optional key_regex
+      - format
+      - query
+    """
+
+    def __init__(
+        self,
+        s3_client,
+        bucket: str,
+        prefixes: List[str],
+        format_override: str,
+        query: str,
+        key_regex: Optional[str] = None,
+    ):
+        self.s3 = s3_client
+        self.bucket = bucket
+        self.prefixes = prefixes
+        self.format_override = format_override
+        self.query = query
+
+        self._key_regex = re.compile(key_regex) if key_regex else None
+
+        self._lock = Lock()
+        self._exhausted = False
+        self._prefix_index = 0
+        self._current_iter: Optional[Iterator[str]] = None
+
+    @staticmethod
+    def _get_format_gz(value: str) -> Tuple[str, bool]:
+        """
+        Split 'json.gz' -> ('json', True), 'json' -> ('json', False), etc.
+        """
+        is_gz = value.endswith(".gz")
+        fmt = value[:-3] if is_gz else value
+        return fmt, is_gz
+
+    def _next_key_iter(self) -> Optional[Iterator[str]]:
+        """
+        Get an iterator over keys for the next prefix in self.prefixes.
+        """
+        if self._prefix_index >= len(self.prefixes):
+            return None
+
+        prefix = self.prefixes[self._prefix_index]
+        self._prefix_index += 1
+
+        paginator = self.s3.get_paginator("list_objects_v2")
+        return (
+            obj.get("Key")
+            for page in paginator.paginate(Bucket=self.bucket, Prefix=prefix)
+            for obj in page.get("Contents", [])
+            if isinstance(obj.get("Key"), str)
+        )
+
+    def _iter_all_keys(self) -> Iterator[str]:
+        """
+        Iterate over all keys for all prefixes, applying key_regex if present.
+        """
+        while True:
+            if self._current_iter is None:
+                self._current_iter = self._next_key_iter()
+                if self._current_iter is None:
+                    return
+
+            try:
+                key = next(self._current_iter)
+            except StopIteration:
+                self._current_iter = None
+                continue
+
+            if self._key_regex and not self._key_regex.search(key):
+                continue
+
+            yield key
+
+    def _infer_format_and_compression(self) -> Tuple[str, bool]:
+        if not self.format_override:
+            raise ValueError("format is required in .s3s config")
+        fmt, is_gz = self._get_format_gz(self.format_override)
+        return fmt, is_gz
+
+    def _build_source_for_key(self, key: str) -> Source:
+        fmt, is_gz = self._infer_format_and_compression()
+        return S3SelectObjectSource(
+            s3_client=self.s3,
+            bucket=self.bucket,
+            key=key,
+            query=self.query,
+            input_format=fmt,
+            is_gz=is_gz,
+        )
+
+    def reserve_next_source(self) -> Optional[Source]:
+        """
+        Atomically reserve and construct the next per-key S3SelectObjectSource.
+        Returns None when exhausted.
+        """
+        if self._exhausted:
+            return None
+
+        with self._lock:
+            if self._exhausted:
+                return None
+
+            try:
+                key = next(self._iter_all_keys())
+            except StopIteration:
+                self._exhausted = True
+                return None
+
+        return self._build_source_for_key(key)
+
+
+# ============================================================
+#  Main S3 Select Source (YAML-driven)
+# ============================================================
+
+class S3SelectSource(Source):
+    """
+    S3 Select source using a .s3s YAML config file.
+
+    Example config.s3s:
+
+        s3_bucket: my-bucket
+        prefix: balancerevent/2025/11/
+        sub_keys:
+          - 01
+          - 02
+        format: json.gz
+        query: |
+          SELECT ...
+          FROM S3Object s
+
+    Parser passes the config file path as ptok.all_but_params.
+    """
+
+    extension = "s3s"
+
+    # ---------- Usage ----------
+
+    @classmethod
+    def usage(cls):
+        usage = Usage(
+            name="s3s",
+            desc=(
+                "S3 select source using <file>.s3s YAML config file.\n"
+                "where <file>.s3s e.g:\n\n"
+                "s3_bucket: my-bucket\n"
+                "prefix: my-prefix\n"
+                "sub_keys: # optional\n"
+                "- 01\n"
+                "- 02\n"
+                "format: format.gz # csv, etc\n"
+                "query: |\n"
+                " SELECT s.FooCol FROM S3Object s\n"
+                " WHERE s.IntCol = 42"
+            ),
+            component_class=cls,
+        )
+        usage.def_example(expr_tokens=["config.s3s", "-"], expect=None)
+        usage.def_syntax(None)
+        return usage
+
+    # ---------- Construction ----------
+
+    def __init__(self, ptok: ParsedToken, usage: Usage):
+        super().__init__(root=None)
+
+        config_path = ptok.all_but_params
+        cfg = self._load_config(config_path)
+
+        bucket = cfg.get("s3_bucket")
+        prefix = cfg.get("prefix")
+        fmt = cfg.get("format")
+        query = cfg.get("query")
+
+        if not bucket:
+            raise ValueError("s3s config must include 's3_bucket'")
+        if not prefix:
+            raise ValueError("s3s config must include 'prefix'")
+        if not fmt:
+            raise ValueError("s3s config must include 'format'")
+        if not query:
+            raise ValueError("s3s config must include 'query'")
+
+        key_regex = cfg.get("key_regex")
+
+        prefixes = self._build_prefixes_from_config(prefix, cfg.get("sub_keys"))
+
+        import boto3  # lazy
+        s3 = boto3.client("s3")
+
+        state = _SharedS3SelectState(
+            s3_client=s3,
+            bucket=bucket,
+            prefixes=prefixes,
+            format_override=fmt,
+            query=query,
+            key_regex=key_regex,
+        )
+
+        reserved = state.reserve_next_source()
+
+        self._state = state
+        self._current: Optional[Source] = reserved
+
+    # Alternate ctor used by deep_copy
+    @classmethod
+    def _from_state(cls, state: _SharedS3SelectState, reserved: Optional[Source]):
+        obj = cls.__new__(cls)
+        Source.__init__(obj, root=None)
+        obj._state = state
+        obj._current = reserved
+        return obj
+
+    # ---------- Iteration / deep_copy ----------
+
+    def __iter__(self):
+        while True:
+            if self._current is None:
+                self._current = self._state.reserve_next_source()
+                if self._current is None:
+                    return  # exhausted
+
+            for record in self._current:
+                yield record
+
+            self._current = None
+
+    def deep_copy(self):
+        reserved = self._state.reserve_next_source()
+        if reserved is None:
+            return None
+        return self._from_state(self._state, reserved)
+
+    # ---------- Config helpers ----------
+
+    @staticmethod
+    def _load_config(path: str) -> Dict[str, Any]:
+        with open(path, "r", encoding="utf-8") as f:
+            cfg = yaml.safe_load(f)
+        if not isinstance(cfg, dict):
+            raise ValueError("s3s config must be a YAML mapping at top level")
+        return cfg
+
+    @staticmethod
+    def _build_prefixes_from_config(prefix: str, sub_keys: Optional[List[Any]]) -> List[str]:
+        """
+        If sub_keys present (list of suffix strings), produce prefix+suffix
+        for each; otherwise just [prefix].
+        """
+        if not sub_keys:
+            return [prefix]
+
+        result: List[str] = []
+        for s in sub_keys:
+            # YAML might give ints or strings; normalize to str and strip
+            suffix = str(s).strip()
+            if not suffix:
+                continue
+            result.append(f"{prefix}{suffix}")
+        return result

pjk/sources/sql_source.py

@@ -15,13 +15,22 @@ class SQLSource(FormatSource):
     def __init__(self, lazy_file: LazyFile):
         super().__init__(root=None)
         self.lazy_file = lazy_file
-        self.num_recs = 0
 
     def __iter__(self):
+        lines = []
         with self.lazy_file.open() as f:
-            sql_text = f.read().strip()
-            sql_text = sql_text.replace("\r", " ").replace("\n", " ").strip()
+            for line in f:
+                line = line.strip()
+                if len(line) == 0:
+                    continue
+
+                if '#' in line:
+                    line = line.split('#')[0]
+                if '--' in line:
+                    line = line.split('--')[0]
+                lines.append(line)
+
+            sql_text = ' '.join(lines)
 
             if sql_text:
-                self.num_recs += 1
                 yield {"query": sql_text}

pjk/version.py

@@ -1,4 +1,4 @@
 # SPDX-License-Identifier: Apache-2.0
 # Copyright 2024 Mike Schultz
 
-__version__ = "0.7.0"
+__version__ = "0.7.4"

{python_jack_knife-0.7.0.dist-info → python_jack_knife-0.7.4.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: python-jack-knife
-Version: 0.7.0
+Version: 0.7.4
 Summary: Python Jack Knife – a command line data processor
 Author-email: Mike Schultz <mike.schultz@gmail.com>
 License: 

{python_jack_knife-0.7.0.dist-info → python_jack_knife-0.7.4.dist-info}/RECORD

@@ -1,7 +1,7 @@
 pjk/__init__.py,sha256=6HGDVcFOFv6VPSNjxVnusm9wHqy01pELX3AyCWFzqWg,128
 pjk/common.py,sha256=8f9oBKF2ydeAPKnP3bd5XG-JDSIb_UcUUMI-_h_MKmY,5125
 pjk/components.py,sha256=0HGN0V-CyAMOciyoRjr6-Y_bLFqVvW3bLRK2ymFGvzk,3655
-pjk/history.py,sha256=ef0kkXSD7uwXP4Y4VDB8SAqJ7E4CxLNryhqZbmaN0xg,2979
+pjk/history.py,sha256=UAzV40GqS3yzTTDDUiV2WAR_CHkMnjDOwKWtj6MuRrs,3076
 pjk/log.py,sha256=LjBboWhUrr2Cz-ygpftCIK17ee2-zNHKZjeJyoW0PlA,2163
 pjk/main.py,sha256=EqU6fWLFK5fD3l701XVv8p6FpvsH2iNcEWHtMYpz7j8,4463
 pjk/man_page.py,sha256=U8s--CH6TYfR98lFr3oiII6uCOHcgxAabPAK_N5VLH8,6327
@@ -9,15 +9,15 @@ pjk/parser.py,sha256=3oFAssyK3BURULPWJegCQcmUI3yiAU_bdKF4LmKcOjM,13667
 pjk/progress.py,sha256=G9igcBy2B9-QuwysKF65KR7inK2l1u8qts6ulsQMPuk,10290
 pjk/registry.py,sha256=u2LJcjZcoD2OfpON6cgZT-ZI_oCwCK6mQL8OTIppo_o,8328
 pjk/usage.py,sha256=Q3BbCld65h9Mn6opjCnlE5YIj8XYWMZk44p_rK-SiJA,11509
-pjk/version.py,sha256=eWZffaiHpGO5lGKztmHNMuxfIDJjh3U-_lvjHFjJQzQ,91
+pjk/version.py,sha256=hjsaIvpSnszaMTm4_LAcj0sFlUX7XHcK8bbB9SZYLxk,91
 pjk/integrations/opensearch_client.py,sha256=2lpN1ZV7qMcpIq-V0Pih2Xs08IBgr0-PQBOZDNxfm5Y,3572
 pjk/integrations/opensearch_index_sink.py,sha256=1heDtR85iHSzxn2j9j2TrzZPPVnxs2sdCx1EaD40iRQ,3537
 pjk/integrations/opensearch_query_pipe.py,sha256=mD9U0Jz6t4I8qE-bzGbyLOiCG_8kwK_Lr6UrxopYCqg,4519
-pjk/integrations/postgres_pipe.py,sha256=WQbHL7_w2338y0SqCDvT_44nTGsdzhda6Mp6gTNaYpg,7466
+pjk/integrations/postgres_pipe.py,sha256=IOROg-jW2C5Rl0vfaizOOjOPKQtvBti5WYTnBvHNnhk,8530
 pjk/integrations/snowflake_pipe.py,sha256=TK7pkLyPREv-LEF2LznFtoe_UWg83n6s8KOjBWM-4C4,7638
 pjk/pipes/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 pjk/pipes/denorm.py,sha256=NBVWO88V-wal4RgtFrrksSU6yVj8qGTG7jcy7VuVFm0,2273
-pjk/pipes/factory.py,sha256=qhWbqDsM-OsBpcorE615iJ8XHScgJ9MZPKrLAsL2wSY,2258
+pjk/pipes/factory.py,sha256=JrnwdCTw-441hTBQXvYh0GfBvjCrHK3Gja5brGfSq3s,2261
 pjk/pipes/filter.py,sha256=K-ScDrUCb7m4Q8bRMahixJ9zuJNfVduNoX3GEHetSR4,2091
 pjk/pipes/head.py,sha256=Dy9EXza8TMi4G0ErO1OE5KzGYUgSF_V0ibB0lNBq2cE,930
 pjk/pipes/join.py,sha256=11Q9VGAQ71L0ywtuy1eGbrxLX_YBJKaPZ9KnbJmbOQc,3822
@@ -25,10 +25,10 @@ pjk/pipes/let_reduce.py,sha256=33x-Zxi6IfD85396-uw3rDiFUro-ZNp9-OHz9W49T-s,7212
 pjk/pipes/map.py,sha256=EDDZ9kV8HZwSpuiH7J8uTcIj3zZ_1AFbECkBISwErHI,5175
 pjk/pipes/move_field.py,sha256=0uljmqurslrJ-rXQuj28QvhWhb-uuuvg_dyYXguNeJY,1169
 pjk/pipes/progress_pipe.py,sha256=d8-QeD7Uo7MoSjZYbzAuy4c-cSIcDRVQnFII-1proto,1245
-pjk/pipes/query_pipe.py,sha256=5TX7Jy1ugQAftGBpfNb0QYii4r_tgQg9MR_UCf_iylY,3577
+pjk/pipes/query_pipe.py,sha256=52LABmTTgzlBECUkhxbhEzCrzyB18JCAhpGNEULtC0w,3601
 pjk/pipes/remove_field.py,sha256=nB6QOT-VvNfvn536tpBTIxgG41VjHKY7sByzmH245pI,1219
 pjk/pipes/sample.py,sha256=Bo0OJC-PFAIZfF0EnfdtoFhnJlK7d2AK8mvfjlPmxSg,2296
-pjk/pipes/select.py,sha256=y856h9ked4WqvavxrCAeubpoDUTRQFdssaNVnngQFkA,1358
+pjk/pipes/select.py,sha256=85mlpFrhSlmhsSoGLn6_A7tTwbf6yc_xSZGM63MpIOI,1364
 pjk/pipes/sort.py,sha256=CKs9uXMeHNrLpuSqCAwLED6y-Ajc-kcA_2n6yJzcy2s,2347
 pjk/pipes/tail.py,sha256=_pAqN3cPz8ByCgfrRxbTg0TNNC2gyWIFt4hhXOxDVW0,1097
 pjk/pipes/user_pipe_factory.py,sha256=aGyU1ig4HOD5I8SVnGIra7f_sCbIZa5SEVJqmzYRl9E,1599
@@ -42,7 +42,7 @@ pjk/sinks/expect.py,sha256=oTJ7MsHn1b_xc6V9KQO4ZkCIPJnYedAoNs6rC-zTMC0,3201
 pjk/sinks/factory.py,sha256=uE6ETzliuRsoU1MV-Rc9f7TfVIX7hePp69QLkSqCEbU,1823
 pjk/sinks/format_sink.py,sha256=HCyBJi2cdfKOOn-ebVXbGKYFBJHL7g_VM_gCz1e0aW0,4986
 pjk/sinks/graph.py,sha256=rQGa-Us2LIaK-gPGX6y7hng3B6Bz5EG4MCB-PQ3b3hk,2132
-pjk/sinks/graph_bar_line.py,sha256=N_tgnAjV0tqzrc060-Xr-YOKzznmGo9uCiqzp8K3mzU,16257
+pjk/sinks/graph_bar_line.py,sha256=b37Yf9brzH_kCnhBdYd2budW_dMJDIKy3KHWFP1zqHc,16614
 pjk/sinks/graph_cumulative.py,sha256=fFXI9MSLhxKk5Xwb4df4-QqrctHeiwyXqj4S_pUNHAw,1711
 pjk/sinks/graph_hist.py,sha256=drkAeAMhSXRM-Qm_xfK7WJ1u_usMVlC_TDP1GF_xppI,1994
 pjk/sinks/graph_scatter.py,sha256=3nnIdux9oy8Na2Nt80UzPm03abEglXZyrrHB6ciJabc,1027
@@ -56,7 +56,7 @@ pjk/sinks/user_sink_factory.py,sha256=HEJVNU7Raf6DeU7wvcfziafA4odeVafn5v5X9VV_xt
 pjk/sources/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 pjk/sources/csv_source.py,sha256=_99z8niT6K8jum2TiHhM0VVrf3VuQLpEHhEIKSRZP0s,765
 pjk/sources/dir_source.py,sha256=gdHVxcDA8qnGEt_ChMbG013y-Pp4I6Asm673S6i54qA,5939
-pjk/sources/factory.py,sha256=PIEVTUsNYi3ha64ngRM08ZVazol-RcT-NJ_KcpLrvgM,1736
+pjk/sources/factory.py,sha256=LfNpsZ1NYyI4Q0jo1lAfxqGS-uU5EkhfrtyFsK7cz8g,2234
 pjk/sources/favorite_source.py,sha256=ZgpK_kNeB_rWf0LvZ9zDzwdwCwVxbyeZ8xE5EPs40fQ,1237
 pjk/sources/format_source.py,sha256=qRq6WoRyt-u9fZQ1sBEzQVvN4-0hQXuB88_eFP0T7MM,5064
 pjk/sources/inline_source.py,sha256=t3J1FEr45hySSYLF5gEUPeNo_Fr0M7G3hXa3-J-56XI,1766
@@ -64,16 +64,17 @@ pjk/sources/json_source.py,sha256=3OQ61YHjhVAITWrCf0J9JNld1pFGY1P1Kd2H4jopDPM,13
 pjk/sources/lazy_file.py,sha256=fQYaQz7bytG9vY4JNtIQJxfHWFowCn5il51H7vQrTNg,400
 pjk/sources/lazy_file_local.py,sha256=giDruMzRJSfUmWtuuJcXb2mUF2Cz0og-l-HOk3tFv0I,588
 pjk/sources/lazy_file_s3.py,sha256=a4PyBM_WoHfmKrbMucTlqxOPF79KGRyJGxECD-dVq5Q,877
-pjk/sources/npy_source.py,sha256=_U5Bbh_ESSWl7Qkt7FVt0WRcah1RH6tNieT04mQ4UQM,2866
+pjk/sources/npy_source.py,sha256=kEGCeuhtaxymQzWBs2PJaMdZX1_kvAECgP5NG5Zpddk,2930
 pjk/sources/parquet_source.py,sha256=Dcou-p3xB1M1NUtDDEj1m0fnb8IaWdiSLr5qp7JGAfw,831
+pjk/sources/s3_select_source.py,sha256=Rcg_pcoBXwy9TsEcXQNymJo5h_iTps3cT2K_EiyvtG8,10928
 pjk/sources/s3_source.py,sha256=CSZYpjxqNoQcqgZCmWJXNG6Q-8XAThz0bF9ycjtuQ6Y,5019
 pjk/sources/source_list.py,sha256=E70ewbGvc0MiD61_nIiRAZ_71ObH7iqs2JZfSR-6BaM,650
-pjk/sources/sql_source.py,sha256=QLEHlbP9HAfPKKbcIBrCJ0DXLXg20kwkkm_oYJH3NUo,806
+pjk/sources/sql_source.py,sha256=p2hqMAqOH0qUUvoM0cvELFX24mign1K-xSXAwI5irEA,1007
 pjk/sources/tsv_source.py,sha256=37nhEblCZ8XeTNHVo-WcdJ8HbIbiwKgsDES_yzz6EdU,306
 pjk/sources/user_source_factory.py,sha256=FiUrdP_ecqYG1btGjDTOWMMYGFMhKuDZYW1BUMzwQEQ,1353
-python_jack_knife-0.7.0.dist-info/licenses/LICENSE,sha256=z8d0m5b2O9McPEK1xHG_dWgUBT6EfBDz6wA0F7xSPTA,11358
-python_jack_knife-0.7.0.dist-info/METADATA,sha256=xom-jljEUznsjT7V2B2Wto-qV6hTKLsrhFgN_bLhnsk,14641
-python_jack_knife-0.7.0.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-python_jack_knife-0.7.0.dist-info/entry_points.txt,sha256=kzZ10zEisvEaG2xYqqw7xRpuV62rAO_dPEHnM6USelk,38
-python_jack_knife-0.7.0.dist-info/top_level.txt,sha256=r-Ef_I9SbVDL9jD-W0WtshstLos_7guWbpItYxxSllQ,4
-python_jack_knife-0.7.0.dist-info/RECORD,,
+python_jack_knife-0.7.4.dist-info/licenses/LICENSE,sha256=z8d0m5b2O9McPEK1xHG_dWgUBT6EfBDz6wA0F7xSPTA,11358
+python_jack_knife-0.7.4.dist-info/METADATA,sha256=CDUdtNS7SK2wbwcnU_BrTGn-REtYKhgaKfqW5kPTrZ0,14641
+python_jack_knife-0.7.4.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+python_jack_knife-0.7.4.dist-info/entry_points.txt,sha256=kzZ10zEisvEaG2xYqqw7xRpuV62rAO_dPEHnM6USelk,38
+python_jack_knife-0.7.4.dist-info/top_level.txt,sha256=r-Ef_I9SbVDL9jD-W0WtshstLos_7guWbpItYxxSllQ,4
+python_jack_knife-0.7.4.dist-info/RECORD,,