Skip to content

Commit fd47e9e

Browse files
Jmoore1127Jmoore1127
authored
Merge pull request catherinedevlin#2 from kacecode/master
Merge upstream (catherinedevlin#1)
2 parents 581b22f + 1838317 commit fd47e9e

10 files changed

Lines changed: 277 additions & 176 deletions

File tree

‎NEWS.txt‎

Lines changed: 9 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -154,4 +154,12 @@ Deleted Plugin import left behind in 0.2.2
154154
0.3.9
155155
-----
156156

157-
* Restored Python 2 compatibility (thanks tokenmathguy)
157+
* Restored Python 2 compatibility (thanks tokenmathguy)
158+
159+
0.4.0
160+
-----
161+
162+
* Changed most non-SQL commands to argparse arguments (thanks pik)
163+
* User can specify a creator for connections (thanks pik)
164+
* Bogus pseudo-SQL command `PERSIST` removed, replaced with `--persist` arg
165+
* Turn off echo of connection information with `displaycon` in config

‎README.rst‎

Lines changed: 76 additions & 33 deletions
Original file line numberDiff line numberDiff line change
@@ -165,10 +165,36 @@ Note that an ``impala`` connection with `impyla`_ for HiveServer2 requires disa
165165

166166
.. _impyla: https://github.com/cloudera/impyla
167167

168-
connect_args can be provided after the connection string either as a local bindvar (prefixed with ':') or as a json string::
168+
Connection arguments not whitelisted by SQLALchemy can be provided as
169+
a flag with the connection string as a JSON string.
170+
Please note the flag parsing engine will not allow spaces in your JSON string::
169171

170-
In [1]: %sql sqlite:// :my_connect_args SELECT * FROM work;
171-
In [2]: %sql sqlite:// {"options": "-cmyarg=value"} SELECT * from work;
172+
%sql --connection_arguments {"timeout":10,"mode":"ro"} sqlite:// SELECT * FROM work;
173+
%sql -a {"timeout":10,"mode":"ro"} sqlite:// SELECT * from work;
174+
175+
.. _SQLAlchemy: https://docs.sqlalchemy.org/en/13/core/engines.html#custom-dbapi-args
176+
177+
DSN connections
178+
~~~~~~~~~~~~~~~
179+
180+
Alternately, you can store connection info in a
181+
configuration file, under a section name chosen to
182+
refer to your database.
183+
184+
For example, if dsn.ini contains
185+
186+
[DB_CONFIG_1]
187+
drivername=postgres
188+
host=my.remote.host
189+
port=5433
190+
database=mydatabase
191+
username=myuser
192+
password=1234
193+
194+
then you can
195+
196+
%config SqlMagic.dsn_filename='./dsn.ini'
197+
%sql --section DB_CONFIG_1
172198

173199
Configuration
174200
-------------
@@ -182,34 +208,45 @@ only the screen display is truncated.
182208

183209
.. code-block:: python
184210
185-
In [2]: %config SqlMagic
186-
SqlMagic options
187-
--------------
188-
SqlMagic.autocommit=<Bool>
189-
Current: True
190-
Set autocommit mode
191-
SqlMagic.autolimit=<Int>
192-
Current: 0
193-
Automatically limit the size of the returned result sets
194-
SqlMagic.autopandas=<Bool>
195-
Current: False
196-
Return Pandas DataFrames instead of regular result sets
197-
SqlMagic.displaylimit=<Int>
198-
Current: 0
199-
Automatically limit the number of rows displayed (full result set is still
200-
stored)
201-
SqlMagic.feedback=<Bool>
202-
Current: True
203-
Print number of rows affected by DML
204-
SqlMagic.short_errors=<Bool>
205-
Current: True
206-
Don't display the full traceback on SQL Programming Error
207-
SqlMagic.style=<Unicode>
208-
Current: 'DEFAULT'
209-
Set the table printing style to any of prettytable's defined styles
210-
(currently DEFAULT, MSWORD_FRIENDLY, PLAIN_COLUMNS, RANDOM)
211-
212-
In[3]: %config SqlMagic.feedback = False
211+
In [2]: %config SqlMagic
212+
SqlMagic options
213+
--------------
214+
SqlMagic.autocommit=<Bool>
215+
Current: True
216+
Set autocommit mode
217+
SqlMagic.autolimit=<Int>
218+
Current: 0
219+
Automatically limit the size of the returned result sets
220+
SqlMagic.autopandas=<Bool>
221+
Current: False
222+
Return Pandas DataFrames instead of regular result sets
223+
SqlMagic.column_local_vars=<Bool>
224+
Current: False
225+
Return data into local variables from column names
226+
SqlMagic.displaycon=<Bool>
227+
Current: False
228+
Show connection string after execute
229+
SqlMagic.displaylimit=<Int>
230+
Current: None
231+
Automatically limit the number of rows displayed (full result set is still
232+
stored)
233+
SqlMagic.dsn_filename=<Unicode>
234+
Current: 'odbc.ini'
235+
Path to DSN file. When the first argument is of the form [section], a
236+
sqlalchemy connection string is formed from the matching section in the DSN
237+
file.
238+
SqlMagic.feedback=<Bool>
239+
Current: False
240+
Print number of rows affected by DML
241+
SqlMagic.short_errors=<Bool>
242+
Current: True
243+
Don't display the full traceback on SQL Programming Error
244+
SqlMagic.style=<Unicode>
245+
Current: 'DEFAULT'
246+
Set the table printing style to any of prettytable's defined styles
247+
(currently DEFAULT, MSWORD_FRIENDLY, PLAIN_COLUMNS, RANDOM)
248+
249+
In[3]: %config SqlMagic.feedback = False
213250
214251
Please note: if you have autopandas set to true, the displaylimit option will not apply. You can set the pandas display limit by using the pandas ``max_rows`` option as described in the `pandas documentation <http://pandas.pydata.org/pandas-docs/version/0.18.1/options.html#frequently-used-options>`_.
215252

@@ -225,12 +262,15 @@ If you have installed ``pandas``, you can use a result set's
225262
226263
In [4]: dataframe = result.DataFrame()
227264
228-
The bogus non-standard pseudo-SQL command ``PERSIST`` will create a table name
265+
266+
The `--persist` argument, with the name of a
267+
DataFrame object in memory,
268+
will create a table name
229269
in the database from the named DataFrame.
230270

231271
.. code-block:: python
232272
233-
In [5]: %sql PERSIST dataframe
273+
In [5]: %sql --persist dataframe
234274
235275
In [6]: %sql SELECT * FROM dataframe;
236276
@@ -308,11 +348,14 @@ Credits
308348
- Mike Wilson for bind variable code
309349
- Thomas Kluyver and Steve Holden for debugging help
310350
- Berton Earnshaw for DSN connection syntax
351+
- Bruno Harbulot for DSN example
311352
- Andrés Celis for SQL Server bugfix
312353
- Michael Erasmus for DataFrame truth bugfix
313354
- Noam Finkelstein for README clarification
314355
- Xiaochuan Yu for `<<` operator, syntax colorization
315356
- Amjith Ramanujam for PGSpecial and incorporating it here
357+
- Alexander Maznev for better arg parsing, connections accepting specified creator
358+
- Jonathan Larkin for configurable displaycon
316359

317360
.. _Distribute: http://pypi.python.org/pypi/distribute
318361
.. _Buildout: http://www.buildout.org/

‎setup.py‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -7,7 +7,7 @@
77
NEWS = open(os.path.join(here, 'NEWS.txt'), encoding='utf-8').read()
88

99

10-
version = '0.3.9'
10+
version = '0.4.0'
1111

1212
install_requires = [
1313
'prettytable<1',

‎src/sql/connection.py‎

Lines changed: 30 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -31,9 +31,12 @@ def tell_format(cls):
3131
postgresql://username:password@hostname/dbname
3232
or an existing connection: %s""" % str(cls.connections.keys())
3333

34-
def __init__(self, connect_str=None, connect_args={}):
34+
def __init__(self, connect_str=None, connect_args={}, creator=None):
3535
try:
36-
engine = sqlalchemy.create_engine(connect_str, connect_args=connect_args)
36+
if creator:
37+
engine = sqlalchemy.create_engine(connect_str, connect_args=connect_args, creator=creator)
38+
else:
39+
engine = sqlalchemy.create_engine(connect_str, connect_args=connect_args)
3740
except: # TODO: bare except; but what's an ArgumentError?
3841
print(self.tell_format())
3942
raise
@@ -42,24 +45,28 @@ def __init__(self, connect_str=None, connect_args={}):
4245
self.name = self.assign_name(engine)
4346
self.session = engine.connect()
4447
self.connections[repr(self.metadata.bind.url)] = self
48+
self.connect_args = connect_args
4549
Connection.current = self
4650

4751
@classmethod
48-
def set(cls, descriptor, connect_args={}):
52+
def set(cls, descriptor, displaycon, connect_args={}, creator=None):
4953
"Sets the current database connection"
5054

5155
if descriptor:
5256
if isinstance(descriptor, Connection):
5357
cls.current = descriptor
5458
else:
5559
existing = rough_dict_get(cls.connections, descriptor)
56-
cls.current = existing or Connection(descriptor, connect_args)
60+
# http://docs.sqlalchemy.org/en/rel_0_9/core/engines.html#custom-dbapi-connect-arguments
61+
cls.current = existing or Connection(descriptor, connect_args, creator)
5762
else:
63+
5864
if cls.connections:
59-
print(cls.connection_list())
65+
if displaycon:
66+
print(cls.connection_list())
6067
else:
6168
if os.getenv('DATABASE_URL'):
62-
cls.current = Connection(os.getenv('DATABASE_URL'))
69+
cls.current = Connection(os.getenv('DATABASE_URL'), connect_args, creator)
6370
else:
6471
raise ConnectionError('Environment variable $DATABASE_URL not set, and no connect string given.')
6572
return cls.current
@@ -80,3 +87,20 @@ def connection_list(cls):
8087
template = ' {}'
8188
result.append(template.format(engine_url.__repr__()))
8289
return '\n'.join(result)
90+
def _close(cls, descriptor):
91+
if isinstance(descriptor, Connection):
92+
conn = descriptor
93+
else:
94+
conn = cls.connections.get(descriptor) or \
95+
cls.connections.get(descriptor.lower())
96+
if not conn:
97+
raise Exception("Could not close connection because it was not found amongst these: %s" \
98+
%str(cls.connections.keys()))
99+
cls.connections.pop(conn.name)
100+
cls.connections.pop(str(conn.metadata.bind.url))
101+
conn.session.close()
102+
103+
def close(self):
104+
self.__class__._close(self)
105+
106+

‎src/sql/magic.py‎

Lines changed: 42 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -1,3 +1,4 @@
1+
import json
12
import re
23
from IPython.core.magic import Magics, magics_class, cell_magic, line_magic, needs_local_scope
34
from IPython.display import display_javascript
@@ -7,6 +8,7 @@
78
except ImportError:
89
from IPython.config.configurable import Configurable
910
from IPython.utils.traitlets import Bool, Int, Unicode
11+
from IPython.core.magic_arguments import argument, magic_arguments, parse_argstring
1012
try:
1113
from pandas.core.frame import DataFrame, Series
1214
except ImportError:
@@ -25,6 +27,7 @@ class SqlMagic(Magics, Configurable):
2527
2628
Provides the %%sql magic."""
2729

30+
displaycon = Bool(True, config=True, help="Show connection string after execute")
2831
autolimit = Int(0, config=True, allow_none=True, help="Automatically limit the size of the returned result sets")
2932
style = Unicode('DEFAULT', config=True, help="Set the table printing style to any of prettytable's defined styles (currently DEFAULT, MSWORD_FRIENDLY, PLAIN_COLUMNS, RANDOM)")
3033
short_errors = Bool(True, config=True, help="Don't display the full traceback on SQL Programming Error")
@@ -49,7 +52,15 @@ def __init__(self, shell):
4952
@needs_local_scope
5053
@line_magic('sql')
5154
@cell_magic('sql')
52-
def execute(self, line, cell='', local_ns={}):
55+
@magic_arguments()
56+
@argument('line', default='', nargs='*', type=str, help='sql')
57+
@argument('-l', '--connections', action='store_true', help="list active connections")
58+
@argument('-x', '--close', type=str, help="close a session by name")
59+
@argument('-c', '--creator', type=str, help="specify creator function for new connection")
60+
@argument('-s', '--section', type=str, help="section of dsn_file to be used for generating a connection string")
61+
@argument('-p', '--persist', action='store_true', help="create a table name in the database from the named DataFrame")
62+
@argument('-a', '--connection_arguments', type=str, help="specify dictionary of connection arguments to pass to SQL driver")
63+
def execute(self, line='', cell='', local_ns={}):
5364
"""Runs SQL statement against a database, specified by SQLAlchemy connect string.
5465
5566
If no database connection has been established, first word
@@ -74,22 +85,46 @@ def execute(self, line, cell='', local_ns={}):
7485
mysql+pymysql://me:mypw@localhost/mydb
7586
7687
"""
88+
args = parse_argstring(self.execute, line)
89+
if args.connections:
90+
return sql.connection.Connection.connections
91+
elif args.close:
92+
return sql.connection.Connection._close(args.close)
93+
7794
# save globals and locals so they can be referenced in bind vars
7895
user_ns = self.shell.user_ns.copy()
7996
user_ns.update(local_ns)
8097

81-
parsed = sql.parse.parse('%s\n%s' % (line, cell), self, user_ns)
82-
flags = parsed['flags']
98+
parsed = sql.parse.parse(' '.join(args.line) + cell, self)
99+
100+
connect_str = parsed['connection']
101+
if args.section:
102+
connect_str = sql.parse.connection_from_dsn_section(args.section, self)
103+
104+
connect_args = {}
105+
if args.connection_arguments:
106+
try:
107+
connect_args = json.loads(args.connection_arguments)
108+
except Exception as e:
109+
print(e)
110+
raise e
111+
112+
if args.creator:
113+
args.creator = user_ns[args.creator]
114+
83115
try:
84-
conn = sql.connection.Connection.set(parsed['connection'], parsed['connect_args'])
116+
conn = sql.connection.Connection.set(parsed['connection'], displaycon=self.displaycon, connect_args=connect_args, creator=args.creator)
85117
except Exception as e:
86118
print(e)
87119
print(sql.connection.Connection.tell_format())
88120
return None
89121

90-
if flags.get('persist'):
122+
if args.persist:
91123
return self._persist_dataframe(parsed['sql'], conn, user_ns)
92124

125+
if not parsed['sql']:
126+
return
127+
93128
try:
94129
result = sql.run.run(conn, parsed['sql'], self, user_ns)
95130

@@ -112,8 +147,8 @@ def execute(self, line, cell='', local_ns={}):
112147
return None
113148
else:
114149

115-
if flags.get('result_var'):
116-
result_var = flags['result_var']
150+
if parsed['result_var']:
151+
result_var = parsed['result_var']
117152
print("Returning data to local variable {}".format(result_var))
118153
self.shell.user_ns.update({result_var: result})
119154
return None

0 commit comments

Comments
 (0)