preset-io / elasticsearch-dbapi

@@ -38,7 +38,7 @@
Loading
38 38
39 39
class Connection(BaseConnection):
40 40
41 -
    """Connection to an ES Cluster """
41 +
    """Connection to an ES Cluster"""
42 42
43 43
    def __init__(
44 44
        self,

@@ -42,7 +42,7 @@
Loading
42 42
43 43
class Connection(BaseConnection):
44 44
45 -
    """Connection to an ES Cluster """
45 +
    """Connection to an ES Cluster"""
46 46
47 47
    def __init__(
48 48
        self,

@@ -109,7 +109,7 @@
Loading
109 109
110 110
class BaseConnection(object):
111 111
112 -
    """Connection to an ES Cluster """
112 +
    """Connection to an ES Cluster"""
113 113
114 114
    def __init__(
115 115
        self,
@@ -192,6 +192,7 @@
Loading
192 192
        self.es = es
193 193
        self.sql_path = kwargs.get("sql_path", DEFAULT_SQL_PATH)
194 194
        self.fetch_size = kwargs.get("fetch_size", DEFAULT_FETCH_SIZE)
195 +
        self.time_zone: Optional[str] = kwargs.get("time_zone")
195 196
        # This read/write attribute specifies the number of rows to fetch at a
196 197
        # time with .fetchmany(). It defaults to 1 meaning to fetch a single
197 198
        # row at a time.
@@ -218,7 +219,7 @@
Loading
218 219
    @check_result
219 220
    @check_closed
220 221
    def rowcount(self) -> int:
221 -
        """ Counts the number of rows on a result """
222 +
        """Counts the number of rows on a result"""
222 223
        if self._results:
223 224
            return len(self._results)
224 225
        return 0
@@ -230,7 +231,7 @@
Loading
230 231
231 232
    @check_closed
232 233
    def execute(self, operation, parameters=None) -> "BaseCursor":
233 -
        """ Children must implement their own custom execute """
234 +
        """Children must implement their own custom execute"""
234 235
        raise NotImplementedError  # pragma: no cover
235 236
236 237
    @check_closed
@@ -311,11 +312,13 @@
Loading
311 312
        payload = {"query": query}
312 313
        if self.fetch_size is not None:
313 314
            payload["fetch_size"] = self.fetch_size
315 +
        if self.time_zone is not None:
316 +
            payload["time_zone"] = self.time_zone
314 317
        path = f"/{self.sql_path}/"
315 318
        try:
316 319
            response = self.es.transport.perform_request("POST", path, body=payload)
317 320
        except es_exceptions.ConnectionError:
318 -
            raise exceptions.OperationalError(f"Error connecting to Elasticsearch")
321 +
            raise exceptions.OperationalError("Error connecting to Elasticsearch")
319 322
        except es_exceptions.RequestError as ex:
320 323
            raise exceptions.ProgrammingError(
321 324
                f"Error ({ex.error}): {ex.info['error']['reason']}"

@@ -7,28 +7,35 @@
Loading
7 7
from es.opendistro.api import connect as open_connect
8 8
9 9
10 +
def convert_bool(value: str) -> bool:
11 +
    return True if value == "True" else False
12 +
13 +
10 14
class TestDBAPI(unittest.TestCase):
11 15
    def setUp(self):
12 16
        self.driver_name = os.environ.get("ES_DRIVER", "elasticsearch")
13 -
        host = os.environ.get("ES_HOST", "localhost")
14 -
        port = int(os.environ.get("ES_PORT", 9200))
15 -
        scheme = os.environ.get("ES_SCHEME", "http")
16 -
        verify_certs = os.environ.get("ES_VERIFY_CERTS", False)
17 -
        user = os.environ.get("ES_USER", None)
18 -
        password = os.environ.get("ES_PASSWORD", None)
17 +
        self.host = os.environ.get("ES_HOST", "localhost")
18 +
        self.port = int(os.environ.get("ES_PORT", 9200))
19 +
        self.scheme = os.environ.get("ES_SCHEME", "http")
20 +
        self.verify_certs = os.environ.get("ES_VERIFY_CERTS", False)
21 +
        self.user = os.environ.get("ES_USER", None)
22 +
        self.password = os.environ.get("ES_PASSWORD", None)
19 23
        self.v2 = bool(os.environ.get("ES_V2", False))
24 +
        self.support_datetime_parse = convert_bool(
25 +
            os.environ.get("ES_SUPPORT_DATETIME_PARSE", "True")
26 +
        )
20 27
21 28
        if self.driver_name == "elasticsearch":
22 29
            self.connect_func = elastic_connect
23 30
        else:
24 31
            self.connect_func = open_connect
25 32
        self.conn = self.connect_func(
26 -
            host=host,
27 -
            port=port,
28 -
            scheme=scheme,
29 -
            verify_certs=verify_certs,
30 -
            user=user,
31 -
            password=password,
33 +
            host=self.host,
34 +
            port=self.port,
35 +
            scheme=self.scheme,
36 +
            verify_certs=self.verify_certs,
37 +
            user=self.user,
38 +
            password=self.password,
32 39
            v2=self.v2,
33 40
        )
34 41
        self.cursor = self.conn.cursor()
@@ -213,3 +220,62 @@
Loading
213 220
        mock_elasticsearch.assert_called_once_with(
214 221
            "https://localhost:9200/", http_auth=("user", "password")
215 222
        )
223 +
224 +
    def test_simple_search_with_time_zone(self):
225 +
        """
226 +
        DBAPI: Test simple search with time zone
227 +
        UTC -> CST
228 +
        2019-10-13T00:00:00.000Z => 2019-10-13T08:00:00.000+08:00
229 +
        2019-10-13T00:00:01.000Z => 2019-10-13T08:01:00.000+08:00
230 +
        2019-10-13T00:00:02.000Z => 2019-10-13T08:02:00.000+08:00
231 +
        """
232 +
233 +
        if not self.support_datetime_parse:
234 +
            return
235 +
236 +
        conn = self.connect_func(
237 +
            host=self.host,
238 +
            port=self.port,
239 +
            scheme=self.scheme,
240 +
            verify_certs=self.verify_certs,
241 +
            user=self.user,
242 +
            password=self.password,
243 +
            v2=self.v2,
244 +
            time_zone="Asia/Shanghai",
245 +
        )
246 +
        cursor = conn.cursor()
247 +
        pattern = "yyyy-MM-dd HH:mm:ss"
248 +
        sql = f"""
249 +
        SELECT timestamp FROM data1
250 +
        WHERE timestamp >= DATETIME_PARSE('2019-10-13 00:08:00', '{pattern}')
251 +
        """
252 +
253 +
        rows = cursor.execute(sql).fetchall()
254 +
        self.assertEqual(len(rows), 3)
255 +
256 +
    def test_simple_search_without_time_zone(self):
257 +
        """
258 +
        DBAPI: Test simple search without time zone
259 +
        """
260 +
261 +
        if not self.support_datetime_parse:
262 +
            return
263 +
264 +
        conn = self.connect_func(
265 +
            host=self.host,
266 +
            port=self.port,
267 +
            scheme=self.scheme,
268 +
            verify_certs=self.verify_certs,
269 +
            user=self.user,
270 +
            password=self.password,
271 +
            v2=self.v2,
272 +
        )
273 +
        cursor = conn.cursor()
274 +
        pattern = "yyyy-MM-dd HH:mm:ss"
275 +
        sql = f"""
276 +
        SELECT * FROM data1
277 +
        WHERE timestamp >= DATETIME_PARSE('2019-10-13 08:00:00', '{pattern}')
278 +
        """
279 +
280 +
        rows = cursor.execute(sql).fetchall()
281 +
        self.assertEqual(len(rows), 0)
Files Coverage
es 94.67%
Project Totals (15 files) 94.67%

No yaml found.

Create your codecov.yml to customize your Codecov experience

Sunburst
The inner-most circle is the entire project, moving away from the center are folders then, finally, a single file. The size and color of each slice is representing the number of statements and the coverage, respectively.
Icicle
The top section represents the entire project. Proceeding with folders and finally individual files. The size and color of each slice is representing the number of statements and the coverage, respectively.
Grid
Each block represents a single file in the project. The size and color of each block is represented by the number of statements and the coverage, respectively.
Loading