-
Notifications
You must be signed in to change notification settings - Fork 13
Expand file tree
/
Copy pathsql_script.py
More file actions
52 lines (47 loc) · 1.75 KB
/
sql_script.py
File metadata and controls
52 lines (47 loc) · 1.75 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
import abc
from typing import Sequence
from exabel.scripts.command_line_script import CommandLineScript
from exabel.services.sql.sql_reader_configuration import SqlReaderConfiguration
from exabel.services.sql.sqlalchemy_reader import SQLAlchemyReader
class SqlScript(CommandLineScript, abc.ABC):
"""
Base class for scripts that perform a query against a SQL database and optionally store the
result to a file.
"""
def __init__(
self,
argv: Sequence[str],
description: str,
reader_configuration_class: type[SqlReaderConfiguration],
):
super().__init__(argv, description)
self.reader_configuration_class = reader_configuration_class
self.parser.add_argument(
"--query",
required=True,
help="The query to execute.",
)
self.parser.add_argument(
"--output-file",
help=(
"The file to write the result to. If no output file is specified, a sample is "
"printed."
),
)
self.parser.add_argument(
"--batch-size",
type=int,
help=(
"The number of rows to read at a time. If not specified, the entire result is "
"read into memory."
),
)
def run(self) -> None:
args = self.parse_arguments()
self.setup_logging()
configuration = self.reader_configuration_class.from_args(args)
connection_string, connect_args = configuration.get_connection_string_and_kwargs()
reader = SQLAlchemyReader(connection_string, kwargs=connect_args)
reader.read_sql_query_and_write_result(
args.query, args.output_file, batch_size=args.batch_size
)