mapepire-python 0.1.4__tar.gz → 0.1.5__tar.gz

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (23) hide show
  1. {mapepire_python-0.1.4 → mapepire_python-0.1.5}/PKG-INFO +195 -18
  2. {mapepire_python-0.1.4 → mapepire_python-0.1.5}/README.md +191 -17
  3. mapepire_python-0.1.5/mapepire_python/base_job.py +42 -0
  4. {mapepire_python-0.1.4 → mapepire_python-0.1.5}/mapepire_python/client/query.py +31 -17
  5. {mapepire_python-0.1.4 → mapepire_python-0.1.5}/mapepire_python/client/sql_job.py +80 -30
  6. mapepire_python-0.1.5/mapepire_python/client/websocket.py +37 -0
  7. mapepire_python-0.1.4/mapepire_python/types.py → mapepire_python-0.1.5/mapepire_python/data_types.py +8 -2
  8. mapepire_python-0.1.5/mapepire_python/pool/pool_client.py +141 -0
  9. mapepire_python-0.1.5/mapepire_python/pool/pool_job.py +303 -0
  10. mapepire_python-0.1.5/mapepire_python/pool/pool_query.py +136 -0
  11. mapepire_python-0.1.5/mapepire_python/query_manager.py +68 -0
  12. {mapepire_python-0.1.4 → mapepire_python-0.1.5}/mapepire_python/ssl.py +1 -1
  13. {mapepire_python-0.1.4 → mapepire_python-0.1.5}/mapepire_python/version.py +1 -1
  14. {mapepire_python-0.1.4 → mapepire_python-0.1.5}/mapepire_python.egg-info/PKG-INFO +195 -18
  15. {mapepire_python-0.1.4 → mapepire_python-0.1.5}/mapepire_python.egg-info/SOURCES.txt +8 -3
  16. {mapepire_python-0.1.4 → mapepire_python-0.1.5}/mapepire_python.egg-info/requires.txt +3 -0
  17. {mapepire_python-0.1.4 → mapepire_python-0.1.5}/pyproject.toml +3 -0
  18. mapepire_python-0.1.4/mapepire_python/pool.py +0 -29
  19. {mapepire_python-0.1.4 → mapepire_python-0.1.5}/LICENSE +0 -0
  20. {mapepire_python-0.1.4 → mapepire_python-0.1.5}/mapepire_python/py.typed +0 -0
  21. {mapepire_python-0.1.4 → mapepire_python-0.1.5}/mapepire_python.egg-info/dependency_links.txt +0 -0
  22. {mapepire_python-0.1.4 → mapepire_python-0.1.5}/mapepire_python.egg-info/top_level.txt +0 -0
  23. {mapepire_python-0.1.4 → mapepire_python-0.1.5}/setup.cfg +0 -0
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.1
2
2
  Name: mapepire-python
3
- Version: 0.1.4
3
+ Version: 0.1.5
4
4
  Author-email: Allen Institute for Artificial Intelligence <contact@allenai.org>
5
5
  License: Apache License
6
6
  Version 2.0, January 2004
@@ -217,12 +217,15 @@ Description-Content-Type: text/markdown
217
217
  License-File: LICENSE
218
218
  Requires-Dist: dataclasses-json>=0.6.4
219
219
  Requires-Dist: websocket-client>=1.2.1
220
+ Requires-Dist: websockets
221
+ Requires-Dist: pyee
220
222
  Provides-Extra: dev
221
223
  Requires-Dist: ruff; extra == "dev"
222
224
  Requires-Dist: mypy<1.5,>=1.0; extra == "dev"
223
225
  Requires-Dist: black<24.0,>=23.0; extra == "dev"
224
226
  Requires-Dist: isort<5.13,>=5.12; extra == "dev"
225
227
  Requires-Dist: pytest; extra == "dev"
228
+ Requires-Dist: pytest-asyncio; extra == "dev"
226
229
  Requires-Dist: pytest-sphinx; extra == "dev"
227
230
  Requires-Dist: pytest-cov; extra == "dev"
228
231
  Requires-Dist: twine>=1.11.0; extra == "dev"
@@ -252,24 +255,36 @@ Requires-Dist: pre-commit; extra == "dev"
252
255
  <br/>
253
256
  </div>
254
257
 
