ping-dataexport 0.1.0__py3-none-any.whl

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (42) hide show
  1. ping_dataexport/__init__.py +11 -0
  2. ping_dataexport/__main__.py +54 -0
  3. ping_dataexport/api.py +164 -0
  4. ping_dataexport/applogger.py +104 -0
  5. ping_dataexport/cli.py +101 -0
  6. ping_dataexport/config.py +233 -0
  7. ping_dataexport/db/__init__.py +0 -0
  8. ping_dataexport/db/connector.py +99 -0
  9. ping_dataexport/db/oracle_connector.py +86 -0
  10. ping_dataexport/db/pyodbc_connector.py +82 -0
  11. ping_dataexport/db/sqlite_connector.py +61 -0
  12. ping_dataexport/db/type_map.py +118 -0
  13. ping_dataexport/errors.py +19 -0
  14. ping_dataexport/export/__init__.py +0 -0
  15. ping_dataexport/export/engine.py +170 -0
  16. ping_dataexport/export/options.py +48 -0
  17. ping_dataexport/export/runner.py +50 -0
  18. ping_dataexport/export/util.py +65 -0
  19. ping_dataexport/job/__init__.py +0 -0
  20. ping_dataexport/job/executor.py +107 -0
  21. ping_dataexport/job/shared_context.py +22 -0
  22. ping_dataexport/job/spec.py +251 -0
  23. ping_dataexport/job/validator.py +87 -0
  24. ping_dataexport/paths.py +71 -0
  25. ping_dataexport/planner/__init__.py +0 -0
  26. ping_dataexport/planner/base.py +71 -0
  27. ping_dataexport/planner/datelist.py +35 -0
  28. ping_dataexport/planner/groupby.py +71 -0
  29. ping_dataexport/planner/monthbydate.py +104 -0
  30. ping_dataexport/planner/single.py +18 -0
  31. ping_dataexport/planner/watermark.py +202 -0
  32. ping_dataexport/writer/__init__.py +0 -0
  33. ping_dataexport/writer/arrow_map.py +174 -0
  34. ping_dataexport/writer/base.py +37 -0
  35. ping_dataexport/writer/delimited.py +108 -0
  36. ping_dataexport/writer/parquet.py +84 -0
  37. ping_dataexport-0.1.0.dist-info/METADATA +228 -0
  38. ping_dataexport-0.1.0.dist-info/RECORD +42 -0
  39. ping_dataexport-0.1.0.dist-info/WHEEL +5 -0
  40. ping_dataexport-0.1.0.dist-info/entry_points.txt +2 -0
  41. ping_dataexport-0.1.0.dist-info/licenses/LICENSE +201 -0
  42. ping_dataexport-0.1.0.dist-info/top_level.txt +1 -0
