feat(sql2csv): Add --execution-option option, use stream_results=True…

… by default, closes #1255
wireservice · Jul 16, 2024 · b8bcf7a · b8bcf7a
1 parent 2fab831
commit b8bcf7a
Show file tree

Hide file tree

Showing 17 changed files with 51 additions and 16 deletions.
diff --git a/CHANGELOG.rst b/CHANGELOG.rst
@@ -1,3 +1,9 @@
+Unreleased
+----------
+
+-  feat: :doc:`/scripts/sql2csv` adds a :code:`--execution-option` option.
+-  feat: :doc:`/scripts/sql2csv` uses the ``stream_results=True`` execution option, by default, to not load all data into memory at once.
+
 2.0.1 - July 12, 2024
 ---------------------
 

diff --git a/csvkit/utilities/sql2csv.py b/csvkit/utilities/sql2csv.py
@@ -1,11 +1,23 @@
 #!/usr/bin/env python
+import ast
 
 import agate
 from sqlalchemy import create_engine
 
 from csvkit.cli import CSVKitUtility
 
 
+def parse_list(pairs):
+    options = {}
+    for key, value in pairs:
+        try:
+            value = ast.literal_eval(value)
+        except ValueError:
+            pass
+        options[key] = value
+    return options
+
+
 class SQL2CSV(CSVKitUtility):
     description = 'Execute a SQL query on a database and output the result to a CSV file.'
     # Overrides all flags except --linenumbers, --verbose, --version.
@@ -19,6 +31,13 @@ def add_arguments(self):
             '--engine-option', dest='engine_option', nargs=2, action='append', default=[],
             help="A keyword argument to SQLAlchemy's create_engine(), as a space-separated pair. "
                  "This option can be specified multiple times. For example: thick_mode True")
+        self.argparser.add_argument(
+            '--execution-option', dest='execution_option', nargs=2, action='append',
+            # https://docs.sqlalchemy.org/en/20/core/connections.html#sqlalchemy.engine.Connection.execution_options.params.no_parameters
+            # https://docs.sqlalchemy.org/en/20/core/connections.html#sqlalchemy.engine.Connection.execution_options.params.stream_results
+            default=[['no_parameters', True], ['stream_results', True]],
+            help="A keyword argument to SQLAlchemy's execution_options(), as a space-separated pair. "
+                 "This option can be specified multiple times. For example: stream_results True")
         self.argparser.add_argument(
             metavar='FILE', nargs='?', dest='input_path',
             help='The file to use as SQL query. If FILE and --query are omitted, the query is piped data via STDIN.')
@@ -49,7 +68,7 @@ def main(self):
             self.argparser.error('You must provide an input file or piped data.')
 
         try:
