# dbtk/writers/json.py
"""
JSON writer for database results.
"""
import json
import logging
from datetime import datetime, date, time
from typing import Union, List, Optional, Any
from pathlib import Path
from .base import BaseWriter, BatchWriter
from ..utils import to_string
logger = logging.getLogger(__name__)
[docs]
class JSONWriter(BaseWriter):
"""JSON writer class that extends BaseWriter."""
[docs]
def __init__(self,
data=None,
file: Optional[Union[str, Path]] = None,
columns: Optional[List[str]] = None,
encoding: str = 'utf-8',
indent: Optional[int] = 2,
compression: str = 'infer',
**json_kwargs):
"""
Initialize JSON writer.
Args:
data: Cursor object or list of records
file: Output file. If None, writes to stdout
columns: Column names for list-of-lists data (optional for other types)
encoding: File encoding
indent: JSON indentation - defaults to 2 (pretty-print), 0 or None for compact
compression: Compression type. 'infer' detects from file extension (.gz, .bz2, .xz).
Pass 'gzip', 'bz2', or 'lzma' to override, or None to disable.
**json_kwargs: Additional arguments passed to json.dump
"""
# Preserve data types for JSON output
super().__init__(data, file, columns, encoding, compression=compression, indent=indent, **json_kwargs)
[docs]
def to_string(self, obj: Any) -> Any:
"""Convert object to string. For JSON just convert dates and times. """
if isinstance(obj, (datetime, date, time)):
return to_string(obj)
else:
return obj
def _write_data(self, file_obj) -> None:
"""Write JSON data to file object."""
records = []
for record in self.data_iterator:
record_dict = self._row_to_dict(record)
records.append(record_dict)
self._row_num = len(records)
json.dump(records, file_obj, **self._format_kwargs)
[docs]
class NDJSONWriter(BatchWriter):
"""NDJSON (newline-delimited JSON) writer."""
[docs]
def __init__(self,
data=None,
file: Optional[Union[str, Path]] = None,
columns: Optional[List[str]] = None,
encoding: str = 'utf-8',
compression: str = 'infer',
**json_kwargs):
"""
Initialize NDJSON writer.
Args:
data: Cursor object or list of records
file: Output file. If None, writes to stdout
columns: Column names for list-of-lists data (optional for other types)
encoding: File encoding
compression: Compression type. 'infer' detects from file extension (.gz, .bz2, .xz).
Pass 'gzip', 'bz2', or 'lzma' to override, or None to disable.
**json_kwargs: Additional arguments passed to json.dumps
"""
# NDJSON doesn't use indentation
super().__init__(data, file, columns=columns, encoding=encoding, compression=compression,
indent=None, **json_kwargs)
[docs]
def to_string(self, obj: Any) -> Any:
"""Convert object to string. For JSON just convert dates and times. """
if isinstance(obj, (datetime, date, time)):
return to_string(obj)
else:
return obj
def _write_data(self, file_obj) -> None:
"""Write NDJSON data to file object."""
for record in self.data_iterator:
record_dict = self._row_to_dict(record)
json_line = json.dumps(record_dict, **self._format_kwargs)
file_obj.write(json_line + '\n')
self._row_num += 1
file_obj.flush()
[docs]
def to_json(data,
file: Optional[Union[str, Path]] = None,
encoding: str = 'utf-8',
indent: Optional[int] = 2,
compression: str = 'infer',
**json_kwargs) -> None:
"""
Export cursor or result set to JSON file as an array of dictionaries.
Args:
data: Cursor object or list of records
file: Output file. If None, writes to stdout
encoding: File encoding
indent: JSON indentation - defaults to 2 (pretty-print), 0 or None for compact
**json_kwargs: Additional arguments passed to json.dump
Example:
# Write to file as JSON array
to_json(cursor, 'users.json')
# Write to stdout
to_json(cursor)
# Compact format
to_json(cursor, 'data.json', indent=None)
"""
with JSONWriter(
data=data,
file=file,
encoding=encoding,
indent=indent,
compression=compression,
**json_kwargs
) as writer:
writer.write()
[docs]
def to_ndjson(data,
file: Optional[Union[str, Path]] = None,
encoding: str = 'utf-8',
compression: str = 'infer',
**json_kwargs) -> None:
"""
Export cursor or result set to NDJSON (newline-delimited JSON) file.
Args:
data: Cursor object or list of records
file: Output file. If None, writes to stdout
encoding: File encoding
**json_kwargs: Additional arguments passed to json.dumps
Example:
# Write to file as NDJSON
to_ndjson(cursor, 'users.ndjson')
# Write to stdout
to_ndjson(cursor)
"""
with NDJSONWriter(
data=data,
file=file,
encoding=encoding,
compression=compression,
**json_kwargs
) as writer:
writer.write()