@@ -0,0 +1,71 @@
1
+ """PartitionPlanner interface + Partition/Plan + for_options() registry.
2
+
3
+ A planner turns ExportOptions into partitions: (query with '?' placeholders,
4
+ bind params, target filename). ALL partition values are bound parameters — a
5
+ group value like O'Brien can never break the generated SQL. Placeholders are
6
+ qmark; the connector rewrites for its paramstyle and adapts param types.
7
+ """
8
+
9
+ from __future__ import annotations
10
+
11
+ from abc import ABC, abstractmethod
12
+ from dataclasses import dataclass, field
13
+ from typing import Any, Callable, List, Optional, Tuple
14
+
15
+ from ..applogger import AppLogger
16
+ from ..export.options import ExportOptions
17
+
18
+
19
+ @dataclass(frozen=True)
20
+ class Partition:
21
+ """One unit of work: query + bind params + target file (None = single output)."""
22
+
23
+ query: str
24
+ params: Tuple[Any, ...] = ()
25
+ filename: Optional[str] = None
26
+ label: str = ""
27
+ where_display: Optional[str] = None # human-readable WHERE for the logs
28
+
29
+
30
+ @dataclass
31
+ class Plan:
32
+ partitions: List[Partition] = field(default_factory=list)
33
+ #: called by the engine after ALL partitions succeed (never in --test);
34
+ #: watermark uses it to advance state — at-least-once delivery.
35
+ on_success: Optional[Callable[[], None]] = None
36
+
37
+
38
+ class PartitionPlanner(ABC):
39
+ @abstractmethod
40
+ def plan(self, options: ExportOptions, conn: Any, connector: Any, logger: AppLogger) -> Plan:
41
+ """May query the DB (groupby DISTINCT pass, watermark MAX/type checks)."""
42
+
43
+
44
+ def wrap_query(query: str) -> str:
45
+ """User query as a dialect-neutral subquery."""
46
+ return f"SELECT * FROM ({query}) QRY"
47
+
48
+
49
+ def for_options(options: ExportOptions) -> PartitionPlanner:
50
+ mode = options.mode
51
+ if mode is None or mode == "splitrow":
52
+ from .single import SinglePlanner
53
+
54
+ return SinglePlanner()
55
+ if mode == "groupby":
56
+ from .groupby import GroupByPlanner
57
+
58
+ return GroupByPlanner()
59
+ if mode in ("date", "relativedate"):
60
+ from .datelist import DateListPlanner
61
+
62
+ return DateListPlanner()
63
+ if mode == "monthbydate":
64
+ from .monthbydate import MonthByDatePlanner
65
+
66
+ return MonthByDatePlanner()
67
+ if mode == "watermark":
68
+ from .watermark import WatermarkPlanner
69
+
70
+ return WatermarkPlanner()
71
+ raise ValueError(f"no planner for mode {mode!r}") # unreachable after validation
@@ -0,0 +1,35 @@
1
+ """DateListPlanner — date and relativedate modes: one file per day.
2
+
3
+ The inclusive date list is pre-built by JobSpecBuilder (Python, never SQL).
4
+ Per date: SELECT * FROM (<q>) QRY WHERE <col> = ? bound as datetime.date.
5
+ Files: base_YYYY-MM-DD.ext
6
+ """
7
+
8
+ from __future__ import annotations
9
+
10
+ from typing import Any
11
+
12
+ from ..applogger import AppLogger
13
+ from ..export.options import ExportOptions
14
+ from ..export.util import output_filename
15
+ from .base import Partition, PartitionPlanner, Plan, wrap_query
16
+
17
+
18
+ class DateListPlanner(PartitionPlanner):
19
+ def plan(self, options: ExportOptions, conn: Any, connector: Any, logger: AppLogger) -> Plan:
20
+ col = options.column
21
+ base = wrap_query(options.query)
22
+ partitions = [
23
+ Partition(
24
+ query=f"{base} WHERE {col} = ?",
25
+ params=(day,),
26
+ filename=output_filename(
27
+ options.base_name, day.isoformat(), options.format, options.gzip
28
+ ),
29
+ label=day.isoformat(),
30
+ where_display=f"{col} = '{day.isoformat()}'",
31
+ )
32
+ for day in options.dates
33
+ ]
34
+ logger.info("EXPORT", "DATELIST", f"{len(partitions)} date partition(s)")
35
+ return Plan(partitions=partitions)
@@ -0,0 +1,71 @@
1
+ """GroupByPlanner — one file per distinct value of the group column.
2
+
3
+ DISTINCT pass: SELECT DISTINCT <col> FROM (<q>) A ORDER BY <col>
4
+ Per value: SELECT * FROM (<q>) QRY WHERE <col> = ? (NULL group -> IS NULL)
5
+ Files: base_<value>.ext, base_NULL.ext
6
+ """
7
+
8
+ from __future__ import annotations
9
+
10
+ import datetime
11
+ from typing import Any
12
+
13
+ from ..applogger import AppLogger
14
+ from ..errors import ExportError
15
+ from ..export.options import ExportOptions
16
+ from ..export.util import output_filename
17
+ from .base import Partition, PartitionPlanner, Plan, wrap_query
18
+
19
+
20
+ def _suffix_for(value: Any) -> str:
21
+ if value is None:
22
+ return "NULL"
23
+ if isinstance(value, datetime.datetime):
24
+ return value.strftime("%Y-%m-%d %H%M%S")
25
+ if isinstance(value, datetime.date):
26
+ return value.isoformat()
27
+ return str(value)
28
+
29
+
30
+ class GroupByPlanner(PartitionPlanner):
31
+ def plan(self, options: ExportOptions, conn: Any, connector: Any, logger: AppLogger) -> Plan:
32
+ col = options.column
33
+ distinct_sql = connector.rewrite_params(
34
+ f"SELECT DISTINCT {col} FROM ({options.query}) A ORDER BY {col}"
35
+ )
36
+ cursor = connector.prepare_cursor(conn, options.fetch_size)
37
+ try:
38
+ cursor.execute(distinct_sql)
39
+ values = [row[0] for row in cursor.fetchall()]
40
+ except Exception as exc:
41
+ raise ExportError(f"groupby DISTINCT query failed: {exc}") from exc
42
+ finally:
43
+ cursor.close()
44
+
45
+ logger.info("EXPORT", "GROUPBY", f"{len(values)} distinct value(s) of {col}")
46
+ base = wrap_query(options.query)
47
+ partitions = []
48
+ for value in values:
49
+ suffix = _suffix_for(value)
50
+ filename = output_filename(options.base_name, suffix, options.format, options.gzip)
51
+ if value is None:
52
+ partitions.append(
53
+ Partition(
54
+ query=f"{base} WHERE {col} IS NULL",
55
+ params=(),
56
+ filename=filename,
57
+ label=suffix,
58
+ where_display=f"{col} IS NULL",
59
+ )
60
+ )
61
+ else:
62
+ partitions.append(
63
+ Partition(
64
+ query=f"{base} WHERE {col} = ?",
65
+ params=(value,),
66
+ filename=filename,
67
+ label=suffix,
68
+ where_display=f"{col} = '{value}'",
69
+ )
70
+ )
71
+ return Plan(partitions=partitions)
@@ -0,0 +1,104 @@
1
+ """MonthByDatePlanner — one file per calendar month.
2
+
3
+ Months: N months before the anchor month .. the anchor month (anchor = --datefrom,
4
+ default today). Boundaries are generated in Python, never SQL:
5
+ full month: WHERE <col> >= ? AND <col> < ? [month start, next month start)
6
+ anchor month: WHERE <col> >= ? AND <col> <= ? (<= anchor date)
7
+ Files: base_YYYYMM.ext
8
+
9
+ Pre-check: a WHERE 1=0 metadata query must show a DATE/TIMESTAMP column —
10
+ a non-date column logs ERROR and skips the export; an inconclusive detection
11
+ (e.g. sqlite, which has no type codes) is only a warning.
12
+ """
13
+
14
+ from __future__ import annotations
15
+
16
+ import datetime
17
+ from typing import Any, List, Tuple
18
+
19
+ from ..applogger import AppLogger
20
+ from ..export.options import ExportOptions
21
+ from ..export.util import output_filename
22
+ from .base import Partition, PartitionPlanner, Plan, wrap_query
23
+
24
+
25
+ def month_starts(anchor: datetime.date, months_back: int) -> List[datetime.date]:
26
+ """First-of-month dates, oldest first, for anchor-month - months_back .. anchor-month."""
27
+ year, month = anchor.year, anchor.month
28
+ total = year * 12 + (month - 1)
29
+ starts = []
30
+ for m in range(total - months_back, total + 1):
31
+ starts.append(datetime.date(m // 12, m % 12 + 1, 1))
32
+ return starts
33
+
34
+
35
+ def next_month(day: datetime.date) -> datetime.date:
36
+ if day.month == 12:
37
+ return datetime.date(day.year + 1, 1, 1)
38
+ return datetime.date(day.year, day.month + 1, 1)
39
+
40
+
41
+ class MonthByDatePlanner(PartitionPlanner):
42
+ def plan(self, options: ExportOptions, conn: Any, connector: Any, logger: AppLogger) -> Plan:
43
+ col = options.column
44
+ kind = self._column_kind(options, conn, connector)
45
+ if kind is None:
46
+ logger.warning(
47
+ "EXPORT", "MONTHBYDATE", f"cannot verify column type of {col} — proceeding"
48
+ )
49
+ elif kind not in ("DATE", "TIMESTAMP"):
50
+ logger.error(
51
+ "EXPORT",
52
+ "MONTHBYDATE",
53
+ f"column {col} is {kind}, not DATE/DATETIME — export skipped",
54
+ )
55
+ return Plan(partitions=[])
56
+
57
+ anchor = options.datefrom
58
+ base = wrap_query(options.query)
59
+ partitions = []
60
+ starts = month_starts(anchor, options.monthsrelative)
61
+ for start in starts:
62
+ is_anchor_month = (start.year, start.month) == (anchor.year, anchor.month)
63
+ if is_anchor_month:
64
+ query = f"{base} WHERE {col} >= ? AND {col} <= ?"
65
+ params: Tuple[Any, ...] = (start, anchor)
66
+ display = f"{col} >= '{start}' AND {col} <= '{anchor}'"
67
+ else:
68
+ end = next_month(start)
69
+ query = f"{base} WHERE {col} >= ? AND {col} < ?"
70
+ params = (start, end)
71
+ display = f"{col} >= '{start}' AND {col} < '{end}'"
72
+ suffix = start.strftime("%Y%m")
73
+ partitions.append(
74
+ Partition(
75
+ query=query,
76
+ params=params,
77
+ filename=output_filename(options.base_name, suffix, options.format, options.gzip),
78
+ label=suffix,
79
+ where_display=display,
80
+ )
81
+ )
82
+ logger.info("EXPORT", "MONTHBYDATE", f"{len(partitions)} month partition(s)")
83
+ return Plan(partitions=partitions)
84
+
85
+ @staticmethod
86
+ def _column_kind(options: ExportOptions, conn: Any, connector: Any):
87
+ """Canonical kind of the mode column via a WHERE 1=0 query; None = inconclusive."""
88
+ sql = connector.rewrite_params(f"{wrap_query(options.query)} WHERE 1=0")
89
+ cursor = connector.prepare_cursor(conn, options.fetch_size)
90
+ try:
91
+ cursor.execute(sql)
92
+ columns = connector.describe(cursor, [])
93
+ except Exception:
94
+ return None # detection failure is a warning, not fatal
95
+ finally:
96
+ cursor.close()
97
+ for c in columns:
98
+ if c.name.lower() == options.column.lower():
99
+ # sample-inferred describes (sqlite) see STRING for everything on a
100
+ # zero-row query -> inconclusive; typed drivers are authoritative
101
+ if not connector.authoritative_types and c.kind == "STRING":
102
+ return None
103
+ return c.kind
104
+ return None
@@ -0,0 +1,18 @@
1
+ """SinglePlanner — default mode and splitrow: the query as-is, one partition.
2
+
3
+ For splitrow the ENGINE rotates the output file every N rows; the plan is still a
4
+ single query so the DB is read once.
5
+ """
6
+
7
+ from __future__ import annotations
8
+
9
+ from typing import Any
10
+
11
+ from ..applogger import AppLogger
12
+ from ..export.options import ExportOptions
13
+ from .base import Partition, PartitionPlanner, Plan
14
+
15
+
16
+ class SinglePlanner(PartitionPlanner):
17
+ def plan(self, options: ExportOptions, conn: Any, connector: Any, logger: AppLogger) -> Plan:
18
+ return Plan(partitions=[Partition(query=options.query, label=options.job_name)])
@@ -0,0 +1,202 @@
1
+ """WatermarkPlanner — incremental export driven by __watermark__.ini in --outputdir.
2
+
3
+ First run (missing file / blank value / FirstRun=True):
4
+ validate column type, MAX(col), export WHERE <col> IS NULL OR <col> <= ?
5
+ Incremental:
6
+ export WHERE <col> > ? AND <col> <= ? (prev, MAX); skip when no new data.
7
+ State advances ONLY after a successful export (Plan.on_success) — at-least-once
8
+ delivery. INT binds as int; DATE/DATETIME bind as date/datetime objects — the
9
+ quoted form appears only in the logged human-readable WHERE.
10
+ Files: base_YYYYMMDD_HHmmss.ext
11
+ """
12
+
13
+ from __future__ import annotations
14
+
15
+ import configparser
16
+ import datetime
17
+ from dataclasses import dataclass
18
+ from pathlib import Path
19
+ from typing import Any, Optional, Tuple
20
+
21
+ from ..applogger import AppLogger
22
+ from ..errors import ExportError
23
+ from ..export.options import ExportOptions
24
+ from ..export.util import output_filename
25
+ from .base import Partition, PartitionPlanner, Plan, wrap_query
26
+
27
+ STATE_FILENAME = "__watermark__.ini"
28
+ _SECTION = "Watermark"
29
+
30
+ # state ColumnType values — the state file format records DATETIME, not TIMESTAMP
31
+ INT_TYPE = "INT"
32
+ DATE_TYPE = "DATE"
33
+ DATETIME_TYPE = "DATETIME"
34
+
35
+
36
+ @dataclass(frozen=True)
37
+ class WatermarkState:
38
+ column: str
39
+ column_type: Optional[str]
40
+ previous_value: str
41
+ first_run: bool
42
+
43
+
44
+ def read_state(path: Path) -> Optional[WatermarkState]:
45
+ if not path.is_file():
46
+ return None
47
+ parser = configparser.RawConfigParser()
48
+ parser.optionxform = str
49
+ parser.read(path, encoding="utf-8-sig")
50
+ if not parser.has_section(_SECTION):
51
+ return None
52
+ get = lambda key: parser.get(_SECTION, key, fallback="").strip() # noqa: E731
53
+ return WatermarkState(
54
+ column=get("Column"),
55
+ column_type=get("ColumnType") or None,
56
+ previous_value=get("PreviousValue"),
57
+ first_run=get("FirstRun").lower() in ("true", "yes", "1", ""),
58
+ )
59
+
60
+
61
+ def write_state(path: Path, column: str, column_type: str, value: str) -> None:
62
+ parser = configparser.RawConfigParser()
63
+ parser.optionxform = str
64
+ parser.add_section(_SECTION)
65
+ parser.set(_SECTION, "Column", column)
66
+ parser.set(_SECTION, "ColumnType", column_type)
67
+ parser.set(_SECTION, "PreviousValue", value)
68
+ parser.set(_SECTION, "FirstRun", "False")
69
+ with open(path, "w", encoding="utf-8", newline="\n") as fh:
70
+ parser.write(fh)
71
+
72
+
73
+ def value_to_state_string(value: Any) -> str:
74
+ if isinstance(value, datetime.datetime):
75
+ return value.isoformat(sep=" ")
76
+ if isinstance(value, datetime.date):
77
+ return value.isoformat()
78
+ return str(value)
79
+
80
+
81
+ def parse_state_value(text: str, column_type: str) -> Any:
82
+ try:
83
+ if column_type == INT_TYPE:
84
+ return int(text)
85
+ if column_type == DATE_TYPE:
86
+ return datetime.date.fromisoformat(text[:10])
87
+ return datetime.datetime.fromisoformat(text.replace("T", " "))
88
+ except ValueError:
89
+ raise ExportError(f"invalid watermark PreviousValue {text!r} for {column_type}") from None
90
+
91
+
92
+ class WatermarkPlanner(PartitionPlanner):
93
+ def plan(self, options: ExportOptions, conn: Any, connector: Any, logger: AppLogger) -> Plan:
94
+ col = options.column
95
+ state_path = Path(options.output_dir) / STATE_FILENAME
96
+ state = read_state(state_path)
97
+ first_run = state is None or state.previous_value == "" or state.first_run
98
+ logger.info(
99
+ "WATERMARK", "STATE",
100
+ "first run" if first_run else f"previous value {state.previous_value}",
101
+ )
102
+
103
+ max_value, column_type = self._max_and_type(options, conn, connector)
104
+ if max_value is None:
105
+ logger.info("WATERMARK", "SKIP", "no data (MAX is NULL) — nothing to export")
106
+ return Plan(partitions=[])
107
+
108
+ base = wrap_query(options.query)
109
+ suffix = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
110
+ filename = output_filename(options.base_name, suffix, options.format, options.gzip)
111
+
112
+ if first_run:
113
+ query = f"{base} WHERE {col} IS NULL OR {col} <= ?"
114
+ params: Tuple[Any, ...] = (max_value,)
115
+ display = f"{col} IS NULL OR {col} <= '{value_to_state_string(max_value)}'"
116
+ else:
117
+ if state.column_type:
118
+ column_type = state.column_type # the state file is authoritative
119
+ prev = parse_state_value(state.previous_value, column_type)
120
+ if not self._is_newer(max_value, prev):
121
+ logger.info(
122
+ "WATERMARK", "SKIP",
123
+ f"no new data (MAX {value_to_state_string(max_value)} "
124
+ f"<= previous {state.previous_value})",
125
+ )
126
+ return Plan(partitions=[])
127
+ query = f"{base} WHERE {col} > ? AND {col} <= ?"
128
+ params = (prev, max_value)
129
+ display = (
130
+ f"{col} > '{value_to_state_string(prev)}' "
131
+ f"AND {col} <= '{value_to_state_string(max_value)}'"
132
+ )
133
+
134
+ logger.info("WATERMARK", "RANGE", display)
135
+
136
+ def advance_state() -> None:
137
+ write_state(state_path, col, column_type, value_to_state_string(max_value))
138
+ logger.info(
139
+ "WATERMARK", "ADVANCED", f"PreviousValue={value_to_state_string(max_value)}"
140
+ )
141
+
142
+ return Plan(
143
+ partitions=[
144
+ Partition(
145
+ query=query, params=params, filename=filename,
146
+ label=suffix, where_display=display,
147
+ )
148
+ ],
149
+ on_success=advance_state,
150
+ )
151
+
152
+ def _max_and_type(
153
+ self, options: ExportOptions, conn: Any, connector: Any
154
+ ) -> Tuple[Any, str]:
155
+ """MAX(col) + canonical watermark type (INT | DATE | DATETIME)."""
156
+ col = options.column
157
+ sql = connector.rewrite_params(f"SELECT MAX({col}) FROM ({options.query}) QRY")
158
+ cursor = connector.prepare_cursor(conn, options.fetch_size)
159
+ try:
160
+ cursor.execute(sql)
161
+ row = cursor.fetchone()
162
+ max_raw = row[0] if row else None
163
+ columns = connector.describe(cursor, [row] if row else [])
164
+ except Exception as exc:
165
+ raise ExportError(f"watermark MAX query failed: {exc}") from exc
166
+ finally:
167
+ cursor.close()
168
+ if max_raw is None:
169
+ return None, ""
170
+
171
+ kind = columns[0].kind if columns else "STRING"
172
+ if kind in ("INT", "BIGINT"):
173
+ return int(max_raw), INT_TYPE
174
+ if kind == "DATE":
175
+ value = (
176
+ max_raw if isinstance(max_raw, datetime.date)
177
+ and not isinstance(max_raw, datetime.datetime)
178
+ else datetime.date.fromisoformat(str(max_raw)[:10])
179
+ )
180
+ return value, DATE_TYPE
181
+ if kind == "TIMESTAMP":
182
+ value = (
183
+ max_raw if isinstance(max_raw, datetime.datetime)
184
+ else datetime.datetime.fromisoformat(str(max_raw).replace("T", " "))
185
+ )
186
+ return value, DATETIME_TYPE
187
+ raise ExportError(
188
+ f"watermark column {col} must be INT, DATE or DATETIME (detected {kind})"
189
+ )
190
+
191
+ @staticmethod
192
+ def _is_newer(max_value: Any, prev: Any) -> bool:
193
+ # DATE state with a DATETIME max (or vice versa) — align before comparing
194
+ if isinstance(prev, datetime.datetime) and not isinstance(max_value, datetime.datetime):
195
+ max_value = datetime.datetime(max_value.year, max_value.month, max_value.day)
196
+ elif (
197
+ isinstance(prev, datetime.date)
198
+ and not isinstance(prev, datetime.datetime)
199
+ and isinstance(max_value, datetime.datetime)
200
+ ):
201
+ max_value = max_value.date()
202
+ return max_value > prev
File without changes
@@ -0,0 +1,174 @@
1
+ """ArrowRowMapper — driver rows -> PyArrow, the ONLY place driver quirks meet Arrow.
2
+
3
+ - Schema: canonical ColumnInfo -> Arrow type (all fields nullable). DECIMAL keeps
4
+ decimal128(p,s) when p <= 38 — full precision survives, never downcast to DOUBLE.
5
+ - One value normalizer per column, chosen ONCE from metadata (no per-value
6
+ isinstance chains on the hot path beyond what a normalizer needs).
7
+ - RecordBatches are built with explicit types: pa.array(col, type=...) — Arrow type
8
+ inference is never used, exactly as pandas never is.
9
+
10
+ pyarrow is imported lazily so delimited-only deployments don't need it installed.
11
+ """
12
+
13
+ from __future__ import annotations
14
+
15
+ import datetime
16
+ import decimal
17
+ from typing import Any, Callable, List, Optional, Sequence
18
+
19
+ from ..db.connector import (
20
+ BIGINT,
21
+ BINARY,
22
+ BOOL,
23
+ DATE,
24
+ DECIMAL,
25
+ DOUBLE,
26
+ FLOAT,
27
+ INT,
28
+ STRING,
29
+ TIME,
30
+ TIMESTAMP,
31
+ ColumnInfo,
32
+ )
33
+ from ..errors import ExportError
34
+
35
+ _MS_PER_DAY = 86_400_000
36
+ _MAX_DECIMAL128_PRECISION = 38
37
+
38
+ Normalizer = Callable[[Any], Any]
39
+
40
+
41
+ class ArrowRowMapper:
42
+ def __init__(self, columns: List[ColumnInfo], clean_text: Callable[[str], str]):
43
+ try:
44
+ import pyarrow as pa
45
+ except ImportError:
46
+ raise ExportError("parquet output requires the pyarrow package") from None
47
+ self._pa = pa
48
+ self.columns = columns
49
+ self._types = [self._arrow_type(c) for c in columns]
50
+ self._normalizers = [self._normalizer(c, clean_text) for c in columns]
51
+ self._schema = pa.schema(
52
+ [pa.field(c.name, t, nullable=True) for c, t in zip(columns, self._types)]
53
+ )
54
+
55
+ @property
56
+ def schema(self): # pyarrow.Schema
57
+ return self._schema
58
+
59
+ def new_buffers(self) -> List[List[Any]]:
60
+ return [[] for _ in self.columns]
61
+
62
+ def append_row(self, buffers: List[List[Any]], row: Sequence[Any]) -> None:
63
+ for buf, normalize, value in zip(buffers, self._normalizers, row):
64
+ buf.append(normalize(value))
65
+
66
+ def to_record_batch(self, buffers: List[List[Any]]):
67
+ pa = self._pa
68
+ arrays = [pa.array(buf, type=t) for buf, t in zip(buffers, self._types)]
69
+ return pa.record_batch(arrays, schema=self._schema)
70
+
71
+ # -- type mapping (concept §5.2 table) ----------------------------------
72
+
73
+ def _arrow_type(self, col: ColumnInfo):
74
+ pa = self._pa
75
+ if col.kind == INT:
76
+ return pa.int32()
77
+ if col.kind == BIGINT:
78
+ return pa.int64()
79
+ if col.kind == FLOAT:
80
+ return pa.float32()
81
+ if col.kind == DOUBLE:
82
+ return pa.float64()
83
+ if col.kind == DECIMAL:
84
+ if col.precision and col.precision <= _MAX_DECIMAL128_PRECISION:
85
+ return pa.decimal128(col.precision, col.scale or 0)
86
+ return pa.float64() # p > 38: downcast, documented
87
+ if col.kind == BOOL:
88
+ return pa.bool_()
89
+ if col.kind == DATE:
90
+ return pa.date32()
91
+ if col.kind == TIME:
92
+ return pa.time32("ms")
93
+ if col.kind == TIMESTAMP:
94
+ return pa.timestamp("us")
95
+ if col.kind == BINARY:
96
+ return pa.binary()
97
+ return pa.string()
98
+
99
+ # -- per-column value normalizers ---------------------------------------
100
+
101
+ def _normalizer(self, col: ColumnInfo, clean_text: Callable[[str], str]) -> Normalizer:
102
+ if col.kind == DECIMAL and col.precision and col.precision <= _MAX_DECIMAL128_PRECISION:
103
+ quantum = decimal.Decimal(1).scaleb(-(col.scale or 0))
104
+
105
+ def norm_decimal(v: Any) -> Optional[decimal.Decimal]:
106
+ if v is None:
107
+ return None
108
+ d = v if isinstance(v, decimal.Decimal) else decimal.Decimal(str(v))
109
+ return d.quantize(quantum)
110
+
111
+ return norm_decimal
112
+
113
+ if col.kind in (DOUBLE, FLOAT) or col.kind == DECIMAL: # DECIMAL p>38 -> float64
114
+ return lambda v: None if v is None else float(v)
115
+
116
+ if col.kind in (INT, BIGINT):
117
+ return lambda v: None if v is None else int(v)
118
+
119
+ if col.kind == BOOL:
120
+ return lambda v: None if v is None else bool(v)
121
+
122
+ if col.kind == DATE:
123
+
124
+ def norm_date(v: Any) -> Optional[datetime.date]:
125
+ if v is None:
126
+ return None
127
+ if isinstance(v, datetime.datetime):
128
+ return v.date()
129
+ if isinstance(v, datetime.date):
130
+ return v
131
+ # loose drivers (sqlite, old FreeTDS) deliver ISO strings
132
+ return datetime.date.fromisoformat(str(v)[:10])
133
+
134
+ return norm_date
135
+
136
+ if col.kind == TIMESTAMP:
137
+
138
+ def norm_ts(v: Any) -> Optional[datetime.datetime]:
139
+ if v is None:
140
+ return None
141
+ if isinstance(v, datetime.datetime):
142
+ return v
143
+ if isinstance(v, datetime.date):
144
+ return datetime.datetime(v.year, v.month, v.day)
145
+ return datetime.datetime.fromisoformat(str(v).replace("T", " "))
146
+
147
+ return norm_ts
148
+
149
+ if col.kind == TIME:
150
+
151
+ def norm_time(v: Any) -> Optional[int]:
152
+ if v is None:
153
+ return None
154
+ if isinstance(v, datetime.time):
155
+ ms = ((v.hour * 60 + v.minute) * 60 + v.second) * 1000 + v.microsecond // 1000
156
+ elif isinstance(v, (int, float)):
157
+ ms = int(v)
158
+ else:
159
+ t = datetime.time.fromisoformat(str(v))
160
+ ms = ((t.hour * 60 + t.minute) * 60 + t.second) * 1000 + t.microsecond // 1000
161
+ return ms % _MS_PER_DAY
162
+
163
+ return norm_time
164
+
165
+ if col.kind == BINARY:
166
+ return lambda v: None if v is None else bytes(v)
167
+
168
+ # STRING / fallback — strip rules apply here (parquet parity with delimited)
169
+ def norm_string(v: Any) -> Optional[str]:
170
+ if v is None:
171
+ return None
172
+ return clean_text(v) if isinstance(v, str) else str(v)
173
+
174
+ return norm_string