MySqlLoaderWriter.write_bool()   A
last analyzed

Complexity

Conditions 2

Size

Total Lines 12
Code Lines 5

Duplication

Lines 0
Ratio 0 %

Code Coverage

Tests 4
CRAP Score 2

Importance

Changes 0
Metric Value
eloc 5
dl 0
loc 12
ccs 4
cts 4
cp 1
rs 10
c 0
b 0
f 0
cc 2
nop 2
crap 2
1 1
import datetime
2 1
import os
3 1
from abc import ABC
4 1
from decimal import Decimal
5 1
from typing import Any, Dict, Optional
6
from uuid import UUID
7 1
8 1
import pytz
9
from etlt.writer.SqlLoaderWriter import SqlLoaderWriter
10
11 1
12
class MySqlLoaderWriter(SqlLoaderWriter, ABC):
13
    """
14
    Writer for storing rows in CSV format optimized for MariaDB and MySQL instances and 'load data local infile'
15
    statement.
16
    """
17
18 1
    # ------------------------------------------------------------------------------------------------------------------
19 1
    @staticmethod
20
    def write_bool(value: bool, file: Any) -> None:
21
        """
22
        Writes a boolean as a field to a CSV file.
23
24
        :param value: The boolean.
25
        :param file: The file like object
26 1
        """
27 1
        if value:
28
            file.write('1')
29 1
        else:
30
            file.write('0')
31
32 1
    # ------------------------------------------------------------------------------------------------------------------
33 1
    @staticmethod
34
    def write_date(value: datetime.date, file: Any) -> None:
35
        """
36
        Writes a date as a field to a CSV file.
37
38
        :param value: The date.
39
        :param file: The file like object
40 1
        """
41
        file.write(value.isoformat())
42
43 1
    # ------------------------------------------------------------------------------------------------------------------
44 1
    @staticmethod
45
    def write_datetime(value: datetime.datetime, file: Any) -> None:
46
        """
47
        Writes a datetime as a field to a CSV file.
48
49
        :param value: The date.
50
        :param file: The file like object
51 1
        """
52 1
        zone = value.tzinfo
53
        if zone and zone != pytz.utc:
54 1
            raise ValueError('Only native and UTC timezone supported')
55
        file.write(value.isoformat())
56
57 1
    # ------------------------------------------------------------------------------------------------------------------
58 1
    @staticmethod
59
    def write_timedelta(value: datetime.timedelta, file: Any) -> None:
60
        """
61
        Writes a timedelta as a field to a CSV file.
62
63
        :param value: The timedelta.
64
        :param file: The file like object
65 1
        """
66
        MySqlLoaderWriter.write_string(str(value), file)
67
68 1
    # ------------------------------------------------------------------------------------------------------------------
69 1
    @staticmethod
70
    def write_decimal(value: Decimal, file: Any) -> None:
71
        """
72
        Writes a decimal as a field to a CSV file.
73
74
        :param value: The decimal.
75
        :param file: The file like object
76 1
        """
77
        file.write(str(value))
78
79 1
    # ------------------------------------------------------------------------------------------------------------------
80 1
    @staticmethod
81
    def write_float(value: float, file: Any) -> None:
82
        """
83
        Writes a float as a field to a CSV file.
84
85
        :param value: The float.
86
        :param file: The file like object
87 1
        """
88
        file.write(str(value))
89
90 1
    # ------------------------------------------------------------------------------------------------------------------
91 1
    @staticmethod
92
    def write_int(value: int, file: Any) -> None:
93
        """
94
        Writes an integer as a field to a CSV file.
95
96
        :param value: The integer.
97
        :param file: The file like object
98 1
        """
99
        file.write(str(value))
100
101 1
    # ------------------------------------------------------------------------------------------------------------------
102 1
    @staticmethod
103
    def write_none(_, file: Any) -> None:
104
        """
105
        Writes None as a field to a CSV file.
106
107
        :param _: The None object.
108
        :param file: The file like object
109 1
        """
110
        file.write('\\N')
111
112 1
    # ------------------------------------------------------------------------------------------------------------------
113 1
    @staticmethod
114
    def write_string(value: str, file: Any) -> None:
115
        """
116
        Writes a string as a field to a CSV file.
117
118
        :param value: The string.
119
        :param file: The file.
120 1
        """
121 1
        if value == '':
122
            file.write('\\N')
123 1
        else:
124 1
            file.write("'")
125 1
            file.write(value.replace('\\', '\\\\').replace("'", "\\'"))
126
            file.write("'")
127
128 1
    # ------------------------------------------------------------------------------------------------------------------
129
    def writerow(self, row: Dict[str, Any]) -> None:
130
        """
131
        Writes a row to the destination file.
132
133
        :param row: The row.
134 1
        """
135 1
        for field in self._fields:
136 1
            self._write_field(row[field])
137
            self._file.write(',')
138 1
139
        self._file.write(os.linesep)
140
141 1
    # ------------------------------------------------------------------------------------------------------------------
142 1
    @staticmethod
143
    def write_uuid(value: UUID, file: Any) -> None:
144
        """
145
        Writes a UUID as a field to a CSV file.
146
147
        :param value: The UUID.
148
        :param file: The file like object
149 1
        """
150
        MySqlLoaderWriter.write_string(str(value), file)
151
152 1
    # ------------------------------------------------------------------------------------------------------------------
153
    def get_bulk_load_sql(self, table_name: str, partition: Optional[str] = None) -> str:
154
        """
155
        Returns a SQL statement for bulk loading the data into a table.
156
157
        :param table_name: The name of the table.
158
        :param partition: When applicable, the name of the partition in which the data must be loaded.`
159
        """
160
        sql = "load data local infile '{FILENAME}'"
161
        sql += ' into table `{TABLE_NAME}`'
162
        if partition is not None:
163
            sql += ' partition ({PARTITION})'
164
        sql += ' character set {ENCODING}'
165
        sql += " fields terminated by ','"
166
        sql += " optionally enclosed by '\\\''"
167
        sql += " escaped by '\\\\'"
168
        sql += " lines terminated by '\\n'"
169
170
        return sql.format(FILENAME=self._filename,  # @todo Quoting of filename
171
                          ENCODING=self._encoding, TABLE_NAME=table_name, PARTITION=partition)
172
173
174 1
# ----------------------------------------------------------------------------------------------------------------------
175 1
MySqlLoaderWriter.register_handler("<class 'bool'>", MySqlLoaderWriter.write_bool)
176 1
MySqlLoaderWriter.register_handler("<class 'datetime.date'>", MySqlLoaderWriter.write_date)
177 1
MySqlLoaderWriter.register_handler("<class 'datetime.datetime'>", MySqlLoaderWriter.write_datetime)
178 1
MySqlLoaderWriter.register_handler("<class 'datetime.timedelta'>", MySqlLoaderWriter.write_timedelta)
179 1
MySqlLoaderWriter.register_handler("<class 'decimal.Decimal'>", MySqlLoaderWriter.write_decimal)
180 1
MySqlLoaderWriter.register_handler("<class 'float'>", MySqlLoaderWriter.write_float)
181 1
MySqlLoaderWriter.register_handler("<class 'int'>", MySqlLoaderWriter.write_int)
182 1
MySqlLoaderWriter.register_handler("<class 'NoneType'>", MySqlLoaderWriter.write_none)
183 1
MySqlLoaderWriter.register_handler("<class 'str'>", MySqlLoaderWriter.write_string)
184
MySqlLoaderWriter.register_handler("<class 'uuid.UUID'>", MySqlLoaderWriter.write_uuid)
185