-            engine = create_engine(self.args.connection_string, **dict(self.args.engine_option))
+            engine = create_engine(self.args.connection_string, **parse_list(self.args.engine_option))
         except ImportError as e:
             raise ImportError(
                 "You don't appear to have the necessary database backend installed for connection string you're "
@@ -73,7 +92,7 @@ def main(self):
 
             self.input_file.close()
 
-        rows = connection.execution_options(no_parameters=True).exec_driver_sql(query)
+        rows = connection.execution_options(**parse_list(self.args.execution_option)).exec_driver_sql(query)
         output = agate.csv.writer(self.output_file, **self.writer_kwargs)
 
         if rows.returns_rows:

diff --git a/docs/scripts/sql2csv.rst b/docs/scripts/sql2csv.rst
@@ -28,6 +28,11 @@ Executes arbitrary commands against a SQL database and outputs the results as a
                            A keyword argument to SQLAlchemy's create_engine(), as
                            a space-separated pair. This option can be specified
                            multiple times. For example: thick_mode True
+     --execution-option EXECUTION_OPTION EXECUTION_OPTION
+                           A keyword argument to SQLAlchemy's
+                           execution_options(), as a space-separated pair. This
+                           option can be specified multiple times. For example:
+                           stream_results True
      --query QUERY         The SQL query to execute. Overrides FILE and STDIN.
      -e ENCODING, --encoding ENCODING
                            Specify the encoding of the input query file.

diff --git a/man/csvclean.1 b/man/csvclean.1
@@ -27,7 +27,7 @@ level margin: \\n[rst2man-indent\\n[rst2man-indent-level]]
 .\" new: \\n[rst2man-indent\\n[rst2man-indent-level]]
 .in \\n[rst2man-indent\\n[rst2man-indent-level]]u
 ..
-.TH "CSVCLEAN" "1" "Jul 12, 2024" "2.0.1" "csvkit"
+.TH "CSVCLEAN" "1" "Jul 16, 2024" "2.0.1" "csvkit"
 .SH NAME
 csvclean \- csvclean Documentation
 .SH DESCRIPTION

diff --git a/man/csvcut.1 b/man/csvcut.1
@@ -27,7 +27,7 @@ level margin: \\n[rst2man-indent\\n[rst2man-indent-level]]
 .\" new: \\n[rst2man-indent\\n[rst2man-indent-level]]
 .in \\n[rst2man-indent\\n[rst2man-indent-level]]u
 ..
-.TH "CSVCUT" "1" "Jul 12, 2024" "2.0.1" "csvkit"
+.TH "CSVCUT" "1" "Jul 16, 2024" "2.0.1" "csvkit"
 .SH NAME
 csvcut \- csvcut Documentation
 .SH DESCRIPTION

diff --git a/man/csvformat.1 b/man/csvformat.1
@@ -27,7 +27,7 @@ level margin: \\n[rst2man-indent\\n[rst2man-indent-level]]
 .\" new: \\n[rst2man-indent\\n[rst2man-indent-level]]
 .in \\n[rst2man-indent\\n[rst2man-indent-level]]u
 ..
-.TH "CSVFORMAT" "1" "Jul 12, 2024" "2.0.1" "csvkit"
+.TH "CSVFORMAT" "1" "Jul 16, 2024" "2.0.1" "csvkit"
 .SH NAME
 csvformat \- csvformat Documentation
 .SH DESCRIPTION

diff --git a/man/csvgrep.1 b/man/csvgrep.1
@@ -27,7 +27,7 @@ level margin: \\n[rst2man-indent\\n[rst2man-indent-level]]
 .\" new: \\n[rst2man-indent\\n[rst2man-indent-level]]
 .in \\n[rst2man-indent\\n[rst2man-indent-level]]u
 ..
-.TH "CSVGREP" "1" "Jul 12, 2024" "2.0.1" "csvkit"
+.TH "CSVGREP" "1" "Jul 16, 2024" "2.0.1" "csvkit"
 .SH NAME
 csvgrep \- csvgrep Documentation
 .SH DESCRIPTION

diff --git a/man/csvjoin.1 b/man/csvjoin.1
@@ -27,7 +27,7 @@ level margin: \\n[rst2man-indent\\n[rst2man-indent-level]]
 .\" new: \\n[rst2man-indent\\n[rst2man-indent-level]]
 .in \\n[rst2man-indent\\n[rst2man-indent-level]]u
 ..
-.TH "CSVJOIN" "1" "Jul 12, 2024" "2.0.1" "csvkit"
+.TH "CSVJOIN" "1" "Jul 16, 2024" "2.0.1" "csvkit"
 .SH NAME
 csvjoin \- csvjoin Documentation
 .SH DESCRIPTION

diff --git a/man/csvjson.1 b/man/csvjson.1
@@ -27,7 +27,7 @@ level margin: \\n[rst2man-indent\\n[rst2man-indent-level]]
 .\" new: \\n[rst2man-indent\\n[rst2man-indent-level]]
 .in \\n[rst2man-indent\\n[rst2man-indent-level]]u
 ..
-.TH "CSVJSON" "1" "Jul 12, 2024" "2.0.1" "csvkit"
+.TH "CSVJSON" "1" "Jul 16, 2024" "2.0.1" "csvkit"
 .SH NAME
 csvjson \- csvjson Documentation
 .SH DESCRIPTION

diff --git a/man/csvlook.1 b/man/csvlook.1
@@ -27,7 +27,7 @@ level margin: \\n[rst2man-indent\\n[rst2man-indent-level]]
 .\" new: \\n[rst2man-indent\\n[rst2man-indent-level]]
 .in \\n[rst2man-indent\\n[rst2man-indent-level]]u
 ..
-.TH "CSVLOOK" "1" "Jul 12, 2024" "2.0.1" "csvkit"
+.TH "CSVLOOK" "1" "Jul 16, 2024" "2.0.1" "csvkit"
 .SH NAME
 csvlook \- csvlook Documentation
 .SH DESCRIPTION

diff --git a/man/csvpy.1 b/man/csvpy.1
@@ -27,7 +27,7 @@ level margin: \\n[rst2man-indent\\n[rst2man-indent-level]]
 .\" new: \\n[rst2man-indent\\n[rst2man-indent-level]]
 .in \\n[rst2man-indent\\n[rst2man-indent-level]]u
 ..
-.TH "CSVPY" "1" "Jul 12, 2024" "2.0.1" "csvkit"
+.TH "CSVPY" "1" "Jul 16, 2024" "2.0.1" "csvkit"
 .SH NAME
 csvpy \- csvpy Documentation
 .SH DESCRIPTION

diff --git a/man/csvsort.1 b/man/csvsort.1
@@ -27,7 +27,7 @@ level margin: \\n[rst2man-indent\\n[rst2man-indent-level]]
 .\" new: \\n[rst2man-indent\\n[rst2man-indent-level]]
 .in \\n[rst2man-indent\\n[rst2man-indent-level]]u
 ..
-.TH "CSVSORT" "1" "Jul 12, 2024" "2.0.1" "csvkit"
+.TH "CSVSORT" "1" "Jul 16, 2024" "2.0.1" "csvkit"
 .SH NAME
 csvsort \- csvsort Documentation
 .SH DESCRIPTION

diff --git a/man/csvsql.1 b/man/csvsql.1
@@ -27,7 +27,7 @@ level margin: \\n[rst2man-indent\\n[rst2man-indent-level]]
 .\" new: \\n[rst2man-indent\\n[rst2man-indent-level]]
 .in \\n[rst2man-indent\\n[rst2man-indent-level]]u
 ..
-.TH "CSVSQL" "1" "Jul 12, 2024" "2.0.1" "csvkit"
+.TH "CSVSQL" "1" "Jul 16, 2024" "2.0.1" "csvkit"
 .SH NAME
 csvsql \- csvsql Documentation
 .SH DESCRIPTION

diff --git a/man/csvstack.1 b/man/csvstack.1
@@ -27,7 +27,7 @@ level margin: \\n[rst2man-indent\\n[rst2man-indent-level]]
 .\" new: \\n[rst2man-indent\\n[rst2man-indent-level]]
 .in \\n[rst2man-indent\\n[rst2man-indent-level]]u
 ..
-.TH "CSVSTACK" "1" "Jul 12, 2024" "2.0.1" "csvkit"
+.TH "CSVSTACK" "1" "Jul 16, 2024" "2.0.1" "csvkit"
 .SH NAME
 csvstack \- csvstack Documentation
 .SH DESCRIPTION

diff --git a/man/csvstat.1 b/man/csvstat.1
@@ -27,7 +27,7 @@ level margin: \\n[rst2man-indent\\n[rst2man-indent-level]]
 .\" new: \\n[rst2man-indent\\n[rst2man-indent-level]]
 .in \\n[rst2man-indent\\n[rst2man-indent-level]]u
 ..
-.TH "CSVSTAT" "1" "Jul 12, 2024" "2.0.1" "csvkit"
+.TH "CSVSTAT" "1" "Jul 16, 2024" "2.0.1" "csvkit"
 .SH NAME
 csvstat \- csvstat Documentation
 .SH DESCRIPTION

diff --git a/man/in2csv.1 b/man/in2csv.1
@@ -27,7 +27,7 @@ level margin: \\n[rst2man-indent\\n[rst2man-indent-level]]
 .\" new: \\n[rst2man-indent\\n[rst2man-indent-level]]
 .in \\n[rst2man-indent\\n[rst2man-indent-level]]u
 ..
-.TH "IN2CSV" "1" "Jul 12, 2024" "2.0.1" "csvkit"
+.TH "IN2CSV" "1" "Jul 16, 2024" "2.0.1" "csvkit"
 .SH NAME
 in2csv \- in2csv Documentation
 .SH DESCRIPTION

diff --git a/man/sql2csv.1 b/man/sql2csv.1
@@ -27,7 +27,7 @@ level margin: \\n[rst2man-indent\\n[rst2man-indent-level]]
 .\" new: \\n[rst2man-indent\\n[rst2man-indent-level]]
 .in \\n[rst2man-indent\\n[rst2man-indent-level]]u
 ..
-.TH "SQL2CSV" "1" "Jul 12, 2024" "2.0.1" "csvkit"
+.TH "SQL2CSV" "1" "Jul 16, 2024" "2.0.1" "csvkit"
 .SH NAME
 sql2csv \- sql2csv Documentation
 .SH DESCRIPTION
@@ -57,6 +57,11 @@ optional arguments:
                         A keyword argument to SQLAlchemy\(aqs create_engine(), as
                         a space\-separated pair. This option can be specified
                         multiple times. For example: thick_mode True
+  \-\-execution\-option EXECUTION_OPTION EXECUTION_OPTION
+                        A keyword argument to SQLAlchemy\(aqs
+                        execution_options(), as a space\-separated pair. This
+                        option can be specified multiple times. For example:
+                        stream_results True
   \-\-query QUERY         The SQL query to execute. Overrides FILE and STDIN.
   \-e ENCODING, \-\-encoding ENCODING
                         Specify the encoding of the input query file.