415 lines
16 KiB
Python
415 lines
16 KiB
Python
#
|
|
# Copyright (c) 2023 Airbyte, Inc., all rights reserved.
|
|
#
|
|
|
|
import json
|
|
import os
|
|
import random
|
|
import shutil
|
|
import string
|
|
from pathlib import Path
|
|
from typing import Any, List, Mapping, Tuple
|
|
|
|
import pendulum
|
|
import pytest
|
|
from smart_open import open as smart_open
|
|
from source_s3.source_files_abstract.file_info import FileInfo
|
|
from source_s3.source_files_abstract.formats.csv_parser import CsvParser
|
|
|
|
from .abstract_test_parser import AbstractTestParser, memory_limit
|
|
from .conftest import TMP_FOLDER
|
|
|
|
SAMPLE_DIRECTORY = Path(__file__).resolve().parent.joinpath("sample_files/")
|
|
|
|
# All possible CSV data types
|
|
CSV_TYPES = {
|
|
# logical_type: (json_type, csv_types, convert_function)
|
|
# standard types
|
|
"string": ("string", ["string"], None),
|
|
"boolean": ("boolean", ["boolean"], None),
|
|
"number": ("number", ["number"], None),
|
|
"integer": ("integer", ["integer"], None),
|
|
}
|
|
|
|
|
|
def _generate_value(typ: str) -> Any:
|
|
if typ == "string":
|
|
if AbstractTestParser._generate_value("boolean"):
|
|
return None
|
|
random_length = random.randint(0, 512)
|
|
return "".join(random.choice(string.ascii_uppercase + string.digits) for _ in range(random_length))
|
|
return AbstractTestParser._generate_value(typ)
|
|
|
|
|
|
def _generate_row(types: List[str]) -> List[Any]:
|
|
"""Generates random values with request types"""
|
|
row = []
|
|
for needed_type in types:
|
|
for json_type in CSV_TYPES:
|
|
if json_type == needed_type:
|
|
value = _generate_value(needed_type)
|
|
if value is None:
|
|
value = ""
|
|
row.append(str(value))
|
|
break
|
|
return row
|
|
|
|
|
|
def generate_csv_file(filename: str, columns: Mapping[str, str], num_rows: int, delimiter: str) -> str:
|
|
"""Generates a random CSV data and save it to a tmp file"""
|
|
header_line = delimiter.join(columns.keys())
|
|
types = list(columns.values()) if num_rows else []
|
|
with open(filename, "w") as f:
|
|
f.write(header_line + "\n")
|
|
for _ in range(num_rows):
|
|
f.write(delimiter.join(_generate_row(types)) + "\n")
|
|
return filename
|
|
|
|
|
|
def generate_big_file(filepath: str, size_in_gigabytes: float, columns_number: int, template_file: str = None) -> Tuple[dict, float]:
|
|
temp_files = [filepath + ".1", filepath + ".2"]
|
|
if template_file:
|
|
shutil.copyfile(template_file, filepath)
|
|
schema = None
|
|
else:
|
|
schema = {f"column {i}": random.choice(["integer", "string", "boolean", "number"]) for i in range(columns_number)}
|
|
generate_csv_file(filepath, schema, 456, ",")
|
|
|
|
skip_headers = False
|
|
with open(filepath, "r") as f:
|
|
with open(temp_files[0], "w") as tf:
|
|
for line in f:
|
|
if not skip_headers:
|
|
skip_headers = True
|
|
continue
|
|
tf.write(str(line))
|
|
|
|
with open(filepath, "ab") as f:
|
|
while True:
|
|
file_size = os.stat(filepath).st_size / (1024**3)
|
|
if file_size > size_in_gigabytes:
|
|
break
|
|
with open(temp_files[0], "rb") as tf: # type: ignore[assignment]
|
|
with open(temp_files[1], "wb") as tf2:
|
|
buf = tf.read(50 * 1024**2) # by 50Mb
|
|
if buf:
|
|
f.write(buf) # type: ignore[arg-type]
|
|
tf2.write(buf) # type: ignore[arg-type]
|
|
temp_files.append(temp_files.pop(0))
|
|
# remove temp files
|
|
for temp_file in temp_files:
|
|
if os.path.exists(temp_file):
|
|
os.remove(temp_file)
|
|
return schema, file_size
|
|
|
|
|
|
class TestCsvParser(AbstractTestParser):
|
|
record_types = CSV_TYPES
|
|
filetype = "csv"
|
|
|
|
@classmethod
|
|
def cases(cls) -> Mapping[str, Any]:
|
|
return {
|
|
"basic_normal_test": {
|
|
"AbstractFileParser": CsvParser(
|
|
format={"filetype": "csv"},
|
|
master_schema={
|
|
"id": "integer",
|
|
"name": "string",
|
|
"valid": "boolean",
|
|
"code": "integer",
|
|
"degrees": "number",
|
|
"birthday": "string",
|
|
"last_seen": "string",
|
|
},
|
|
),
|
|
"filepath": os.path.join(SAMPLE_DIRECTORY, "csv/test_file_1.csv"),
|
|
"num_records": 8,
|
|
"inferred_schema": {
|
|
"id": "integer",
|
|
"name": "string",
|
|
"valid": "boolean",
|
|
"code": "integer",
|
|
"degrees": "number",
|
|
"birthday": "string",
|
|
"last_seen": "string",
|
|
},
|
|
"line_checks": {},
|
|
"fails": [],
|
|
},
|
|
"custom_csv_parameters": {
|
|
# tests custom CSV parameters (odd delimiter, quote_char, escape_char & newlines in values in the file)
|
|
"AbstractFileParser": CsvParser(
|
|
format={"filetype": "csv", "delimiter": "^", "quote_char": "|", "escape_char": "!", "newlines_in_values": True},
|
|
master_schema={
|
|
"id": "integer",
|
|
"name": "string",
|
|
"valid": "boolean",
|
|
"code": "integer",
|
|
"degrees": "number",
|
|
"birthday": "string",
|
|
"last_seen": "string",
|
|
},
|
|
),
|
|
"filepath": os.path.join(SAMPLE_DIRECTORY, "csv/test_file_2_params.csv"),
|
|
"num_records": 8,
|
|
"inferred_schema": {
|
|
"id": "integer",
|
|
"name": "string",
|
|
"valid": "boolean",
|
|
"code": "integer",
|
|
"degrees": "number",
|
|
"birthday": "string",
|
|
"last_seen": "string",
|
|
},
|
|
"line_checks": {},
|
|
"fails": [],
|
|
},
|
|
"encoding_Big5": {
|
|
# tests encoding: Big5
|
|
"AbstractFileParser": CsvParser(
|
|
format={"filetype": "csv", "encoding": "big5"}, master_schema={"id": "integer", "name": "string", "valid": "boolean"}
|
|
),
|
|
"filepath": os.path.join(SAMPLE_DIRECTORY, "csv/test_file_3_enc_Big5.csv"),
|
|
"num_records": 8,
|
|
"inferred_schema": {"id": "integer", "name": "string", "valid": "boolean"},
|
|
"line_checks": {
|
|
3: {
|
|
"id": 3,
|
|
"name": "變形金剛，偽裝的機器人",
|
|
"valid": False,
|
|
}
|
|
},
|
|
"fails": [],
|
|
},
|
|
"encoding_Arabic_(Windows 1256)": {
|
|
# tests encoding: Arabic (Windows 1256)
|
|
"AbstractFileParser": CsvParser(
|
|
format={"filetype": "csv", "encoding": "windows-1256"},
|
|
master_schema={"id": "integer", "notes": "string", "valid": "boolean"},
|
|
),
|
|
"filepath": os.path.join(SAMPLE_DIRECTORY, "csv/test_file_4_enc_Arabic.csv"),
|
|
"num_records": 2,
|
|
"inferred_schema": {"id": "integer", "notes": "string", "valid": "boolean"},
|
|
"line_checks": {
|
|
1: {
|
|
"id": 1,
|
|
"notes": "البايت الجوي هو الأفضل",
|
|
"valid": False,
|
|
}
|
|
},
|
|
"fails": [],
|
|
},
|
|
"compression_gzip": {
|
|
# tests compression: gzip
|
|
"AbstractFileParser": CsvParser(
|
|
format={"filetype": "csv"},
|
|
master_schema={
|
|
"id": "integer",
|
|
"name": "string",
|
|
"valid": "boolean",
|
|
"code": "integer",
|
|
"degrees": "number",
|
|
"birthday": "string",
|
|
"last_seen": "string",
|
|
},
|
|
),
|
|
"filepath": os.path.join(SAMPLE_DIRECTORY, "csv/test_file_5.csv.gz"),
|
|
"num_records": 8,
|
|
"inferred_schema": {
|
|
"id": "integer",
|
|
"name": "string",
|
|
"valid": "boolean",
|
|
"code": "integer",
|
|
"degrees": "number",
|
|
"birthday": "string",
|
|
"last_seen": "string",
|
|
},
|
|
"line_checks": {
|
|
7: {
|
|
"id": 7,
|
|
"name": "xZhh1Kyl",
|
|
"valid": False,
|
|
"code": 10,
|
|
"degrees": -9.2,
|
|
"birthday": "2021-07-14",
|
|
"last_seen": "2021-07-14 15:30:09.225145",
|
|
}
|
|
},
|
|
"fails": [],
|
|
},
|
|
"compression_bz2": {
|
|
# tests compression: bz2
|
|
"AbstractFileParser": CsvParser(
|
|
format={"filetype": "csv"},
|
|
master_schema={
|
|
"id": "integer",
|
|
"name": "string",
|
|
"valid": "boolean",
|
|
"code": "integer",
|
|
"degrees": "number",
|
|
"birthday": "string",
|
|
"last_seen": "string",
|
|
},
|
|
),
|
|
"filepath": os.path.join(SAMPLE_DIRECTORY, "csv/test_file_7_bz2.csv.bz2"),
|
|
"num_records": 8,
|
|
"inferred_schema": {
|
|
"id": "integer",
|
|
"name": "string",
|
|
"valid": "boolean",
|
|
"code": "integer",
|
|
"degrees": "number",
|
|
"birthday": "string",
|
|
"last_seen": "string",
|
|
},
|
|
"line_checks": {
|
|
7: {
|
|
"id": 7,
|
|
"name": "xZhh1Kyl",
|
|
"valid": False,
|
|
"code": 10,
|
|
"degrees": -9.2,
|
|
"birthday": "2021-07-14",
|
|
"last_seen": "2021-07-14 15:30:09.225145",
|
|
}
|
|
},
|
|
"fails": [],
|
|
},
|
|
"extra_columns_in_master_schema": {
|
|
# tests extra columns in master schema
|
|
"AbstractFileParser": CsvParser(
|
|
format={"filetype": "csv"},
|
|
master_schema={
|
|
"EXTRA_COLUMN_1": "boolean",
|
|
"EXTRA_COLUMN_2": "number",
|
|
"id": "integer",
|
|
"name": "string",
|
|
"valid": "boolean",
|
|
"code": "integer",
|
|
"degrees": "number",
|
|
"birthday": "string",
|
|
"last_seen": "string",
|
|
},
|
|
),
|
|
"filepath": os.path.join(SAMPLE_DIRECTORY, "csv/test_file_1.csv"),
|
|
"num_records": 8,
|
|
"inferred_schema": {
|
|
"id": "integer",
|
|
"name": "string",
|
|
"valid": "boolean",
|
|
"code": "integer",
|
|
"degrees": "number",
|
|
"birthday": "string",
|
|
"last_seen": "string",
|
|
},
|
|
"line_checks": {},
|
|
"fails": [],
|
|
},
|
|
"missing_columns_in_master_schema": {
|
|
# tests missing columns in master schema
|
|
# TODO: maybe this should fail read_records, but it does pick up all the columns from file despite missing from master schema
|
|
"AbstractFileParser": CsvParser(format={"filetype": "csv"}, master_schema={"id": "integer", "name": "string"}),
|
|
"filepath": os.path.join(SAMPLE_DIRECTORY, "csv/test_file_1.csv"),
|
|
"num_records": 8,
|
|
"inferred_schema": {
|
|
"id": "integer",
|
|
"name": "string",
|
|
"valid": "boolean",
|
|
"code": "integer",
|
|
"degrees": "number",
|
|
"birthday": "string",
|
|
"last_seen": "string",
|
|
},
|
|
"line_checks": {},
|
|
"fails": [],
|
|
},
|
|
"empty_csv_file": {
|
|
# tests empty file, SHOULD FAIL INFER & STREAM RECORDS
|
|
"AbstractFileParser": CsvParser(format={"filetype": "csv"}, master_schema={}),
|
|
"filepath": os.path.join(SAMPLE_DIRECTORY, "csv/test_file_6_empty.csv"),
|
|
"num_records": 0,
|
|
"inferred_schema": {},
|
|
"line_checks": {},
|
|
"fails": ["test_get_inferred_schema", "test_stream_records"],
|
|
},
|
|
"empty_advanced_options": {
|
|
"AbstractFileParser": CsvParser(
|
|
format={"filetype": "csv", "advanced_options": ""},
|
|
master_schema={
|
|
"id": "integer",
|
|
"name": "string",
|
|
"valid": "boolean",
|
|
"code": "integer",
|
|
"degrees": "number",
|
|
"birthday": "string",
|
|
"last_seen": "string",
|
|
},
|
|
),
|
|
"filepath": os.path.join(SAMPLE_DIRECTORY, "csv/test_file_1.csv"),
|
|
"num_records": 8,
|
|
"inferred_schema": {
|
|
"id": "integer",
|
|
"name": "string",
|
|
"valid": "boolean",
|
|
"code": "integer",
|
|
"degrees": "number",
|
|
"birthday": "string",
|
|
"last_seen": "string",
|
|
},
|
|
"line_checks": {},
|
|
"fails": [],
|
|
},
|
|
"no_header_csv_file": {
|
|
# no header test
|
|
"AbstractFileParser": CsvParser(
|
|
format={
|
|
"filetype": "csv",
|
|
"advanced_options": json.dumps(
|
|
{"column_names": ["id", "name", "valid", "code", "degrees", "birthday", "last_seen"]}
|
|
),
|
|
},
|
|
master_schema={},
|
|
),
|
|
"filepath": os.path.join(SAMPLE_DIRECTORY, "csv/test_file_8_no_header.csv"),
|
|
"num_records": 8,
|
|
"inferred_schema": {
|
|
"id": "integer",
|
|
"name": "string",
|
|
"valid": "boolean",
|
|
"code": "integer",
|
|
"degrees": "number",
|
|
"birthday": "string",
|
|
"last_seen": "string",
|
|
},
|
|
"line_checks": {},
|
|
"fails": [],
|
|
},
|
|
}
|
|
|
|
@memory_limit(20)
|
|
@pytest.mark.order(1)
|
|
def test_big_file(self) -> None:
|
|
"""tests a big csv file (>= 1.5G records)"""
|
|
filepath = os.path.join(TMP_FOLDER, "big_csv_file." + self.filetype)
|
|
schema, file_size = generate_big_file(filepath, 0.1, 123)
|
|
expected_count = sum(1 for _ in open(filepath)) - 1
|
|
self.logger.info(f"generated file {filepath} with size {file_size}Gb, lines: {expected_count}")
|
|
for _ in range(3):
|
|
parser = CsvParser(
|
|
format={"filetype": self.filetype, "block_size": 5 * 1024**2},
|
|
master_schema=schema,
|
|
)
|
|
expected_file = open(filepath, "r")
|
|
# skip the first header line
|
|
next(expected_file)
|
|
read_count = 0
|
|
with smart_open(filepath, self._get_readmode({"AbstractFileParser": parser})) as f:
|
|
for record in parser.stream_records(f, FileInfo(key=filepath, size=file_size, last_modified=pendulum.now())):
|
|
record_line = ",".join("" if v is None else str(v) for v in record.values())
|
|
expected_line = next(expected_file).strip("\n")
|
|
assert record_line == expected_line
|
|
read_count += 1
|
|
assert read_count == expected_count
|
|
expected_file.close()
|