Skip to content

Commit

Permalink
Adds --engine-option to csvsql. (#1256)
Browse files Browse the repository at this point in the history
feat(csvsql): adds --engine-option
  • Loading branch information
sgpeter1 authored Jul 23, 2024
1 parent f63a247 commit 6773cc0
Show file tree
Hide file tree
Showing 4 changed files with 20 additions and 17 deletions.
1 change: 1 addition & 0 deletions CHANGELOG.rst
Original file line number Diff line number Diff line change
@@ -1,6 +1,7 @@
Unreleased
----------

- feat: :doc:`/scripts/csvsql` adds a :code:`--engine-option` option.
- 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.

Expand Down
12 changes: 12 additions & 0 deletions csvkit/cli.py
Original file line number Diff line number Diff line change
@@ -1,4 +1,5 @@
#!/usr/bin/env python
import ast
import argparse
import bz2
import csv
Expand Down Expand Up @@ -560,6 +561,17 @@ def parse_column_identifiers(ids, column_names, column_offset=1, excluded_column
return [c for c in columns if c not in excludes]


def parse_list(pairs):
options = {}
for key, value in pairs:
try:
value = ast.literal_eval(value)
except ValueError:
pass
options[key] = value
return options


# Adapted from https://github.com/pallets/click/blame/main/src/click/utils.py
def _expand_args(args):
out = []
Expand Down
9 changes: 6 additions & 3 deletions csvkit/utilities/csvsql.py
Original file line number Diff line number Diff line change
@@ -1,13 +1,12 @@
#!/usr/bin/env python

import os.path
import sys

import agate
import agatesql # noqa: F401
from sqlalchemy import create_engine, dialects

from csvkit.cli import CSVKitUtility, isatty
from csvkit.cli import CSVKitUtility, isatty, parse_list

try:
import importlib_metadata
Expand All @@ -33,6 +32,10 @@ def add_arguments(self):
self.argparser.add_argument(
'--db', dest='connection_string',
help='If present, a SQLAlchemy connection string to use to directly execute generated SQL on a database.')
self.argparser.add_argument(
'--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(
'--query', dest='queries', action='append',
help='Execute one or more SQL queries delimited by ";" and output the result of the last query as CSV. '
Expand Down Expand Up @@ -137,7 +140,7 @@ def main(self):
# Establish database validity before reading CSV files
if self.args.connection_string:
try:
engine = create_engine(self.args.connection_string)
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 "
Expand Down
15 changes: 1 addition & 14 deletions csvkit/utilities/sql2csv.py
Original file line number Diff line number Diff line change
@@ -1,21 +1,8 @@
#!/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
from csvkit.cli import CSVKitUtility, parse_list


class SQL2CSV(CSVKitUtility):
Expand Down

0 comments on commit 6773cc0

Please sign in to comment.