255
- - [Overview](#overview)
256
- - [Setup](#setup)
257
- - [Install with `pip`](#install-with-pip)
258
- - [Server Component Setup (Forthcoming)](#server-component-setup-forthcoming)
259
- - [Example usage](#example-usage)
258
+ - [mapepire-python](#mapepire-python)
259
+ - [Overview](#overview)
260
+ - [Setup](#setup)
261
+ - [Install with `pip`](#install-with-pip)
262
+ - [Server Component Setup](#server-component-setup)
263
+ - [Example usage](#example-usage)
264
+ - [Query and run](#query-and-run)
265
+ - [Asynchronous Query Execution](#asynchronous-query-execution)
266
+ - [Pooling (beta)](#pooling-beta)
260
267
  - [Development Setup](#development-setup)
261
268
  - [Setup python virtual environment with pip and venv](#setup-python-virtual-environment-with-pip-and-venv)
262
269
  - [Create a new virtual environment](#create-a-new-virtual-environment)
263
270
  - [Unix/macOS](#unixmacos)
264
271
  - [Windows](#windows)
265
272
  - [Activate the virtual environment:](#activate-the-virtual-environment)
273
+ - [Unix.macOS](#unixmacos-1)
274
+ - [Windows](#windows-1)
275
+ - [Unix/macOS](#unixmacos-2)
276
+ - [Windows](#windows-2)
266
277
  - [Prepare pip](#prepare-pip)
278
+ - [Unix/macOS](#unixmacos-3)
279
+ - [Windows](#windows-3)
267
280
  - [Install Dependencies using `requirements-dev.txt`](#install-dependencies-using-requirements-devtxt)
281
+ - [Unix/macOS](#unixmacos-4)
282
+ - [Windows](#windows-4)
268
283
  - [Setup Python virtual environment with Conda](#setup-python-virtual-environment-with-conda)
269
284
  - [Create an environment from an environment-dev.yml file](#create-an-environment-from-an-environment-devyml-file)
270
285
  - [1. Activate the new environment:](#1-activate-the-new-environment)
271
286
  - [2. Verify the new environment was installed:](#2-verify-the-new-environment-was-installed)
272
- - [Run local test suite](#run-local-test-suite)
287
+ - [Run local test suite](#run-local-test-suite)
273
288
 
274
289
 
275
290
 
@@ -299,26 +314,38 @@ To use mapire-python, you will need to have the Mapepire Server Component runnin
299
314
 
300
315
  ## Example usage
301
316
 
302
- The following script sets up a `DaemonServer` object that will be used to connect with the Server Component. Then a single `SQLJob` is created to facilitate the connection from the client side.
317
+ Setup the server credentials used to connect to the server. One way to do this is to create a `mapepire.ini` file in the root of your project with the following content:
318
+
319
+ ```ini
320
+ [mapepire]
321
+ SERVER="SERVER"
322
+ PORT="PORT"
323
+ USER="USER"
324
+ PASSWORD="PASSWORD"
325
+ ```
326
+
327
+ The following script sets up a `DaemonServer` object that will be used to connect with the Server Component. Then a single `SQLJob` is created to facilitate the connection from the client side.
303
328
 
304
329
  ```python
330
+ import configparser
305
331
  from mapepire_python.client.sql_job import SQLJob
306
332
  from mapepire_python.types import DaemonServer
307
333
 
334
+ config = configparser.ConfigParser()
335
+ config.read('mapepire.ini')
336
+
308
337
  creds = DaemonServer(
309
- host="localhost",
310
- port=8085,
311
- user="USER",
312
- password="PASSWORD",
338
+ host=config['mapepire']['SERVER'],
339
+ port=config['mapepire']['PORT'],
340
+ user=config['mapepire']['USER'],
341
+ password=config['mapepire']['PASSWORD'],
313
342
  ignoreUnauthorized=True
314
343
  )
315
344
 
316
-
317
- job = SQLJob()
318
- res = job.connect(creds)
319
- query = job.query("select * from sample.employee")
320
- result = query.run(rows_to_fetch=3)
321
- print(result)
345
+ with SQLJob(creds) as sql_job:
346
+ with sql_job.query("select * from sample.employee") as query:
347
+ result = query.run(rows_to_fetch=1)
348
+ print(result)
322
349
  ```
323
350
 
324
351
  Here is the output from the script above:
@@ -441,6 +468,156 @@ Here is the output from the script above:
441
468
  }
442
469
 
443
470
  ```
471
+ The results object is a JSON object that contains the metadata and data from the query. Here are the different fields returned:
472
+ - `id` field contains the query ID
473
+ - `has_results` field indicates whether the query returned any results
474
+ - `update_count` field indicates the number of rows updated by the query (-1 if the query did not update any rows)
475
+ - `metadata` field contains information about the columns returned by the query
476
+ - `data` field contains the results of the query
477
+ - `is_done` field indicates whether the query has finished executing
478
+ - `success` field indicates whether the query was successful.
479
+
480
+ In the ouput above, the query was successful and returned one row of data.
481
+
482
+ ### Query and run
483
+
484
+ To create and run a query in a single step, use the `query_and_run` method:
485
+
486
+ ```python
487
+ import configparser
488
+ from mapepire_python.client.sql_job import SQLJob
489
+ from mapepire_python.types import DaemonServer
490
+
491
+ config = configparser.ConfigParser()
492
+ config.read('mapepire.ini')
493
+
494
+ creds = DaemonServer(
495
+ host=config['mapepire']['SERVER'],
496
+ port=config['mapepire']['PORT'],
497
+ user=config['mapepire']['USER'],
498
+ password=config['mapepire']['PASSWORD'],
499
+ ignoreUnauthorized=True
500
+ )
501
+
502
+ with SQLJob(creds) as sql_job:
503
+ # query automatically closed after running
504
+ results = sql_job.query_and_run("select * from sample.employee", rows_to_fetch=1)
505
+ print(result)
506
+ ```
507
+
508
+ ### Asynchronous Query Execution
509
+
510
+ The `PoolJob` object can be used to create and run queries asynchronously:
511
+
512
+ ```python
513
+ import asyncio
514
+ import configparser
515
+ from mapepire_python.pool.pool_job import PoolJob
516
+ from mapepire_python.types import DaemonServer
517
+
518
+ config = configparser.ConfigParser()
519
+ config.read('mapepire.ini')
520
+
521
+ creds = DaemonServer(
522
+ host=config['mapepire']['SERVER'],
523
+ port=config['mapepire']['PORT'],
524
+ user=config['mapepire']['USER'],
525
+ password=config['mapepire']['PASSWORD'],
526
+ ignoreUnauthorized=True
527
+ )
528
+
529
+ async def main():
530
+ async with PoolJob(creds=creds) as pool_job:
531
+ async with pool_job.query('select * from sample.employee') as query:
532
+ res = await query.run(rows_to_fetch=1)
533
+
534
+ if __name__ == '__main__':
535
+ asyncio.run(main())
536
+
537
+ ```
538
+
539
+ To run a create and run a query asynchronously in a single step, use the `query_and_run` method:
540
+
541
+ ```python
542
+ import asyncio
543
+ import configparser
544
+ from mapepire_python.pool.pool_job import PoolJob
545
+ from mapepire_python.types import DaemonServer
546
+
547
+ config = configparser.ConfigParser()
548
+ config.read('mapepire.ini')
549
+
550
+ creds = DaemonServer(
551
+ host=config['mapepire']['SERVER'],
552
+ port=config['mapepire']['PORT'],
553
+ user=config['mapepire']['USER'],
554
+ password=config['mapepire']['PASSWORD'],
555
+ ignoreUnauthorized=True
556
+ )
557
+
558
+ async def main():
559
+ async with PoolJob(creds=creds) as pool_job:
560
+ res = await pool_job.query_and_run(rows_to_fetch=1)
561
+ print(res)
562
+
563
+ if __name__ == '__main__':
564
+ asyncio.run(main())
565
+
566
+ ```
567
+
568
+
569
+ ## Pooling (beta)
570
+
571
+ The `Pool` object can be used to create a pool of `PoolJob` objects to run queries concurrently.
572
+
573
+ ```python
574
+ import asyncio
575
+ import configparser
576
+ from mapepire_python.pool.pool_client import Pool, PoolOptions
577
+ from mapepire_python.types import DaemonServer
578
+
579
+ config = configparser.ConfigParser()
580
+ config.read('mapepire.ini')
581
+
582
+ creds = DaemonServer(
583
+ host=config['mapepire']['SERVER'],
584
+ port=config['mapepire']['PORT'],
585
+ user=config['mapepire']['USER'],
586
+ password=config['mapepire']['PASSWORD'],
587
+ ignoreUnauthorized=True
588
+ )
589
+
590
+
591
+ async def main():
592
+ async with Pool(
593
+ options=PoolOptions(
594
+ creds=creds,
595
+ opts=None,
596
+ max_size=5,
597
+ starting_size=3
598
+ )
599
+ ) as pool:
600
+ job_names = []
601
+ resultsA = await asyncio.gather(
602
+ pool.execute('values (job_name)'),
603
+ pool.execute('values (job_name)'),
604
+ pool.execute('values (job_name)')
605
+ )
606
+ job_names = [res['data'][0]['00001'] for res in resultsA]
607
+
608
+ print(job_names)
609
+
610
+
611
+ if __name__ == '__main__':
612
+ asyncio.run(main())
613
+
614
+ ```
615
+ This script will create a pool of 3 `PoolJob` objects and run the query `values (job_name)` concurrently. The results will be printed to the console.
616
+
617
+ ```bash
618
+ ['004460/QUSER/QZDASOINIT', '005096/QUSER/QZDASOINIT', '005319/QUSER/QZDASOINIT']
619
+ ```
620
+
444
621
 
445
622
  # Development Setup
446
623
 
@@ -12,24 +12,36 @@
12
12
  <br/>
13
13
  </div>
14
14
 
15
- - [Overview](#overview)
16
- - [Setup](#setup)
17
- - [Install with `pip`](#install-with-pip)
18
- - [Server Component Setup (Forthcoming)](#server-component-setup-forthcoming)
19
- - [Example usage](#example-usage)
15
+ - [mapepire-python](#mapepire-python)
16
+ - [Overview](#overview)
17
+ - [Setup](#setup)
18
+ - [Install with `pip`](#install-with-pip)
19
+ - [Server Component Setup](#server-component-setup)
20
+ - [Example usage](#example-usage)
21
+ - [Query and run](#query-and-run)
22
+ - [Asynchronous Query Execution](#asynchronous-query-execution)
23
+ - [Pooling (beta)](#pooling-beta)
20
24
  - [Development Setup](#development-setup)
21
25
  - [Setup python virtual environment with pip and venv](#setup-python-virtual-environment-with-pip-and-venv)
22
26
  - [Create a new virtual environment](#create-a-new-virtual-environment)
23
27
  - [Unix/macOS](#unixmacos)
24
28
  - [Windows](#windows)
25
29
  - [Activate the virtual environment:](#activate-the-virtual-environment)
30
+ - [Unix.macOS](#unixmacos-1)
31
+ - [Windows](#windows-1)
32
+ - [Unix/macOS](#unixmacos-2)
33
+ - [Windows](#windows-2)
26
34
  - [Prepare pip](#prepare-pip)
35
+ - [Unix/macOS](#unixmacos-3)
36
+ - [Windows](#windows-3)
27
37
  - [Install Dependencies using `requirements-dev.txt`](#install-dependencies-using-requirements-devtxt)
38
+ - [Unix/macOS](#unixmacos-4)
39
+ - [Windows](#windows-4)
28
40
  - [Setup Python virtual environment with Conda](#setup-python-virtual-environment-with-conda)
29
41
  - [Create an environment from an environment-dev.yml file](#create-an-environment-from-an-environment-devyml-file)
30
42
  - [1. Activate the new environment:](#1-activate-the-new-environment)
31
43
  - [2. Verify the new environment was installed:](#2-verify-the-new-environment-was-installed)
32
- - [Run local test suite](#run-local-test-suite)
44
+ - [Run local test suite](#run-local-test-suite)
33
45
 
34
46
 
35
47
 
@@ -59,26 +71,38 @@ To use mapire-python, you will need to have the Mapepire Server Component runnin
59
71
 
60
72
  ## Example usage
61
73
 
62
- The following script sets up a `DaemonServer` object that will be used to connect with the Server Component. Then a single `SQLJob` is created to facilitate the connection from the client side.
74
+ Setup the server credentials used to connect to the server. One way to do this is to create a `mapepire.ini` file in the root of your project with the following content:
75
+
76
+ ```ini
77
+ [mapepire]
78
+ SERVER="SERVER"
79
+ PORT="PORT"
80
+ USER="USER"
81
+ PASSWORD="PASSWORD"
82
+ ```
83
+
84
+ The following script sets up a `DaemonServer` object that will be used to connect with the Server Component. Then a single `SQLJob` is created to facilitate the connection from the client side.
63
85
 
64
86
  ```python
87
+ import configparser
65
88
  from mapepire_python.client.sql_job import SQLJob
66
89
  from mapepire_python.types import DaemonServer
67
90
 
91
+ config = configparser.ConfigParser()
92
+ config.read('mapepire.ini')
93
+
68
94
  creds = DaemonServer(
69
- host="localhost",
70
- port=8085,
71
- user="USER",
72
- password="PASSWORD",
95
+ host=config['mapepire']['SERVER'],
96
+ port=config['mapepire']['PORT'],
97
+ user=config['mapepire']['USER'],
98
+ password=config['mapepire']['PASSWORD'],
73
99
  ignoreUnauthorized=True
74
100
  )
75
101
 
76
-
77
- job = SQLJob()
78
- res = job.connect(creds)
79
- query = job.query("select * from sample.employee")
80
- result = query.run(rows_to_fetch=3)
81
- print(result)
102
+ with SQLJob(creds) as sql_job:
103
+ with sql_job.query("select * from sample.employee") as query:
104
+ result = query.run(rows_to_fetch=1)
105
+ print(result)
82
106
  ```
83
107
 
84
108
  Here is the output from the script above:
@@ -201,6 +225,156 @@ Here is the output from the script above:
201
225
  }
202
226
 
203
227
  ```
228
+ The results object is a JSON object that contains the metadata and data from the query. Here are the different fields returned:
229
+ - `id` field contains the query ID
230
+ - `has_results` field indicates whether the query returned any results
231
+ - `update_count` field indicates the number of rows updated by the query (-1 if the query did not update any rows)
232
+ - `metadata` field contains information about the columns returned by the query
233
+ - `data` field contains the results of the query
234
+ - `is_done` field indicates whether the query has finished executing
235
+ - `success` field indicates whether the query was successful.
236
+
237
+ In the ouput above, the query was successful and returned one row of data.
238
+
239
+ ### Query and run
240
+
241
+ To create and run a query in a single step, use the `query_and_run` method:
242
+
243
+ ```python
244
+ import configparser
245
+ from mapepire_python.client.sql_job import SQLJob
246
+ from mapepire_python.types import DaemonServer
247
+
248
+ config = configparser.ConfigParser()
249
+ config.read('mapepire.ini')
250
+
251
+ creds = DaemonServer(
252
+ host=config['mapepire']['SERVER'],
253
+ port=config['mapepire']['PORT'],
254
+ user=config['mapepire']['USER'],
255
+ password=config['mapepire']['PASSWORD'],
256
+ ignoreUnauthorized=True
257
+ )
258
+
259
+ with SQLJob(creds) as sql_job:
260
+ # query automatically closed after running
261
+ results = sql_job.query_and_run("select * from sample.employee", rows_to_fetch=1)
262
+ print(result)
263
+ ```
264
+
265
+ ### Asynchronous Query Execution
266
+
267
+ The `PoolJob` object can be used to create and run queries asynchronously:
268
+
269
+ ```python
270
+ import asyncio
271
+ import configparser
272
+ from mapepire_python.pool.pool_job import PoolJob
273
+ from mapepire_python.types import DaemonServer
274
+
275
+ config = configparser.ConfigParser()
276
+ config.read('mapepire.ini')
277
+
278
+ creds = DaemonServer(
279
+ host=config['mapepire']['SERVER'],
280
+ port=config['mapepire']['PORT'],
281
+ user=config['mapepire']['USER'],
282
+ password=config['mapepire']['PASSWORD'],
283
+ ignoreUnauthorized=True
284
+ )
285
+
286
+ async def main():
287
+ async with PoolJob(creds=creds) as pool_job:
288
+ async with pool_job.query('select * from sample.employee') as query:
289
+ res = await query.run(rows_to_fetch=1)
290
+
291
+ if __name__ == '__main__':
292
+ asyncio.run(main())
293
+
294
+ ```
295
+
296
+ To run a create and run a query asynchronously in a single step, use the `query_and_run` method:
297
+
298
+ ```python
299
+ import asyncio
300
+ import configparser
301
+ from mapepire_python.pool.pool_job import PoolJob
302
+ from mapepire_python.types import DaemonServer
303
+
304
+ config = configparser.ConfigParser()
305
+ config.read('mapepire.ini')
306
+
307
+ creds = DaemonServer(
308
+ host=config['mapepire']['SERVER'],
309
+ port=config['mapepire']['PORT'],
310
+ user=config['mapepire']['USER'],
311
+ password=config['mapepire']['PASSWORD'],
312
+ ignoreUnauthorized=True
313
+ )
314
+
315
+ async def main():
316
+ async with PoolJob(creds=creds) as pool_job:
317
+ res = await pool_job.query_and_run(rows_to_fetch=1)
318
+ print(res)
319
+
320
+ if __name__ == '__main__':
321
+ asyncio.run(main())
322
+
323
+ ```
324
+
325
+
326
+ ## Pooling (beta)
327
+
328
+ The `Pool` object can be used to create a pool of `PoolJob` objects to run queries concurrently.
329
+
330
+ ```python
331
+ import asyncio
332
+ import configparser
333
+ from mapepire_python.pool.pool_client import Pool, PoolOptions
334
+ from mapepire_python.types import DaemonServer
335
+
336
+ config = configparser.ConfigParser()
337
+ config.read('mapepire.ini')
338
+
339
+ creds = DaemonServer(
340
+ host=config['mapepire']['SERVER'],
341
+ port=config['mapepire']['PORT'],
342
+ user=config['mapepire']['USER'],
343
+ password=config['mapepire']['PASSWORD'],
344
+ ignoreUnauthorized=True
345
+ )
346
+
347
+
348
+ async def main():
349
+ async with Pool(
350
+ options=PoolOptions(
351
+ creds=creds,
352
+ opts=None,
353
+ max_size=5,
354
+ starting_size=3
355
+ )
356
+ ) as pool:
357
+ job_names = []
358
+ resultsA = await asyncio.gather(
359
+ pool.execute('values (job_name)'),
360
+ pool.execute('values (job_name)'),
361
+ pool.execute('values (job_name)')
362
+ )
363
+ job_names = [res['data'][0]['00001'] for res in resultsA]
364
+
365
+ print(job_names)
366
+
367
+
368
+ if __name__ == '__main__':
369
+ asyncio.run(main())
370
+
371
+ ```
372
+ This script will create a pool of 3 `PoolJob` objects and run the query `values (job_name)` concurrently. The results will be printed to the console.
373
+
374
+ ```bash
375
+ ['004460/QUSER/QZDASOINIT', '005096/QUSER/QZDASOINIT', '005319/QUSER/QZDASOINIT']
376
+ ```
377
+
204
378
 
205
379
  # Development Setup
206
380
 
@@ -0,0 +1,42 @@
1
+ from typing import Any, Dict, Optional, Union
2
+
3
+ from .data_types import DaemonServer, JobStatus, QueryOptions
4
+
5
+
6
+ class BaseJob:
7
+ def __init__(self, creds: DaemonServer = None, options: Dict[Any, Any] = {}) -> None:
8
+ self.creds = creds
9
+ self.options = options
10
+
11
+ def connect(self, db2_server: Union[DaemonServer, Dict[str, Any]]) -> Dict[str, Any]:
12
+ raise NotImplementedError()
13
+
14
+ def close(self) -> None:
15
+ raise NotImplementedError()
16
+
17
+ def get_status(self) -> JobStatus:
18
+ raise NotImplementedError()
19
+
20
+ def query(
21
+ self,
22
+ sql: str,
23
+ opts: Optional[Union[Dict[str, Any], QueryOptions]] = None,
24
+ ):
25
+ raise NotImplementedError()
26
+
27
+ def query_and_run(
28
+ self, sql: str, opts: Optional[Dict[str, Any]] = None, **kwargs
29
+ ) -> Dict[str, Any]:
30
+ raise NotImplementedError()
31
+
32
+ def __enter__(self):
33
+ raise NotImplementedError()
34
+
35
+ def __exit__(self, *args, **kwargs):
36
+ raise NotImplementedError()
37
+
38
+ async def __aenter__(self):
39
+ raise NotImplementedError()
40
+
41
+ async def __aexit__(self, *args, **kwargs):
42
+ raise NotImplementedError()
@@ -1,8 +1,8 @@
1
1
  import json
2
2
  from enum import Enum
3
- from typing import Any, Dict, Generic, List, Optional, TypeVar, Union
3
+ from typing import Any, Dict, Generic, List, Optional, TypeVar
4
4
 
5
- from ..types import QueryOptions
5
+ from ..data_types import QueryOptions
6
6
  from .sql_job import SQLJob
7
7
 
8
8
  T = TypeVar("T")
@@ -15,23 +15,14 @@ class QueryState(Enum):
15
15
  ERROR = 4
16
16
 
17
17
 
18
- def get_query_options(opts: Optional[Union[Dict[str, Any], QueryOptions]] = None) -> QueryOptions:
19
- if isinstance(opts, QueryOptions):
20
- return opts
21
- elif opts:
22
- return QueryOptions(**opts)
23
- else:
24
- return QueryOptions(isClCommand=False, parameters=None, autoClose=False)
25
-
26
-
27
18
  class Query(Generic[T]):
28
19
  global_query_list: List["Query[Any]"] = []
29
20
 
30
21
  def __init__(self, job: SQLJob, query: str, opts: QueryOptions) -> None:
31
22
  self.job = job
23
+ self.sql: str = query
32
24
  self.is_prepared: bool = True if opts.parameters is not None else False
33
25
  self.parameters: Optional[List[str]] = opts.parameters
34
- self.sql: str = query
35
26
  self.is_cl_command: Optional[bool] = opts.isClCommand
36
27
  self.should_auto_close: Optional[bool] = opts.autoClose
37
28
  self.is_terse_results: Optional[bool] = opts.isTerseResults
@@ -41,6 +32,17 @@ class Query(Generic[T]):
41
32
 
42
33
  Query.global_query_list.append(self)
43
34
 
35
+ def __enter__(self):
36
+ return self
37
+
38
+ def __exit__(self, exc_type, exc_value, traceback):
39
+ self.close()
40
+
41
+ def _execute_query(self, qeury_object: Dict[str, Any]) -> Dict[str, Any]:
42
+ self.job.send(json.dumps(qeury_object))
43
+ query_result: Dict[str, Any] = json.loads(self.job._socket.recv())
44
+ return query_result
45
+
44
46
  def run(self, rows_to_fetch: Optional[int] = None) -> Dict[str, Any]:
45
47
  if rows_to_fetch is None:
46
48
  rows_to_fetch = self._rows_to_fetch
@@ -71,8 +73,7 @@ class Query(Generic[T]):
71
73
  "parameters": self.parameters,
72
74
  }
73
75
 
74
- self.job.send(json.dumps(query_object))
75
- query_result: Dict[str, Any] = json.loads(self.job._socket.recv())
76
+ query_result: Dict[str, Any] = self._execute_query(query_object)
76
77
 
77
78
  self.state = (
78
79
  QueryState.RUN_DONE
@@ -81,7 +82,6 @@ class Query(Generic[T]):
81
82
  )
82
83
 
83
84
  if not query_result.get("success", False) and not self.is_cl_command:
84
- print(query_result)
85
85
  self.state = QueryState.ERROR
86
86
  error_keys = ["error", "sql_state", "sql_rc"]
87
87
  error_list = {
@@ -116,8 +116,7 @@ class Query(Generic[T]):
116
116
  }
117
117
 
118
118
  self._rows_to_fetch = rows_to_fetch
119
- self.job.send(json.dumps(query_object))
120
- query_result: Dict[str, Any] = json.loads(self.job._socket.recv())
119
+ query_result: Dict[str, Any] = self._execute_query(query_object)
121
120
 
122
121
  self.state = (
123
122
  QueryState.RUN_DONE
@@ -130,3 +129,18 @@ class Query(Generic[T]):
130
129
  raise Exception(query_result["error"] or "Failed to run Query (unknown error)")
131
130
 
132
131
  return query_result
132
+
133
+ def close(self):
134
+ if not self.job._socket.connected:
135
+ raise Exception("SQL Job not connected")
136
+ if self._correlation_id and self.state is not QueryState.RUN_DONE:
137
+ self.state = QueryState.RUN_DONE
138
+ query_object = {
139
+ "id": self.job._get_unique_id("sqlclose"),
140
+ "cont_id": self._correlation_id,
141
+ "type": "sqlclose",
142
+ }
143
+
144
+ return self._execute_query(query_object)
145
+ elif not self._correlation_id:
146
+ self.state = QueryState.RUN_DONE