diff --git a/databend_sqlalchemy/dml.py b/databend_sqlalchemy/dml.py index cb2ad2b..4856d64 100644 --- a/databend_sqlalchemy/dml.py +++ b/databend_sqlalchemy/dml.py @@ -277,6 +277,7 @@ def __repr__(self): class CSVFormat(CopyFormat): format_type = "CSV" + inherit_cache = False def __init__( self, @@ -340,12 +341,13 @@ def __init__( if binary_format not in ["HEX", "BASE64"]: raise TypeError('Binary Format should be "HEX" or "BASE64".') self.options["BINARY_FORMAT"] = binary_format - if compression: + if compression and compression is not Compression.NONE: self.options["COMPRESSION"] = compression.value class TSVFormat(CopyFormat): format_type = "TSV" + inherit_cache = False def __init__( self, @@ -366,12 +368,13 @@ def __init__( if len(str(field_delimiter).encode().decode("unicode_escape")) != 1: raise TypeError("Field Delimiter should be a single character") self.options["FIELD_DELIMITER"] = f"{repr(field_delimiter)}" - if compression: + if compression and compression is not Compression.NONE: self.options["COMPRESSION"] = compression.value class NDJSONFormat(CopyFormat): format_type = "NDJSON" + inherit_cache = False def __init__( self, @@ -396,12 +399,13 @@ def __init__( 'Missing Field As should be "ERROR", "NULL", "FIELD_DEFAULT" or "TYPE_DEFAULT".' ) self.options["MISSING_FIELD_AS"] = f"{missing_field_as}" - if compression: + if compression and compression is not Compression.NONE: self.options["COMPRESSION"] = compression.value class ParquetFormat(CopyFormat): format_type = "PARQUET" + inherit_cache = False def __init__( self, @@ -416,7 +420,7 @@ def __init__( 'Missing Field As should be "ERROR" or "FIELD_DEFAULT".' ) self.options["MISSING_FIELD_AS"] = f"{missing_field_as}" - if compression: + if compression and compression is not Compression.NONE: if compression not in [Compression.ZSTD, Compression.SNAPPY]: raise TypeError( 'Compression should be None, ZStd, or Snappy.' diff --git a/tests/test_copy_format.py b/tests/test_copy_format.py new file mode 100644 index 0000000..9144e6d --- /dev/null +++ b/tests/test_copy_format.py @@ -0,0 +1,57 @@ +#!/usr/bin/env python + +from sqlalchemy.testing import fixtures, assert_raises_message +from sqlalchemy.testing.assertions import AssertsCompiledSQL + +from databend_sqlalchemy import ( + CSVFormat, + TSVFormat, + NDJSONFormat, + ParquetFormat, + Compression, +) + + +class CompileCopyFormatCompressionTest(fixtures.TestBase, AssertsCompiledSQL): + + __only_on__ = "databend" + + _FORMATS = [ + (CSVFormat, "CSV"), + (TSVFormat, "TSV"), + (NDJSONFormat, "NDJSON"), + (ParquetFormat, "PARQUET"), + ] + + def test_no_compression_emits_no_option(self): + # ``None`` and the ``Compression.NONE`` sentinel both mean uncompressed. + # Regression: ``Compression.NONE`` is a truthy Enum member, so the old + # ``if compression:`` guard treated it as "a codec was set" — emitting a + # spurious ``COMPRESSION = NONE`` option and, for ParquetFormat, raising. + for fmt_cls, type_ in self._FORMATS: + expected = f"FILE_FORMAT = (TYPE = {type_})" + self.assert_compile(fmt_cls(), expected) + self.assert_compile(fmt_cls(compression=Compression.NONE), expected) + + def test_explicit_codec_is_emitted(self): + for fmt_cls, type_, codec in [ + (CSVFormat, "CSV", Compression.GZIP), + (TSVFormat, "TSV", Compression.GZIP), + (NDJSONFormat, "NDJSON", Compression.GZIP), + (ParquetFormat, "PARQUET", Compression.ZSTD), + (ParquetFormat, "PARQUET", Compression.SNAPPY), + ]: + self.assert_compile( + fmt_cls(compression=codec), + f"FILE_FORMAT = (TYPE = {type_}, COMPRESSION = {codec.value})", + ) + + def test_parquet_rejects_unsupported_codec(self): + # ParquetFormat still rejects codecs its writer can't use. + for codec in [Compression.GZIP, Compression.BZ2, Compression.ZIP]: + assert_raises_message( + TypeError, + "Compression should be None, ZStd, or Snappy.", + ParquetFormat, + compression=codec, + )