├── .gitignore ├── HACKING.txt ├── LICENSE ├── MANIFEST.in ├── NEWS.rst ├── README.rst ├── bootstrap.py ├── buildout.cfg ├── examples ├── wordcount.png ├── writers.ipynb └── writers.png ├── ipython-sql.wpr ├── requirements-dev.txt ├── requirements.txt ├── run_tests.sh ├── setup.py ├── src ├── sql │ ├── __init__.py │ ├── column_guesser.py │ ├── connection.py │ ├── magic.py │ ├── parse.py │ └── run.py └── tests │ ├── test_column_guesser.py │ ├── test_dsn_config.ini │ ├── test_magic.py │ └── test_parse.py └── tox.ini /.gitignore: -------------------------------------------------------------------------------- 1 | *.py[cod] 2 | 3 | # C extensions 4 | *.so 5 | 6 | # Packages 7 | *.egg 8 | *.egg-info 9 | dist 10 | build 11 | eggs 12 | parts 13 | bin 14 | var 15 | sdist 16 | develop-eggs 17 | .installed.cfg 18 | lib 19 | lib64 20 | 21 | # Installer logs 22 | pip-log.txt 23 | 24 | # Unit test / coverage reports 25 | .coverage 26 | .tox 27 | nosetests.xml 28 | 29 | # Translations 30 | *.mo 31 | 32 | # Mr Developer 33 | .mr.developer.cfg 34 | .project 35 | .pydevproject 36 | 37 | # Pycharm 38 | /.idea 39 | 40 | .venv 41 | -------------------------------------------------------------------------------- /HACKING.txt: -------------------------------------------------------------------------------- 1 | Development setup 2 | ================= 3 | 4 | Running nose tests with IPython is tricky, so there's a 5 | run_tests.sh script for it. 6 | 7 | pip install -e . 8 | ./run_tests.sh 9 | 10 | To temporarily insert breakpoints for debugging: `from nose.tools import set_trace; set_trace()`. 11 | Or, if running tests, use `pytest.set_trace()`. 12 | 13 | Tests have requirements not installed by setup.py: 14 | 15 | - nose 16 | - pandas 17 | 18 | Release HOWTO 19 | ============= 20 | 21 | To make a release, 22 | 23 | 1) Update release date/version in NEWS.txt and setup.py 24 | 2) Run 'python setup.py sdist' 25 | 3) Test the generated source distribution in dist/ 26 | 4) Upload to PyPI: 'python setup.py sdist register upload' 27 | 5) Increase version in setup.py (for next release) 28 | -------------------------------------------------------------------------------- /LICENSE: -------------------------------------------------------------------------------- 1 | MIT License 2 | 3 | Copyright (c) 2014 Catherine Devlin 4 | 5 | Permission is hereby granted, free of charge, to any person obtaining a copy 6 | of this software and associated documentation files (the "Software"), to deal 7 | in the Software without restriction, including without limitation the rights 8 | to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 9 | copies of the Software, and to permit persons to whom the Software is 10 | furnished to do so, subject to the following conditions: 11 | 12 | The above copyright notice and this permission notice shall be included in all 13 | copies or substantial portions of the Software. 14 | 15 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 16 | IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 17 | FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 18 | AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 19 | LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 20 | OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE 21 | SOFTWARE. 22 | -------------------------------------------------------------------------------- /MANIFEST.in: -------------------------------------------------------------------------------- 1 | include README.rst 2 | include NEWS.rst 3 | include LICENSE 4 | -------------------------------------------------------------------------------- /NEWS.rst: -------------------------------------------------------------------------------- 1 | News 2 | ---- 3 | 4 | 0.1 5 | ~~~ 6 | 7 | *Release date: 21-Mar-2013* 8 | 9 | * Initial release 10 | 11 | 0.1.1 12 | ~~~~~ 13 | 14 | *Release date: 29-Mar-2013* 15 | 16 | * Release to PyPI 17 | 18 | * Results returned as lists 19 | 20 | * print(_) to get table form in text console 21 | 22 | * set autolimit and text wrap in configuration 23 | 24 | 25 | 0.1.2 26 | ~~~~~ 27 | 28 | *Release date: 29-Mar-2013* 29 | 30 | * Python 3 compatibility 31 | 32 | * use prettyprint package 33 | 34 | * allow multiple SQL per cell 35 | 36 | 0.2.0 37 | ~~~~~ 38 | 39 | *Release date: 30-May-2013* 40 | 41 | * Accept bind variables (Thanks Mike Wilson!) 42 | 43 | 0.2.1 44 | ~~~~~ 45 | 46 | *Release date: 15-June-2013* 47 | 48 | * Recognize socket connection strings 49 | 50 | * Bugfix - issue 4 (remember existing connections by case) 51 | 52 | 0.2.2 53 | ~~~~~ 54 | 55 | *Release date: 30-July-2013* 56 | 57 | Converted from an IPython Plugin to an Extension for 1.0 compatibility 58 | 59 | 0.2.2.1 60 | ~~~~~~~ 61 | 62 | *Release date: 01-Aug-2013* 63 | 64 | Deleted Plugin import left behind in 0.2.2 65 | 66 | 0.2.3 67 | ~~~~~ 68 | 69 | *Release date: 20-Sep-2013* 70 | 71 | * Contributions from Olivier Le Thanh Duong: 72 | 73 | - SQL errors reported without internal IPython error stack 74 | 75 | - Proper handling of configuration 76 | 77 | * Added .DataFrame(), .pie(), .plot(), and .bar() methods to 78 | result sets 79 | 80 | 0.3.0 81 | ~~~~~ 82 | 83 | *Release date: 13-Oct-2013* 84 | 85 | * displaylimit config parameter 86 | 87 | * reports number of rows affected by each query 88 | 89 | * test suite working again 90 | 91 | * dict-style access for result sets by primary key 92 | 93 | 0.3.1 94 | ~~~~~ 95 | 96 | * Reporting of number of rows affected configurable with ``feedback`` 97 | 98 | * Local variables usable as SQL bind variables 99 | 100 | 0.3.2 101 | ~~~~~ 102 | 103 | * ``.csv(filename=None)`` method added to result sets 104 | 105 | 0.3.3 106 | ~~~~~ 107 | 108 | * Python 3 compatibility restored 109 | * DSN access supported (thanks Berton Earnshaw) 110 | 111 | 0.3.4 112 | ~~~~~ 113 | 114 | * PERSIST pseudo-SQL command added 115 | 116 | 0.3.5 117 | ~~~~~ 118 | 119 | * Indentations visible in HTML cells 120 | * COMMIT each SQL statement immediately - prevent locks 121 | 122 | 0.3.6 123 | ~~~~~ 124 | 125 | * Fixed issue #30, commit failures for sqlite (thanks stonebig, jandot) 126 | 127 | 0.3.7 128 | ~~~~~ 129 | 130 | * New `column_local_vars` config option submitted by darikg 131 | * Avoid contaminating user namespace from locals (thanks alope107) 132 | 133 | 0.3.7.1 134 | ~~~~~~~ 135 | 136 | * Avoid "connection busy" error for SQL Server (thanks Andrés Celis) 137 | 138 | 0.3.8 139 | ~~~~~ 140 | 141 | * Stop warnings for deprecated use of IPython 3 traitlets in IPython 4 (thanks graphaelli; also stonebig, aebrahim, mccahill) 142 | * README update for keeping connection info private, from eshilts 143 | 144 | 0.3.9 145 | ~~~~~ 146 | 147 | * Fix truth value of DataFrame error (thanks michael-erasmus) 148 | * `<<` operator (thanks xiaochuanyu) 149 | * added README example (thanks tanhuil) 150 | * bugfix in executing column_local_vars (thanks tebeka) 151 | * pgspecial installation optional (thanks jstoebel and arjoe) 152 | * conceal passwords in connection strings (thanks jstoebel) 153 | 154 | 0.3.9 155 | ~~~~~ 156 | 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 166 | * Consistent support for {} variables (thanks Lucas) 167 | 168 | 0.4.1 169 | ~~~~~ 170 | 171 | * Fixed .rst file location in MANIFEST.in 172 | * Parse SQL comments in first line 173 | * Bugfixes for DSN, `--close`, others 174 | 175 | 0.5.0 176 | ~~~~~ 177 | 178 | * Use SQLAlchemy 2.0 179 | * Drop undocumented support for dict-style access to raw row instances -------------------------------------------------------------------------------- /README.rst: -------------------------------------------------------------------------------- 1 | =========== 2 | ipython-sql 3 | =========== 4 | 5 | :Author: Catherine Devlin, http://catherinedevlin.blogspot.com 6 | 7 | Introduces a %sql (or %%sql) magic. 8 | 9 | Legacy project 10 | -------------- 11 | 12 | IPython-SQL's functionality and maintenance have been eclipsed by JupySQL_, a fork maintained and developed by the Ploomber team. Future work will be directed into JupySQL - please file issues there, as well! 13 | 14 | Description 15 | ----------- 16 | 17 | Connect to a database, using `SQLAlchemy URL`_ connect strings, then issue SQL 18 | commands within IPython or IPython Notebook. 19 | 20 | .. image:: https://raw.github.com/catherinedevlin/ipython-sql/master/examples/writers.png 21 | :width: 600px 22 | :alt: screenshot of ipython-sql in the Notebook 23 | 24 | Examples 25 | -------- 26 | 27 | .. code-block:: python 28 | 29 | In [1]: %load_ext sql 30 | 31 | In [2]: %%sql postgresql://will:longliveliz@localhost/shakes 32 | ...: select * from character 33 | ...: where abbrev = 'ALICE' 34 | ...: 35 | Out[2]: [(u'Alice', u'Alice', u'ALICE', u'a lady attending on Princess Katherine', 22)] 36 | 37 | In [3]: result = _ 38 | 39 | In [4]: print(result) 40 | charid charname abbrev description speechcount 41 | ================================================================================= 42 | Alice Alice ALICE a lady attending on Princess Katherine 22 43 | 44 | In [4]: result.keys 45 | Out[5]: [u'charid', u'charname', u'abbrev', u'description', u'speechcount'] 46 | 47 | In [6]: result[0][0] 48 | Out[6]: u'Alice' 49 | 50 | In [7]: result[0].description 51 | Out[7]: u'a lady attending on Princess Katherine' 52 | 53 | After the first connection, connect info can be omitted:: 54 | 55 | In [8]: %sql select count(*) from work 56 | Out[8]: [(43L,)] 57 | 58 | Connections to multiple databases can be maintained. You can refer to 59 | an existing connection by username@database 60 | 61 | .. code-block:: python 62 | 63 | In [9]: %%sql will@shakes 64 | ...: select charname, speechcount from character 65 | ...: where speechcount = (select max(speechcount) 66 | ...: from character); 67 | ...: 68 | Out[9]: [(u'Poet', 733)] 69 | 70 | In [10]: print(_) 71 | charname speechcount 72 | ====================== 73 | Poet 733 74 | 75 | If no connect string is supplied, ``%sql`` will provide a list of existing connections; 76 | however, if no connections have yet been made and the environment variable ``DATABASE_URL`` 77 | is available, that will be used. 78 | 79 | For secure access, you may dynamically access your credentials (e.g. from your system environment or `getpass.getpass`) to avoid storing your password in the notebook itself. Use the `$` before any variable to access it in your `%sql` command. 80 | 81 | .. code-block:: python 82 | 83 | In [11]: user = os.getenv('SOME_USER') 84 | ....: password = os.getenv('SOME_PASSWORD') 85 | ....: connection_string = "postgresql://{user}:{password}@localhost/some_database".format(user=user, password=password) 86 | ....: %sql $connection_string 87 | Out[11]: u'Connected: some_user@some_database' 88 | 89 | You may use multiple SQL statements inside a single cell, but you will 90 | only see any query results from the last of them, so this really only 91 | makes sense for statements with no output 92 | 93 | .. code-block:: python 94 | 95 | In [11]: %%sql sqlite:// 96 | ....: CREATE TABLE writer (first_name, last_name, year_of_death); 97 | ....: INSERT INTO writer VALUES ('William', 'Shakespeare', 1616); 98 | ....: INSERT INTO writer VALUES ('Bertold', 'Brecht', 1956); 99 | ....: 100 | Out[11]: [] 101 | 102 | 103 | As a convenience, dict-style access for result sets is supported, with the 104 | leftmost column serving as key, for unique values. 105 | 106 | .. code-block:: python 107 | 108 | In [12]: result = %sql select * from work 109 | 43 rows affected. 110 | 111 | In [13]: result['richard2'] 112 | Out[14]: (u'richard2', u'Richard II', u'History of Richard II', 1595, u'h', None, u'Moby', 22411, 628) 113 | 114 | Results can also be retrieved as an iterator of dictionaries (``result.dicts()``) 115 | or a single dictionary with a tuple of scalar values per key (``result.dict()``) 116 | 117 | Variable substitution 118 | --------------------- 119 | 120 | Bind variables (bind parameters) can be used in the "named" (:x) style. 121 | The variable names used should be defined in the local namespace. 122 | 123 | .. code-block:: python 124 | 125 | In [15]: name = 'Countess' 126 | 127 | In [16]: %sql select description from character where charname = :name 128 | Out[16]: [(u'mother to Bertram',)] 129 | 130 | In [17]: %sql select description from character where charname = '{name}' 131 | Out[17]: [(u'mother to Bertram',)] 132 | 133 | Alternately, ``$variable_name`` or ``{variable_name}`` can be 134 | used to inject variables from the local namespace into the SQL 135 | statement before it is formed and passed to the SQL engine. 136 | (Using ``$`` and ``{}`` together, as in ``${variable_name}``, 137 | is not supported.) 138 | 139 | Bind variables are passed through to the SQL engine and can only 140 | be used to replace strings passed to SQL. ``$`` and ``{}`` are 141 | substituted before passing to SQL and can be used to form SQL 142 | statements dynamically. 143 | 144 | Assignment 145 | ---------- 146 | 147 | Ordinary IPython assignment works for single-line `%sql` queries: 148 | 149 | .. code-block:: python 150 | 151 | In [18]: works = %sql SELECT title, year FROM work 152 | 43 rows affected. 153 | 154 | The `<<` operator captures query results in a local variable, and 155 | can be used in multi-line ``%%sql``: 156 | 157 | .. code-block:: python 158 | 159 | In [19]: %%sql works << SELECT title, year 160 | ...: FROM work 161 | ...: 162 | 43 rows affected. 163 | Returning data to local variable works 164 | 165 | Connecting 166 | ---------- 167 | 168 | Connection strings are `SQLAlchemy URL`_ standard. 169 | 170 | Some example connection strings:: 171 | 172 | mysql+pymysql://scott:tiger@localhost/foo 173 | oracle://scott:tiger@127.0.0.1:1521/sidname 174 | sqlite:// 175 | sqlite:///foo.db 176 | mssql+pyodbc://username:password@host/database?driver=SQL+Server+Native+Client+11.0 177 | 178 | .. _`SQLAlchemy URL`: http://docs.sqlalchemy.org/en/latest/core/engines.html#database-urls 179 | 180 | Note that ``mysql`` and ``mysql+pymysql`` connections (and perhaps others) 181 | don't read your client character set information from .my.cnf. You need 182 | to specify it in the connection string:: 183 | 184 | mysql+pymysql://scott:tiger@localhost/foo?charset=utf8 185 | 186 | Note that an ``impala`` connection with `impyla`_ for HiveServer2 requires disabling autocommit:: 187 | 188 | %config SqlMagic.autocommit=False 189 | %sql impala://hserverhost:port/default?kerberos_service_name=hive&auth_mechanism=GSSAPI 190 | 191 | .. _impyla: https://github.com/cloudera/impyla 192 | 193 | Connection arguments not whitelisted by SQLALchemy can be provided as 194 | a flag with (-a|--connection_arguments)the connection string as a JSON string. 195 | See `SQLAlchemy Args`_. 196 | 197 | | %sql --connection_arguments {"timeout":10,"mode":"ro"} sqlite:// SELECT * FROM work; 198 | | %sql -a '{"timeout":10, "mode":"ro"}' sqlite:// SELECT * from work; 199 | 200 | .. _`SQLAlchemy Args`: https://docs.sqlalchemy.org/en/13/core/engines.html#custom-dbapi-args 201 | 202 | DSN connections 203 | ~~~~~~~~~~~~~~~ 204 | 205 | Alternately, you can store connection info in a 206 | configuration file, under a section name chosen to 207 | refer to your database. 208 | 209 | For example, if dsn.ini contains 210 | 211 | | [DB_CONFIG_1] 212 | | drivername=postgres 213 | | host=my.remote.host 214 | | port=5433 215 | | database=mydatabase 216 | | username=myuser 217 | | password=1234 218 | 219 | then you can 220 | 221 | | %config SqlMagic.dsn_filename='./dsn.ini' 222 | | %sql --section DB_CONFIG_1 223 | 224 | Configuration 225 | ------------- 226 | 227 | Query results are loaded as lists, so very large result sets may use up 228 | your system's memory and/or hang your browser. There is no autolimit 229 | by default. However, `autolimit` (if set) limits the size of the result 230 | set (usually with a `LIMIT` clause in the SQL). `displaylimit` is similar, 231 | but the entire result set is still pulled into memory (for later analysis); 232 | only the screen display is truncated. 233 | 234 | .. code-block:: python 235 | 236 | In [2]: %config SqlMagic 237 | SqlMagic options 238 | -------------- 239 | SqlMagic.autocommit= 240 | Current: True 241 | Set autocommit mode 242 | SqlMagic.autolimit= 243 | Current: 0 244 | Automatically limit the size of the returned result sets 245 | SqlMagic.autopandas= 246 | Current: False 247 | Return Pandas DataFrames instead of regular result sets 248 | SqlMagic.column_local_vars= 249 | Current: False 250 | Return data into local variables from column names 251 | SqlMagic.displaycon= 252 | Current: False 253 | Show connection string after execute 254 | SqlMagic.displaylimit= 255 | Current: None 256 | Automatically limit the number of rows displayed (full result set is still 257 | stored) 258 | SqlMagic.dsn_filename= 259 | Current: 'odbc.ini' 260 | Path to DSN file. When the first argument is of the form [section], a 261 | sqlalchemy connection string is formed from the matching section in the DSN 262 | file. 263 | SqlMagic.feedback= 264 | Current: False 265 | Print number of rows affected by DML 266 | SqlMagic.short_errors= 267 | Current: True 268 | Don't display the full traceback on SQL Programming Error 269 | SqlMagic.style= 270 | Current: 'DEFAULT' 271 | Set the table printing style to any of prettytable's defined styles 272 | (currently DEFAULT, MSWORD_FRIENDLY, PLAIN_COLUMNS, RANDOM) 273 | 274 | In[3]: %config SqlMagic.feedback = False 275 | 276 | 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 `_. 277 | 278 | Pandas 279 | ------ 280 | 281 | If you have installed ``pandas``, you can use a result set's 282 | ``.DataFrame()`` method 283 | 284 | .. code-block:: python 285 | 286 | In [3]: result = %sql SELECT * FROM character WHERE speechcount > 25 287 | 288 | In [4]: dataframe = result.DataFrame() 289 | 290 | 291 | The ``--persist`` argument, with the name of a 292 | DataFrame object in memory, 293 | will create a table name 294 | in the database from the named DataFrame. 295 | Or use ``--append`` to add rows to an existing 296 | table by that name. 297 | 298 | .. code-block:: python 299 | 300 | In [5]: %sql --persist dataframe 301 | 302 | In [6]: %sql SELECT * FROM dataframe; 303 | 304 | .. _Pandas: http://pandas.pydata.org/ 305 | 306 | Graphing 307 | -------- 308 | 309 | If you have installed ``matplotlib``, you can use a result set's 310 | ``.plot()``, ``.pie()``, and ``.bar()`` methods for quick plotting 311 | 312 | .. code-block:: python 313 | 314 | In[5]: result = %sql SELECT title, totalwords FROM work WHERE genretype = 'c' 315 | 316 | In[6]: %matplotlib inline 317 | 318 | In[7]: result.pie() 319 | 320 | .. image:: https://raw.github.com/catherinedevlin/ipython-sql/master/examples/wordcount.png 321 | :alt: pie chart of word count of Shakespeare's comedies 322 | 323 | Dumping 324 | ------- 325 | 326 | Result sets come with a ``.csv(filename=None)`` method. This generates 327 | comma-separated text either as a return value (if ``filename`` is not 328 | specified) or in a file of the given name. 329 | 330 | .. code-block:: python 331 | 332 | In[8]: result = %sql SELECT title, totalwords FROM work WHERE genretype = 'c' 333 | 334 | In[9]: result.csv(filename='work.csv') 335 | 336 | PostgreSQL features 337 | ------------------- 338 | 339 | ``psql``-style "backslash" `meta-commands`_ commands (``\d``, ``\dt``, etc.) 340 | are provided by `PGSpecial`_. Example: 341 | 342 | .. code-block:: python 343 | 344 | In[9]: %sql \d 345 | 346 | .. _PGSpecial: https://pypi.python.org/pypi/pgspecial 347 | 348 | .. _meta-commands: https://www.postgresql.org/docs/9.6/static/app-psql.html#APP-PSQL-META-COMMANDS 349 | 350 | 351 | Options 352 | ------- 353 | 354 | ``-l`` / ``--connections`` 355 | List all active connections 356 | 357 | ``-x`` / ``--close `` 358 | Close named connection 359 | 360 | ``-c`` / ``--creator `` 361 | Specify creator function for new connection 362 | 363 | ``-s`` / ``--section `` 364 | Section of dsn_file to be used for generating a connection string 365 | 366 | ``-p`` / ``--persist`` 367 | Create a table name in the database from the named DataFrame 368 | 369 | ``--append`` 370 | Like ``--persist``, but appends to the table if it already exists 371 | 372 | ``-a`` / ``--connection_arguments <"{connection arguments}">`` 373 | Specify dictionary of connection arguments to pass to SQL driver 374 | 375 | ``-f`` / ``--file `` 376 | Run SQL from file at this path 377 | 378 | Caution 379 | ------- 380 | 381 | Comments 382 | ~~~~~~~~ 383 | 384 | Because ipyton-sql accepts ``--``-delimited options like ``--persist``, but ``--`` 385 | is also the syntax to denote a SQL comment, the parser needs to make some assumptions. 386 | 387 | - If you try to pass an unsupported argument, like ``--lutefisk``, it will 388 | be interpreted as a SQL comment and will not throw an unsupported argument 389 | exception. 390 | - If the SQL statement begins with a first-line comment that looks like one 391 | of the accepted arguments - like ``%sql --persist is great!`` - it will be 392 | parsed like an argument, not a comment. Moving the comment to the second 393 | line or later will avoid this. 394 | 395 | Installing 396 | ---------- 397 | 398 | Install the latest release with:: 399 | 400 | pip install ipython-sql 401 | 402 | or download from https://github.com/catherinedevlin/ipython-sql and:: 403 | 404 | cd ipython-sql 405 | sudo python setup.py install 406 | 407 | Development 408 | ----------- 409 | 410 | https://github.com/catherinedevlin/ipython-sql 411 | 412 | Credits 413 | ------- 414 | 415 | - Matthias Bussonnier for help with configuration 416 | - Olivier Le Thanh Duong for ``%config`` fixes and improvements 417 | - Distribute_ 418 | - Buildout_ 419 | - modern-package-template_ 420 | - Mike Wilson for bind variable code 421 | - Thomas Kluyver and Steve Holden for debugging help 422 | - Berton Earnshaw for DSN connection syntax 423 | - Bruno Harbulot for DSN example 424 | - Andrés Celis for SQL Server bugfix 425 | - Michael Erasmus for DataFrame truth bugfix 426 | - Noam Finkelstein for README clarification 427 | - Xiaochuan Yu for `<<` operator, syntax colorization 428 | - Amjith Ramanujam for PGSpecial and incorporating it here 429 | - Alexander Maznev for better arg parsing, connections accepting specified creator 430 | - Jonathan Larkin for configurable displaycon 431 | - Jared Moore for ``connection-arguments`` support 432 | - Gilbert Brault for ``--append`` 433 | - Lucas Zeer for multi-line bugfixes for var substitution, ``<<`` 434 | - vkk800 for ``--file`` 435 | - Jens Albrecht for MySQL DatabaseError bugfix 436 | - meihkv for connection-closing bugfix 437 | - Abhinav C for SQLAlchemy 2.0 compatibility 438 | 439 | .. _Distribute: http://pypi.python.org/pypi/distribute 440 | .. _Buildout: http://www.buildout.org/ 441 | .. _modern-package-template: http://pypi.python.org/pypi/modern-package-template 442 | .. _JupySQL: https://github.com/ploomber/jupysql 443 | -------------------------------------------------------------------------------- /bootstrap.py: -------------------------------------------------------------------------------- 1 | ############################################################################## 2 | # 3 | # Copyright (c) 2006 Zope Corporation and Contributors. 4 | # All Rights Reserved. 5 | # 6 | # This software is subject to the provisions of the Zope Public License, 7 | # Version 2.1 (ZPL). A copy of the ZPL should accompany this distribution. 8 | # THIS SOFTWARE IS PROVIDED "AS IS" AND ANY AND ALL EXPRESS OR IMPLIED 9 | # WARRANTIES ARE DISCLAIMED, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED 10 | # WARRANTIES OF TITLE, MERCHANTABILITY, AGAINST INFRINGEMENT, AND FITNESS 11 | # FOR A PARTICULAR PURPOSE. 12 | # 13 | ############################################################################## 14 | """Bootstrap a buildout-based project 15 | 16 | Simply run this script in a directory containing a buildout.cfg. 17 | The script accepts buildout command-line options, so you can 18 | use the -c option to specify an alternate configuration file. 19 | 20 | $Id: bootstrap.py 102545 2009-08-06 14:49:47Z chrisw $ 21 | """ 22 | 23 | import os, shutil, sys, tempfile, urllib2 24 | from optparse import OptionParser 25 | 26 | tmpeggs = tempfile.mkdtemp() 27 | 28 | is_jython = sys.platform.startswith('java') 29 | 30 | # parsing arguments 31 | parser = OptionParser() 32 | parser.add_option("-v", "--version", dest="version", 33 | help="use a specific zc.buildout version") 34 | parser.add_option("-d", "--distribute", 35 | action="store_true", dest="distribute", default=True, 36 | help="Use Disribute rather than Setuptools.") 37 | 38 | options, args = parser.parse_args() 39 | 40 | if options.version is not None: 41 | VERSION = '==%s' % options.version 42 | else: 43 | VERSION = '' 44 | 45 | USE_DISTRIBUTE = options.distribute 46 | args = args + ['bootstrap'] 47 | 48 | to_reload = False 49 | try: 50 | import pkg_resources 51 | if not hasattr(pkg_resources, '_distribute'): 52 | to_reload = True 53 | raise ImportError 54 | except ImportError: 55 | ez = {} 56 | if USE_DISTRIBUTE: 57 | exec urllib2.urlopen('http://python-distribute.org/distribute_setup.py' 58 | ).read() in ez 59 | ez['use_setuptools'](to_dir=tmpeggs, download_delay=0, no_fake=True) 60 | else: 61 | exec urllib2.urlopen('http://peak.telecommunity.com/dist/ez_setup.py' 62 | ).read() in ez 63 | ez['use_setuptools'](to_dir=tmpeggs, download_delay=0) 64 | 65 | if to_reload: 66 | reload(pkg_resources) 67 | else: 68 | import pkg_resources 69 | 70 | if sys.platform == 'win32': 71 | def quote(c): 72 | if ' ' in c: 73 | return '"%s"' % c # work around spawn lamosity on windows 74 | else: 75 | return c 76 | else: 77 | def quote (c): 78 | return c 79 | 80 | cmd = 'from setuptools.command.easy_install import main; main()' 81 | ws = pkg_resources.working_set 82 | 83 | if USE_DISTRIBUTE: 84 | requirement = 'distribute' 85 | else: 86 | requirement = 'setuptools' 87 | 88 | if is_jython: 89 | import subprocess 90 | 91 | assert subprocess.Popen([sys.executable] + ['-c', quote(cmd), '-mqNxd', 92 | quote(tmpeggs), 'zc.buildout' + VERSION], 93 | env=dict(os.environ, 94 | PYTHONPATH= 95 | ws.find(pkg_resources.Requirement.parse(requirement)).location 96 | ), 97 | ).wait() == 0 98 | 99 | else: 100 | assert os.spawnle( 101 | os.P_WAIT, sys.executable, quote (sys.executable), 102 | '-c', quote (cmd), '-mqNxd', quote (tmpeggs), 'zc.buildout' + VERSION, 103 | dict(os.environ, 104 | PYTHONPATH= 105 | ws.find(pkg_resources.Requirement.parse(requirement)).location 106 | ), 107 | ) == 0 108 | 109 | ws.add_entry(tmpeggs) 110 | ws.require('zc.buildout' + VERSION) 111 | import zc.buildout.buildout 112 | zc.buildout.buildout.main(args) 113 | shutil.rmtree(tmpeggs) 114 | -------------------------------------------------------------------------------- /buildout.cfg: -------------------------------------------------------------------------------- 1 | [buildout] 2 | parts = python scripts 3 | develop = . 4 | eggs = ipython-sql 5 | 6 | [python] 7 | recipe = zc.recipe.egg 8 | interpreter = python 9 | eggs = ${buildout:eggs} 10 | 11 | [scripts] 12 | recipe = zc.recipe.egg:scripts 13 | eggs = ${buildout:eggs} 14 | -------------------------------------------------------------------------------- /examples/wordcount.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/catherinedevlin/ipython-sql/eb274844b4a619463149e0d57df705e1bba47635/examples/wordcount.png -------------------------------------------------------------------------------- /examples/writers.ipynb: -------------------------------------------------------------------------------- 1 | { 2 | "cells": [ 3 | { 4 | "cell_type": "code", 5 | "execution_count": 2, 6 | "metadata": { 7 | "collapsed": false, 8 | "jupyter": { 9 | "outputs_hidden": false 10 | } 11 | }, 12 | "outputs": [], 13 | "source": [ 14 | "%load_ext sql" 15 | ] 16 | }, 17 | { 18 | "cell_type": "code", 19 | "execution_count": 3, 20 | "metadata": { 21 | "collapsed": false, 22 | "jupyter": { 23 | "outputs_hidden": false 24 | } 25 | }, 26 | "outputs": [], 27 | "source": [ 28 | "%sql sqlite://" 29 | ] 30 | }, 31 | { 32 | "cell_type": "code", 33 | "execution_count": 3, 34 | "metadata": { 35 | "collapsed": false, 36 | "jupyter": { 37 | "outputs_hidden": false 38 | } 39 | }, 40 | "outputs": [ 41 | { 42 | "name": "stdout", 43 | "output_type": "stream", 44 | "text": [ 45 | " * sqlite://\n", 46 | "Done.\n", 47 | "1 rows affected.\n", 48 | "1 rows affected.\n" 49 | ] 50 | }, 51 | { 52 | "data": { 53 | "text/plain": [ 54 | "[]" 55 | ] 56 | }, 57 | "execution_count": 3, 58 | "metadata": {}, 59 | "output_type": "execute_result" 60 | } 61 | ], 62 | "source": [ 63 | "%%sql\n", 64 | "CREATE TABLE writer (first_name, last_name, year_of_death);\n", 65 | "INSERT INTO writer VALUES ('William', 'Shakespeare', 1616);\n", 66 | "INSERT INTO writer VALUES ('Bertold', 'Brecht', 1956);" 67 | ] 68 | }, 69 | { 70 | "cell_type": "code", 71 | "execution_count": 4, 72 | "metadata": { 73 | "collapsed": false, 74 | "jupyter": { 75 | "outputs_hidden": false 76 | } 77 | }, 78 | "outputs": [ 79 | { 80 | "name": "stdout", 81 | "output_type": "stream", 82 | "text": [ 83 | " * sqlite://\n", 84 | "Done.\n" 85 | ] 86 | }, 87 | { 88 | "data": { 89 | "text/html": [ 90 | "\n", 91 | " \n", 92 | " \n", 93 | " \n", 94 | " \n", 95 | " \n", 96 | " \n", 97 | " \n", 98 | " \n", 99 | " \n", 100 | " \n", 101 | " \n", 102 | " \n", 103 | " \n", 104 | " \n", 105 | " \n", 106 | "
first_namelast_nameyear_of_death
WilliamShakespeare1616
BertoldBrecht1956
" 107 | ], 108 | "text/plain": [ 109 | "[('William', 'Shakespeare', 1616), ('Bertold', 'Brecht', 1956)]" 110 | ] 111 | }, 112 | "execution_count": 4, 113 | "metadata": {}, 114 | "output_type": "execute_result" 115 | } 116 | ], 117 | "source": [ 118 | "%sql select * from writer" 119 | ] 120 | }, 121 | { 122 | "cell_type": "code", 123 | "execution_count": 5, 124 | "metadata": { 125 | "collapsed": false, 126 | "jupyter": { 127 | "outputs_hidden": false 128 | } 129 | }, 130 | "outputs": [ 131 | { 132 | "name": "stdout", 133 | "output_type": "stream", 134 | "text": [ 135 | " * sqlite://\n", 136 | "Done.\n", 137 | "Returning data to local variable writers\n" 138 | ] 139 | } 140 | ], 141 | "source": [ 142 | "%%sql writers << select first_name, year_of_death\n", 143 | "from writer" 144 | ] 145 | }, 146 | { 147 | "cell_type": "code", 148 | "execution_count": 6, 149 | "metadata": {}, 150 | "outputs": [ 151 | { 152 | "data": { 153 | "text/html": [ 154 | "\n", 155 | " \n", 156 | " \n", 157 | " \n", 158 | " \n", 159 | " \n", 160 | " \n", 161 | " \n", 162 | " \n", 163 | " \n", 164 | " \n", 165 | " \n", 166 | " \n", 167 | "
first_nameyear_of_death
William1616
Bertold1956
" 168 | ], 169 | "text/plain": [ 170 | "[('William', 1616), ('Bertold', 1956)]" 171 | ] 172 | }, 173 | "execution_count": 6, 174 | "metadata": {}, 175 | "output_type": "execute_result" 176 | } 177 | ], 178 | "source": [ 179 | "writers" 180 | ] 181 | }, 182 | { 183 | "cell_type": "code", 184 | "execution_count": 7, 185 | "metadata": {}, 186 | "outputs": [], 187 | "source": [ 188 | "var = 'last_name'" 189 | ] 190 | }, 191 | { 192 | "cell_type": "code", 193 | "execution_count": 8, 194 | "metadata": {}, 195 | "outputs": [ 196 | { 197 | "name": "stdout", 198 | "output_type": "stream", 199 | "text": [ 200 | " * sqlite://\n", 201 | "Done.\n" 202 | ] 203 | }, 204 | { 205 | "data": { 206 | "text/html": [ 207 | "\n", 208 | " \n", 209 | " \n", 210 | " \n", 211 | " \n", 212 | " \n", 213 | " \n", 214 | " \n", 215 | " \n", 216 | " \n", 217 | " \n", 218 | "
first_namelast_nameyear_of_death
BertoldBrecht1956
" 219 | ], 220 | "text/plain": [ 221 | "[('Bertold', 'Brecht', 1956)]" 222 | ] 223 | }, 224 | "execution_count": 8, 225 | "metadata": {}, 226 | "output_type": "execute_result" 227 | } 228 | ], 229 | "source": [ 230 | "%sql select * from writer where {var} = 'Brecht'" 231 | ] 232 | }, 233 | { 234 | "cell_type": "code", 235 | "execution_count": 9, 236 | "metadata": {}, 237 | "outputs": [ 238 | { 239 | "name": "stdout", 240 | "output_type": "stream", 241 | "text": [ 242 | " * sqlite://\n", 243 | "Done.\n" 244 | ] 245 | }, 246 | { 247 | "data": { 248 | "text/html": [ 249 | "\n", 250 | " \n", 251 | " \n", 252 | " \n", 253 | " \n", 254 | " \n", 255 | " \n", 256 | " \n", 257 | " \n", 258 | " \n", 259 | " \n", 260 | "
first_namelast_nameyear_of_death
BertoldBrecht1956
" 261 | ], 262 | "text/plain": [ 263 | "[('Bertold', 'Brecht', 1956)]" 264 | ] 265 | }, 266 | "execution_count": 9, 267 | "metadata": {}, 268 | "output_type": "execute_result" 269 | } 270 | ], 271 | "source": [ 272 | "%%sql select * from writer \n", 273 | "where {var} = 'Brecht'" 274 | ] 275 | }, 276 | { 277 | "cell_type": "code", 278 | "execution_count": null, 279 | "metadata": {}, 280 | "outputs": [], 281 | "source": [] 282 | } 283 | ], 284 | "metadata": { 285 | "kernelspec": { 286 | "display_name": "Python 3", 287 | "language": "python", 288 | "name": "python3" 289 | }, 290 | "language_info": { 291 | "codemirror_mode": { 292 | "name": "ipython", 293 | "version": 3 294 | }, 295 | "file_extension": ".py", 296 | "mimetype": "text/x-python", 297 | "name": "python", 298 | "nbconvert_exporter": "python", 299 | "pygments_lexer": "ipython3", 300 | "version": "3.8.0" 301 | } 302 | }, 303 | "nbformat": 4, 304 | "nbformat_minor": 4 305 | } 306 | -------------------------------------------------------------------------------- /examples/writers.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/catherinedevlin/ipython-sql/eb274844b4a619463149e0d57df705e1bba47635/examples/writers.png -------------------------------------------------------------------------------- /ipython-sql.wpr: -------------------------------------------------------------------------------- 1 | #!wing 2 | #!version=5.0 3 | ################################################################## 4 | # Wing IDE project file # 5 | ################################################################## 6 | [project attributes] 7 | proj.directory-list = [{'dirloc': loc('.'), 8 | 'excludes': (), 9 | 'filter': '*', 10 | 'include_hidden': False, 11 | 'recursive': True, 12 | 'watch_for_changes': True}] 13 | proj.file-type = 'shared' 14 | -------------------------------------------------------------------------------- /requirements-dev.txt: -------------------------------------------------------------------------------- 1 | psycopg2-binary 2 | pandas 3 | pytest 4 | wheel 5 | twine 6 | readme-renderer 7 | black 8 | isort 9 | -------------------------------------------------------------------------------- /requirements.txt: -------------------------------------------------------------------------------- 1 | prettytable 2 | ipython 3 | sqlalchemy>=2.0 4 | sqlparse 5 | six 6 | ipython-genutils 7 | traitlets 8 | matplotlib 9 | -------------------------------------------------------------------------------- /run_tests.sh: -------------------------------------------------------------------------------- 1 | #!/bin/bash 2 | ipython -c "import pytest; pytest.main(['.', '-x', '--pdb'])" 3 | # Insert breakpoints with `import pytest; pytest.set_trace()` 4 | -------------------------------------------------------------------------------- /setup.py: -------------------------------------------------------------------------------- 1 | import os 2 | from io import open 3 | 4 | from setuptools import find_packages, setup 5 | 6 | here = os.path.abspath(os.path.dirname(__file__)) 7 | README = open(os.path.join(here, "README.rst"), encoding="utf-8").read() 8 | NEWS = open(os.path.join(here, "NEWS.rst"), encoding="utf-8").read() 9 | 10 | 11 | version = "0.5.0" 12 | 13 | install_requires = [ 14 | "prettytable", 15 | "ipython", 16 | "sqlalchemy>=2.0", 17 | "sqlparse", 18 | "six", 19 | "ipython-genutils", 20 | ] 21 | 22 | 23 | setup( 24 | name="ipython-sql", 25 | version=version, 26 | description="RDBMS access via IPython", 27 | long_description=README + "\n\n" + NEWS, 28 | long_description_content_type="text/x-rst", 29 | classifiers=[ 30 | "Development Status :: 3 - Alpha", 31 | "Environment :: Console", 32 | "License :: OSI Approved :: MIT License", 33 | "Topic :: Database", 34 | "Topic :: Database :: Front-Ends", 35 | "Programming Language :: Python :: 3", 36 | ], 37 | keywords="database ipython postgresql mysql", 38 | author="Catherine Devlin", 39 | author_email="catherine.devlin@gmail.com", 40 | url="https://github.com/catherinedevlin/ipython-sql", 41 | project_urls={ 42 | "Source": "https://github.com/catherinedevlin/ipython-sql", 43 | }, 44 | license="MIT", 45 | packages=find_packages("src"), 46 | package_dir={"": "src"}, 47 | include_package_data=True, 48 | zip_safe=False, 49 | install_requires=install_requires, 50 | ) 51 | -------------------------------------------------------------------------------- /src/sql/__init__.py: -------------------------------------------------------------------------------- 1 | from .magic import * 2 | -------------------------------------------------------------------------------- /src/sql/column_guesser.py: -------------------------------------------------------------------------------- 1 | """ 2 | Splits tabular data in the form of a list of rows into columns; 3 | makes guesses about the role of each column for plotting purposes 4 | (X values, Y values, and text labels). 5 | """ 6 | 7 | 8 | class Column(list): 9 | """Store a column of tabular data; record its name and whether it is numeric""" 10 | is_quantity = True 11 | name = "" 12 | 13 | def __init__(self, *arg, **kwarg): 14 | pass 15 | 16 | 17 | def is_quantity(val): 18 | """Is ``val`` a quantity (int, float, datetime, etc) (not str, bool)? 19 | 20 | Relies on presence of __sub__. 21 | """ 22 | return hasattr(val, "__sub__") 23 | 24 | 25 | class ColumnGuesserMixin(object): 26 | """ 27 | plot: [x, y, y...], y 28 | pie: ... y 29 | """ 30 | 31 | def __init__(self): 32 | self.keys = None 33 | 34 | def _build_columns(self): 35 | self.columns = [Column() for col in self.keys] 36 | for row in self: 37 | for (col_idx, col_val) in enumerate(row): 38 | col = self.columns[col_idx] 39 | col.append(col_val) 40 | if (col_val is not None) and (not is_quantity(col_val)): 41 | col.is_quantity = False 42 | 43 | for (idx, key_name) in enumerate(self.keys): 44 | self.columns[idx].name = key_name 45 | 46 | self.x = Column() 47 | self.ys = [] 48 | 49 | def _get_y(self): 50 | for idx in range(len(self.columns) - 1, -1, -1): 51 | if self.columns[idx].is_quantity: 52 | self.ys.insert(0, self.columns.pop(idx)) 53 | return True 54 | 55 | def _get_x(self): 56 | for idx in range(len(self.columns)): 57 | if self.columns[idx].is_quantity: 58 | self.x = self.columns.pop(idx) 59 | return True 60 | 61 | def _get_xlabel(self, xlabel_sep=" "): 62 | self.xlabels = [] 63 | if self.columns: 64 | for row_idx in range(len(self.columns[0])): 65 | self.xlabels.append( 66 | xlabel_sep.join(str(c[row_idx]) for c in self.columns) 67 | ) 68 | self.xlabel = ", ".join(c.name for c in self.columns) 69 | 70 | def _guess_columns(self): 71 | self._build_columns() 72 | self._get_y() 73 | if not self.ys: 74 | raise AttributeError("No quantitative columns found for chart") 75 | 76 | def guess_pie_columns(self, xlabel_sep=" "): 77 | """ 78 | Assigns x, y, and x labels from the data set for a pie chart. 79 | 80 | Pie charts simply use the last quantity column as 81 | the pie slice size, and everything else as the 82 | pie slice labels. 83 | """ 84 | self._guess_columns() 85 | self._get_xlabel(xlabel_sep) 86 | 87 | def guess_plot_columns(self): 88 | """ 89 | Assigns ``x`` and ``y`` series from the data set for a plot. 90 | 91 | Plots use: 92 | the rightmost quantity column as a Y series 93 | optionally, the leftmost quantity column as the X series 94 | any other quantity columns as additional Y series 95 | """ 96 | self._guess_columns() 97 | self._get_x() 98 | while self._get_y(): 99 | pass 100 | -------------------------------------------------------------------------------- /src/sql/connection.py: -------------------------------------------------------------------------------- 1 | import os 2 | import traceback 3 | 4 | import sqlalchemy 5 | 6 | 7 | class ConnectionError(Exception): 8 | pass 9 | 10 | 11 | def rough_dict_get(dct, sought, default=None): 12 | """ 13 | Like dct.get(sought), but any key containing sought will do. 14 | 15 | If there is a `@` in sought, seek each piece separately. 16 | This lets `me@server` match `me:***@myserver/db` 17 | """ 18 | 19 | sought = sought.split("@") 20 | for (key, val) in dct.items(): 21 | if not any(s.lower() not in key.lower() for s in sought): 22 | return val 23 | return default 24 | 25 | 26 | class Connection(object): 27 | current = None 28 | connections = {} 29 | 30 | @classmethod 31 | def tell_format(cls): 32 | return """Connection info needed in SQLAlchemy format, example: 33 | postgresql://username:password@hostname/dbname 34 | or an existing connection: %s""" % str( 35 | cls.connections.keys() 36 | ) 37 | 38 | def __init__(self, connect_str=None, connect_args={}, creator=None): 39 | try: 40 | if creator: 41 | engine = sqlalchemy.create_engine( 42 | connect_str, connect_args=connect_args, creator=creator 43 | ) 44 | else: 45 | engine = sqlalchemy.create_engine( 46 | connect_str, connect_args=connect_args 47 | ) 48 | except Exception as ex: # TODO: bare except; but what's an ArgumentError? 49 | print(traceback.format_exc()) 50 | print(self.tell_format()) 51 | raise 52 | self.url = engine.url 53 | self.dialect = engine.url.get_dialect() 54 | self.name = self.assign_name(engine) 55 | self.internal_connection = engine.connect() 56 | self.connections[repr(self.url)] = self 57 | self.connect_args = connect_args 58 | Connection.current = self 59 | 60 | @classmethod 61 | def set(cls, descriptor, displaycon, connect_args={}, creator=None): 62 | """Sets the current database connection""" 63 | 64 | if descriptor: 65 | if isinstance(descriptor, Connection): 66 | cls.current = descriptor 67 | else: 68 | existing = rough_dict_get(cls.connections, descriptor) 69 | # http://docs.sqlalchemy.org/en/rel_0_9/core/engines.html#custom-dbapi-connect-arguments 70 | cls.current = existing or Connection(descriptor, connect_args, creator) 71 | else: 72 | 73 | if cls.connections: 74 | if displaycon: 75 | print(cls.connection_list()) 76 | else: 77 | if os.getenv("DATABASE_URL"): 78 | cls.current = Connection( 79 | os.getenv("DATABASE_URL"), connect_args, creator 80 | ) 81 | else: 82 | raise ConnectionError( 83 | "Environment variable $DATABASE_URL not set, and no connect string given." 84 | ) 85 | return cls.current 86 | 87 | @classmethod 88 | def assign_name(cls, engine): 89 | name = "%s@%s" % (engine.url.username or "", engine.url.database) 90 | return name 91 | 92 | @classmethod 93 | def connection_list(cls): 94 | result = [] 95 | for key in sorted(cls.connections): 96 | engine_url = cls.connections[ 97 | key 98 | ].url # type: sqlalchemy.engine.url.URL 99 | if cls.connections[key] == cls.current: 100 | template = " * {}" 101 | else: 102 | template = " {}" 103 | result.append(template.format(engine_url.__repr__())) 104 | return "\n".join(result) 105 | 106 | @classmethod 107 | def close(cls, descriptor): 108 | if isinstance(descriptor, Connection): 109 | conn = descriptor 110 | else: 111 | conn = cls.connections.get(descriptor) or cls.connections.get( 112 | descriptor.lower() 113 | ) 114 | if not conn: 115 | raise Exception( 116 | "Could not close connection because it was not found amongst these: %s" 117 | % str(cls.connections.keys()) 118 | ) 119 | cls.connections.pop(str(conn.url)) 120 | conn.internal_connection.close() 121 | -------------------------------------------------------------------------------- /src/sql/magic.py: -------------------------------------------------------------------------------- 1 | import json 2 | import re 3 | import traceback 4 | 5 | from IPython.core.magic import ( 6 | Magics, 7 | cell_magic, 8 | line_magic, 9 | magics_class, 10 | needs_local_scope, 11 | ) 12 | from IPython.core.magic_arguments import argument, magic_arguments, parse_argstring 13 | from sqlalchemy.exc import OperationalError, ProgrammingError, DatabaseError 14 | 15 | import sql.connection 16 | import sql.parse 17 | import sql.run 18 | 19 | try: 20 | from traitlets.config.configurable import Configurable 21 | from traitlets import Bool, Int, Unicode 22 | except ImportError: 23 | from IPython.config.configurable import Configurable 24 | from IPython.utils.traitlets import Bool, Int, Unicode 25 | try: 26 | from pandas.core.frame import DataFrame, Series 27 | except ImportError: 28 | DataFrame = None 29 | Series = None 30 | 31 | 32 | @magics_class 33 | class SqlMagic(Magics, Configurable): 34 | """Runs SQL statement on a database, specified by SQLAlchemy connect string. 35 | 36 | Provides the %%sql magic.""" 37 | 38 | displaycon = Bool(True, config=True, help="Show connection string after execute") 39 | autolimit = Int( 40 | 0, 41 | config=True, 42 | allow_none=True, 43 | help="Automatically limit the size of the returned result sets", 44 | ) 45 | style = Unicode( 46 | "DEFAULT", 47 | config=True, 48 | help="Set the table printing style to any of prettytable's defined styles " 49 | "(currently DEFAULT, MSWORD_FRIENDLY, PLAIN_COLUMNS, RANDOM)", 50 | ) 51 | short_errors = Bool( 52 | True, 53 | config=True, 54 | help="Don't display the full traceback on SQL Programming Error", 55 | ) 56 | displaylimit = Int( 57 | None, 58 | config=True, 59 | allow_none=True, 60 | help="Automatically limit the number of rows displayed (full result set is still stored)", 61 | ) 62 | autopandas = Bool( 63 | False, 64 | config=True, 65 | help="Return Pandas DataFrames instead of regular result sets", 66 | ) 67 | column_local_vars = Bool( 68 | False, config=True, help="Return data into local variables from column names" 69 | ) 70 | feedback = Bool(True, config=True, help="Print number of rows affected by DML") 71 | dsn_filename = Unicode( 72 | "odbc.ini", 73 | config=True, 74 | help="Path to DSN file. " 75 | "When the first argument is of the form [section], " 76 | "a sqlalchemy connection string is formed from the " 77 | "matching section in the DSN file.", 78 | ) 79 | autocommit = Bool(True, config=True, help="Set autocommit mode") 80 | 81 | def __init__(self, shell): 82 | Configurable.__init__(self, config=shell.config) 83 | Magics.__init__(self, shell=shell) 84 | 85 | # Add ourselves to the list of module configurable via %config 86 | self.shell.configurables.append(self) 87 | 88 | @needs_local_scope 89 | @line_magic("sql") 90 | @cell_magic("sql") 91 | @magic_arguments() 92 | @argument("line", default="", nargs="*", type=str, help="sql") 93 | @argument( 94 | "-l", "--connections", action="store_true", help="list active connections" 95 | ) 96 | @argument("-x", "--close", type=str, help="close a session by name") 97 | @argument( 98 | "-c", "--creator", type=str, help="specify creator function for new connection" 99 | ) 100 | @argument( 101 | "-s", 102 | "--section", 103 | type=str, 104 | help="section of dsn_file to be used for generating a connection string", 105 | ) 106 | @argument( 107 | "-p", 108 | "--persist", 109 | action="store_true", 110 | help="create a table name in the database from the named DataFrame", 111 | ) 112 | @argument( 113 | "--append", 114 | action="store_true", 115 | help="create, or append to, a table name in the database from the named DataFrame", 116 | ) 117 | @argument( 118 | "-a", 119 | "--connection_arguments", 120 | type=str, 121 | help="specify dictionary of connection arguments to pass to SQL driver", 122 | ) 123 | @argument("-f", "--file", type=str, help="Run SQL from file at this path") 124 | def execute(self, line="", cell="", local_ns=None): 125 | """Runs SQL statement against a database, specified by SQLAlchemy connect string. 126 | 127 | If no database connection has been established, first word 128 | should be a SQLAlchemy connection string, or the user@db name 129 | of an established connection. 130 | 131 | Examples:: 132 | 133 | %%sql postgresql://me:mypw@localhost/mydb 134 | SELECT * FROM mytable 135 | 136 | %%sql me@mydb 137 | DELETE FROM mytable 138 | 139 | %%sql 140 | DROP TABLE mytable 141 | 142 | SQLAlchemy connect string syntax examples: 143 | 144 | postgresql://me:mypw@localhost/mydb 145 | sqlite:// 146 | mysql+pymysql://me:mypw@localhost/mydb 147 | 148 | """ 149 | # Parse variables (words wrapped in {}) for %%sql magic (for %sql this is done automatically) 150 | if local_ns is None: 151 | local_ns = {} 152 | cell = self.shell.var_expand(cell) 153 | line = sql.parse.without_sql_comment(parser=self.execute.parser, line=line) 154 | args = parse_argstring(self.execute, line) 155 | if args.connections: 156 | return sql.connection.Connection.connections 157 | elif args.close: 158 | return sql.connection.Connection.close(args.close) 159 | 160 | # save globals and locals, so they can be referenced in bind vars 161 | user_ns = self.shell.user_ns.copy() 162 | user_ns.update(local_ns) 163 | 164 | command_text = " ".join(args.line) + "\n" + cell 165 | 166 | if args.file: 167 | with open(args.file, "r") as infile: 168 | command_text = infile.read() + "\n" + command_text 169 | 170 | parsed = sql.parse.parse(command_text, self) 171 | 172 | connect_str = parsed["connection"] 173 | if args.section: 174 | connect_str = sql.parse.connection_from_dsn_section(args.section, self) 175 | 176 | if args.connection_arguments: 177 | try: 178 | # check for string delineators, we need to strip them for json parse 179 | raw_args = args.connection_arguments 180 | if len(raw_args) > 1: 181 | targets = ['"', "'"] 182 | head = raw_args[0] 183 | tail = raw_args[-1] 184 | if head in targets and head == tail: 185 | raw_args = raw_args[1:-1] 186 | args.connection_arguments = json.loads(raw_args) 187 | except Exception as e: 188 | print(traceback.format_exc()) 189 | raise e 190 | else: 191 | args.connection_arguments = {} 192 | if args.creator: 193 | args.creator = user_ns[args.creator] 194 | 195 | try: 196 | conn = sql.connection.Connection.set( 197 | connect_str, 198 | displaycon=self.displaycon, 199 | connect_args=args.connection_arguments, 200 | creator=args.creator, 201 | ) 202 | # Rollback just in case there was an error in previous statement 203 | conn.internal_connection.rollback() 204 | except Exception: 205 | print(traceback.format_exc()) 206 | print(sql.connection.Connection.tell_format()) 207 | return None 208 | 209 | if args.persist: 210 | return self._persist_dataframe(parsed["sql"], conn, user_ns, append=False) 211 | 212 | if args.append: 213 | return self._persist_dataframe(parsed["sql"], conn, user_ns, append=True) 214 | 215 | if not parsed["sql"]: 216 | return 217 | 218 | try: 219 | result = sql.run.run(conn, parsed["sql"], self, user_ns) 220 | 221 | if ( 222 | result is not None 223 | and not isinstance(result, str) 224 | and self.column_local_vars 225 | ): 226 | # Instead of returning values, set variables directly in the 227 | # user's namespace. Variable names given by column names 228 | 229 | if self.autopandas: 230 | keys = result.keys() 231 | else: 232 | keys = result.keys 233 | result = result.dict() 234 | 235 | if self.feedback: 236 | print( 237 | "Returning data to local variables [{}]".format(", ".join(keys)) 238 | ) 239 | 240 | self.shell.user_ns.update(result) 241 | 242 | return None 243 | else: 244 | 245 | if parsed["result_var"]: 246 | result_var = parsed["result_var"] 247 | print("Returning data to local variable {}".format(result_var)) 248 | self.shell.user_ns.update({result_var: result}) 249 | return None 250 | 251 | # Return results into the default ipython _ variable 252 | return result 253 | 254 | # JA: added DatabaseError for MySQL 255 | except (ProgrammingError, OperationalError, DatabaseError) as e: 256 | # Sqlite apparently return all errors as OperationalError :/ 257 | if self.short_errors: 258 | print(e) 259 | else: 260 | print(traceback.format_exc()) 261 | raise e 262 | 263 | legal_sql_identifier = re.compile(r"^[A-Za-z0-9#_$]+") 264 | 265 | def _persist_dataframe(self, raw, conn, user_ns, append=False): 266 | """Implements PERSIST, which writes a DataFrame to the RDBMS""" 267 | if not DataFrame: 268 | raise ImportError("Must `pip install pandas` to use DataFrames") 269 | 270 | frame_name = raw.strip(";") 271 | 272 | # Get the DataFrame from the user namespace 273 | if not frame_name: 274 | raise SyntaxError("Syntax: %sql --persist ") 275 | try: 276 | frame = eval(frame_name, user_ns) 277 | except SyntaxError: 278 | raise SyntaxError("Syntax: %sql --persist ") 279 | if not isinstance(frame, DataFrame) and not isinstance(frame, Series): 280 | raise TypeError("%s is not a Pandas DataFrame or Series" % frame_name) 281 | 282 | # Make a suitable name for the resulting database table 283 | table_name = frame_name.lower() 284 | table_name = self.legal_sql_identifier.search(table_name).group(0) 285 | 286 | if_exists = "append" if append else "fail" 287 | frame.to_sql(table_name, conn.internal_connection.engine, if_exists=if_exists) 288 | return "Persisted %s" % table_name 289 | 290 | 291 | def load_ipython_extension(ip): 292 | """Load the extension in IPython.""" 293 | 294 | # this fails in both Firefox and Chrome for OS X. 295 | # I get the error: TypeError: IPython.CodeCell.config_defaults is undefined 296 | 297 | # js = "IPython.CodeCell.config_defaults.highlight_modes['magic_sql'] = {'reg':[/^%%sql/]};" 298 | # display_javascript(js, raw=True) 299 | ip.register_magics(SqlMagic) 300 | -------------------------------------------------------------------------------- /src/sql/parse.py: -------------------------------------------------------------------------------- 1 | import itertools 2 | import shlex 3 | from os.path import expandvars 4 | 5 | from six.moves import configparser as CP 6 | from sqlalchemy.engine.url import URL 7 | 8 | 9 | def connection_from_dsn_section(section, config): 10 | parser = CP.ConfigParser() 11 | parser.read(config.dsn_filename) 12 | cfg_dict = dict(parser.items(section)) 13 | return URL.create(**cfg_dict) 14 | 15 | 16 | def _connection_string(s, config): 17 | s = expandvars(s) # for environment variables 18 | if "@" in s or "://" in s: 19 | return s 20 | if s.startswith("[") and s.endswith("]"): 21 | section = s.lstrip("[").rstrip("]") 22 | parser = CP.ConfigParser() 23 | parser.read(config.dsn_filename) 24 | cfg_dict = dict(parser.items(section)) 25 | return str(URL(**cfg_dict)) 26 | return "" 27 | 28 | 29 | def parse(cell, config): 30 | """Extract connection info and result variable from SQL 31 | 32 | Please don't add any more syntax requiring 33 | special parsing. 34 | Instead, add @arguments to SqlMagic.execute. 35 | 36 | We're grandfathering the 37 | connection string and `<<` operator in. 38 | """ 39 | 40 | result = {"connection": "", "sql": "", "result_var": None} 41 | 42 | pieces = cell.split(None, 1) 43 | if not pieces: 44 | return result 45 | result["connection"] = _connection_string(pieces[0], config) 46 | if result["connection"]: 47 | if len(pieces) == 1: 48 | return result 49 | cell = pieces[1] 50 | 51 | pieces = cell.split(None, 2) 52 | if len(pieces) > 1 and pieces[1] == "<<": 53 | result["result_var"] = pieces[0] 54 | if len(pieces) == 2: 55 | return result 56 | cell = pieces[2] 57 | 58 | result["sql"] = cell 59 | return result 60 | 61 | 62 | def _option_strings_from_parser(parser): 63 | """Extracts the expected option strings (-a, --append, etc) from argparse parser 64 | 65 | Thanks Martijn Pieters 66 | https://stackoverflow.com/questions/28881456/how-can-i-list-all-registered-arguments-from-an-argumentparser-instance 67 | 68 | :param parser: [description] 69 | :type parser: IPython.core.magic_arguments.MagicArgumentParser 70 | """ 71 | opts = [a.option_strings for a in parser._actions] 72 | return list(itertools.chain.from_iterable(opts)) 73 | 74 | 75 | def without_sql_comment(parser, line): 76 | """Strips -- comment from a line 77 | 78 | The argparser unfortunately expects -- to precede an option, 79 | but in SQL that delineates a comment. So this removes comments 80 | so a line can safely be fed to the argparser. 81 | 82 | :param line: A line of SQL, possibly mixed with option strings 83 | :type line: str 84 | """ 85 | 86 | args = _option_strings_from_parser(parser) 87 | result = itertools.takewhile( 88 | lambda word: (not word.startswith("--")) or (word in args), 89 | shlex.split(line, posix=False), 90 | ) 91 | return " ".join(result) 92 | -------------------------------------------------------------------------------- /src/sql/run.py: -------------------------------------------------------------------------------- 1 | import codecs 2 | import csv 3 | import operator 4 | import os.path 5 | import re 6 | import traceback 7 | from functools import reduce 8 | 9 | import prettytable 10 | import six 11 | import sqlalchemy 12 | import sqlparse 13 | 14 | from .column_guesser import ColumnGuesserMixin 15 | 16 | try: 17 | from pgspecial.main import PGSpecial 18 | except ImportError: 19 | PGSpecial = None 20 | 21 | 22 | def unduplicate_field_names(field_names): 23 | """Append a number to duplicate field names to make them unique. """ 24 | res = [] 25 | for k in field_names: 26 | if k in res: 27 | i = 1 28 | while k + "_" + str(i) in res: 29 | i += 1 30 | k += "_" + str(i) 31 | res.append(k) 32 | return res 33 | 34 | 35 | class UnicodeWriter(object): 36 | """ 37 | A CSV writer which will write rows to CSV file "f", 38 | which is encoded in the given encoding. 39 | """ 40 | 41 | def __init__(self, f, dialect=csv.excel, encoding="utf-8", **kwds): 42 | # Redirect output to a queue 43 | self.queue = six.StringIO() 44 | self.writer = csv.writer(self.queue, dialect=dialect, **kwds) 45 | self.stream = f 46 | self.encoder = codecs.getincrementalencoder(encoding)() 47 | 48 | def writerow(self, row): 49 | if six.PY2: 50 | _row = [s.encode("utf-8") if hasattr(s, "encode") else s for s in row] 51 | else: 52 | _row = row 53 | self.writer.writerow(_row) 54 | # Fetch UTF-8 output from the queue ... 55 | data = self.queue.getvalue() 56 | if six.PY2: 57 | data = data.decode("utf-8") 58 | # ... and re-encode it into the target encoding 59 | data = self.encoder.encode(data) 60 | # write to the target stream 61 | self.stream.write(data) 62 | # empty queue 63 | self.queue.truncate(0) 64 | self.queue.seek(0) 65 | 66 | def writerows(self, rows): 67 | for row in rows: 68 | self.writerow(row) 69 | 70 | 71 | class CsvResultDescriptor(object): 72 | """Provides IPython Notebook-friendly output for the feedback after a ``.csv`` called.""" 73 | 74 | def __init__(self, file_path): 75 | self.file_path = file_path 76 | 77 | def __repr__(self): 78 | return "CSV results at %s" % os.path.join(os.path.abspath("."), self.file_path) 79 | 80 | def _repr_html_(self): 81 | return 'CSV results' % os.path.join( 82 | ".", "files", self.file_path 83 | ) 84 | 85 | 86 | def _nonbreaking_spaces(match_obj): 87 | """ 88 | Make spaces visible in HTML by replacing all `` `` with `` `` 89 | 90 | Call with a ``re`` match object. Retain group 1, replace group 2 91 | with nonbreaking speaces. 92 | """ 93 | spaces = " " * len(match_obj.group(2)) 94 | return "%s%s" % (match_obj.group(1), spaces) 95 | 96 | 97 | _cell_with_spaces_pattern = re.compile(r"()( {2,})") 98 | 99 | 100 | class ResultSet(list, ColumnGuesserMixin): 101 | """ 102 | Results of a SQL query. 103 | 104 | Can access rows listwise, or by string value of leftmost column. 105 | """ 106 | 107 | def __init__(self, sqlaproxy, config): 108 | self.config = config 109 | if sqlaproxy.returns_rows: 110 | self.keys = sqlaproxy.keys() 111 | if config.autolimit: 112 | list.__init__(self, sqlaproxy.fetchmany(size=config.autolimit)) 113 | else: 114 | list.__init__(self, sqlaproxy.fetchall()) 115 | self.field_names = unduplicate_field_names(self.keys) 116 | self.pretty = PrettyTable(self.field_names, style=prettytable.__dict__[config.style.upper()]) 117 | else: 118 | list.__init__(self, []) 119 | self.pretty = None 120 | 121 | def _repr_html_(self): 122 | _cell_with_spaces_pattern = re.compile(r"()( {2,})") 123 | if self.pretty: 124 | self.pretty.add_rows(self) 125 | result = self.pretty.get_html_string() 126 | result = _cell_with_spaces_pattern.sub(_nonbreaking_spaces, result) 127 | if self.config.displaylimit and len(self) > self.config.displaylimit: 128 | result = ( 129 | '%s\n%d rows, truncated to displaylimit of %d' 130 | % (result, len(self), self.config.displaylimit) 131 | ) 132 | return result 133 | else: 134 | return None 135 | 136 | def __str__(self, *arg, **kwarg): 137 | self.pretty.add_rows(self) 138 | return str(self.pretty or "") 139 | 140 | def __getitem__(self, key): 141 | """ 142 | Access by integer (row position within result set) 143 | or by string (value of leftmost column) 144 | """ 145 | try: 146 | return list.__getitem__(self, key) 147 | except TypeError: 148 | result = [row for row in self if row[0] == key] 149 | if not result: 150 | raise KeyError(key) 151 | if len(result) > 1: 152 | raise KeyError('%d results for "%s"' % (len(result), key)) 153 | return result[0] 154 | 155 | def dict(self): 156 | """Returns a single dict built from the result set 157 | 158 | Keys are column names; values are a tuple""" 159 | return dict(zip(self.keys, zip(*self))) 160 | 161 | def dicts(self): 162 | """Iterator yielding a dict for each row""" 163 | for row in self: 164 | yield dict(zip(self.keys, row)) 165 | 166 | def DataFrame(self): 167 | """Returns a Pandas DataFrame instance built from the result set.""" 168 | import pandas as pd 169 | 170 | frame = pd.DataFrame(self, columns=(self and self.keys) or []) 171 | return frame 172 | 173 | def pie(self, key_word_sep=" ", title=None, **kwargs): 174 | """Generates a pylab pie chart from the result set. 175 | 176 | ``matplotlib`` must be installed, and in an 177 | IPython Notebook, inlining must be on:: 178 | 179 | %%matplotlib inline 180 | 181 | Values (pie slice sizes) are taken from the 182 | rightmost column (numerical values required). 183 | All other columns are used to label the pie slices. 184 | 185 | Parameters 186 | ---------- 187 | key_word_sep: string used to separate column values 188 | from each other in pie labels 189 | title: Plot title, defaults to name of value column 190 | 191 | Any additional keyword arguments will be passsed 192 | through to ``matplotlib.pylab.pie``. 193 | """ 194 | self.guess_pie_columns(xlabel_sep=key_word_sep) 195 | import matplotlib.pylab as plt 196 | 197 | pie = plt.pie(self.ys[0], labels=self.xlabels, **kwargs) 198 | plt.title(title or self.ys[0].name) 199 | return pie 200 | 201 | def plot(self, title=None, **kwargs): 202 | """Generates a pylab plot from the result set. 203 | 204 | ``matplotlib`` must be installed, and in an 205 | IPython Notebook, inlining must be on:: 206 | 207 | %%matplotlib inline 208 | 209 | The first and last columns are taken as the X and Y 210 | values. Any columns between are ignored. 211 | 212 | Parameters 213 | ---------- 214 | title: Plot title, defaults to names of Y value columns 215 | 216 | Any additional keyword arguments will be passsed 217 | through to ``matplotlib.pylab.plot``. 218 | """ 219 | import matplotlib.pylab as plt 220 | 221 | self.guess_plot_columns() 222 | self.x = self.x or range(len(self.ys[0])) 223 | coords = reduce(operator.add, [(self.x, y) for y in self.ys]) 224 | plot = plt.plot(*coords, **kwargs) 225 | if hasattr(self.x, "name"): 226 | plt.xlabel(self.x.name) 227 | ylabel = ", ".join(y.name for y in self.ys) 228 | plt.title(title or ylabel) 229 | plt.ylabel(ylabel) 230 | return plot 231 | 232 | def bar(self, key_word_sep=" ", title=None, **kwargs): 233 | """Generates a pylab bar plot from the result set. 234 | 235 | ``matplotlib`` must be installed, and in an 236 | IPython Notebook, inlining must be on:: 237 | 238 | %%matplotlib inline 239 | 240 | The last quantitative column is taken as the Y values; 241 | all other columns are combined to label the X axis. 242 | 243 | Parameters 244 | ---------- 245 | title: Plot title, defaults to names of Y value columns 246 | key_word_sep: string used to separate column values 247 | from each other in labels 248 | 249 | Any additional keyword arguments will be passsed 250 | through to ``matplotlib.pylab.bar``. 251 | """ 252 | import matplotlib.pylab as plt 253 | 254 | self.guess_pie_columns(xlabel_sep=key_word_sep) 255 | plot = plt.bar(range(len(self.ys[0])), self.ys[0], **kwargs) 256 | if self.xlabels: 257 | plt.xticks(range(len(self.xlabels)), self.xlabels, rotation=45) 258 | plt.xlabel(self.xlabel) 259 | plt.ylabel(self.ys[0].name) 260 | return plot 261 | 262 | def csv(self, filename=None, **format_params): 263 | """Generate results in comma-separated form. Write to ``filename`` if given. 264 | Any other parameters will be passed on to csv.writer.""" 265 | if not self.pretty: 266 | return None # no results 267 | self.pretty.add_rows(self) 268 | if filename: 269 | encoding = format_params.get("encoding", "utf-8") 270 | if six.PY2: 271 | outfile = open(filename, "wb") 272 | else: 273 | outfile = open(filename, "w", newline="", encoding=encoding) 274 | else: 275 | outfile = six.StringIO() 276 | writer = UnicodeWriter(outfile, **format_params) 277 | writer.writerow(self.field_names) 278 | for row in self: 279 | writer.writerow(row) 280 | if filename: 281 | outfile.close() 282 | return CsvResultDescriptor(filename) 283 | else: 284 | return outfile.getvalue() 285 | 286 | 287 | def interpret_rowcount(rowcount): 288 | if rowcount < 0: 289 | result = "Done." 290 | else: 291 | result = "%d rows affected." % rowcount 292 | return result 293 | 294 | 295 | class FakeResultProxy(object): 296 | """A fake class that pretends to behave like the ResultProxy from 297 | SqlAlchemy. 298 | """ 299 | 300 | def __init__(self, cursor, headers): 301 | if cursor is None: 302 | cursor = [] 303 | headers = [] 304 | if isinstance(cursor, list): 305 | self.from_list(source_list=cursor) 306 | else: 307 | self.fetchall = cursor.fetchall 308 | self.fetchmany = cursor.fetchmany 309 | self.rowcount = cursor.rowcount 310 | self.keys = lambda: headers 311 | self.returns_rows = True 312 | 313 | def from_list(self, source_list): 314 | """Simulates SQLA ResultProxy from a list.""" 315 | 316 | self.fetchall = lambda: source_list 317 | self.rowcount = len(source_list) 318 | 319 | def fetchmany(size): 320 | pos = 0 321 | while pos < len(source_list): 322 | yield source_list[pos: pos + size] 323 | pos += size 324 | 325 | self.fetchmany = fetchmany 326 | 327 | 328 | # some dialects have autocommit 329 | # specific dialects break when commit is used: 330 | 331 | _COMMIT_BLACKLIST_DIALECTS = ("athena", "bigquery", "clickhouse", "ingres", "mssql", "teradata", "vertica") 332 | 333 | 334 | def _commit(conn, config): 335 | """Issues a commit, if appropriate for current config and dialect""" 336 | 337 | _should_commit = config.autocommit and all( 338 | dialect not in str(conn.dialect) for dialect in _COMMIT_BLACKLIST_DIALECTS 339 | ) 340 | 341 | if _should_commit: 342 | try: 343 | conn.internal_connection.commit() 344 | except sqlalchemy.exc.OperationalError: 345 | pass # not all engines can commit 346 | except Exception as ex: 347 | conn.internal_connection.rollback() 348 | print(traceback.format_exc()) 349 | raise ex 350 | 351 | 352 | def run(conn, sql, config, user_namespace): 353 | if sql.strip(): 354 | for statement in sqlparse.split(sql): 355 | first_word = sql.strip().split()[0].lower() 356 | if first_word == "begin": 357 | raise Exception("ipython_sql does not support transactions") 358 | if first_word.startswith("\\") and \ 359 | ("postgres" in str(conn.dialect) or 360 | "redshift" in str(conn.dialect)): 361 | if not PGSpecial: 362 | raise ImportError("pgspecial not installed") 363 | pgspecial = PGSpecial() 364 | _, cur, headers, _ = pgspecial.execute( 365 | conn.internal_connection.connection.cursor(), statement 366 | )[0] 367 | result = FakeResultProxy(cur, headers) 368 | else: 369 | txt = sqlalchemy.sql.text(statement) 370 | result = conn.internal_connection.execute(txt, user_namespace) 371 | _commit(conn=conn, config=config) 372 | if result and config.feedback: 373 | print(interpret_rowcount(result.rowcount)) 374 | resultset = ResultSet(result, config) 375 | if config.autopandas: 376 | return resultset.DataFrame() 377 | else: 378 | return resultset 379 | # returning only last result, intentionally 380 | else: 381 | return "Connected: %s" % conn.name 382 | 383 | 384 | class PrettyTable(prettytable.PrettyTable): 385 | def __init__(self, *args, **kwargs): 386 | self.row_count = 0 387 | self.displaylimit = None 388 | return super(PrettyTable, self).__init__(*args, **kwargs) 389 | 390 | def add_rows(self, data): 391 | if self.row_count and (data.config.displaylimit == self.displaylimit): 392 | return # correct number of rows already present 393 | self.clear_rows() 394 | self.displaylimit = data.config.displaylimit 395 | if self.displaylimit == 0: 396 | self.displaylimit = None # TODO: remove this to make 0 really 0 397 | if self.displaylimit in (None, 0): 398 | self.row_count = len(data) 399 | else: 400 | self.row_count = min(len(data), self.displaylimit) 401 | for row in data[: self.displaylimit]: 402 | self.add_row(row) 403 | -------------------------------------------------------------------------------- /src/tests/test_column_guesser.py: -------------------------------------------------------------------------------- 1 | import pytest 2 | 3 | from sql.magic import SqlMagic 4 | 5 | ip = get_ipython() 6 | 7 | 8 | class SqlEnv(object): 9 | def __init__(self, connectstr): 10 | self.connectstr = connectstr 11 | 12 | def query(self, txt): 13 | return ip.run_line_magic("sql", "%s %s" % (self.connectstr, txt)) 14 | 15 | 16 | sql_env = SqlEnv("sqlite://") 17 | 18 | 19 | @pytest.fixture 20 | def tbl(): 21 | sqlmagic = SqlMagic(shell=ip) 22 | ip.register_magics(sqlmagic) 23 | creator = """ 24 | DROP TABLE IF EXISTS manycoltbl; 25 | CREATE TABLE manycoltbl 26 | (name TEXT, y1 REAL, y2 REAL, name2 TEXT, y3 INT); 27 | INSERT INTO manycoltbl VALUES 28 | ('r1-txt1', 1.01, 1.02, 'r1-txt2', 1.04); 29 | INSERT INTO manycoltbl VALUES 30 | ('r2-txt1', 2.01, 2.02, 'r2-txt2', 2.04); 31 | INSERT INTO manycoltbl VALUES ('r3-txt1', 3.01, 3.02, 'r3-txt2', 3.04); 32 | """ 33 | for qry in creator.split(";"): 34 | sql_env.query(qry) 35 | yield 36 | sql_env.query("DROP TABLE manycoltbl") 37 | 38 | 39 | class Harness(object): 40 | def run_query(self): 41 | return sql_env.query(self.query) 42 | 43 | 44 | class TestOneNum(Harness): 45 | query = "SELECT y1 FROM manycoltbl" 46 | 47 | def test_pie(self, tbl): 48 | results = self.run_query() 49 | results.guess_pie_columns(xlabel_sep="//") 50 | assert results.ys[0].is_quantity 51 | assert results.ys == [[1.01, 2.01, 3.01]] 52 | assert results.x == [] 53 | assert results.xlabels == [] 54 | assert results.xlabel == "" 55 | 56 | def test_plot(self, tbl): 57 | results = self.run_query() 58 | results.guess_plot_columns() 59 | assert results.ys == [[1.01, 2.01, 3.01]] 60 | assert results.x == [] 61 | assert results.x.name == "" 62 | 63 | 64 | class TestOneStrOneNum(Harness): 65 | query = "SELECT name, y1 FROM manycoltbl" 66 | 67 | def test_pie(self, tbl): 68 | results = self.run_query() 69 | results.guess_pie_columns(xlabel_sep="//") 70 | assert results.ys[0].is_quantity 71 | assert results.ys == [[1.01, 2.01, 3.01]] 72 | assert results.xlabels == ["r1-txt1", "r2-txt1", "r3-txt1"] 73 | assert results.xlabel == "name" 74 | 75 | def test_plot(self, tbl): 76 | results = self.run_query() 77 | results.guess_plot_columns() 78 | assert results.ys == [[1.01, 2.01, 3.01]] 79 | assert results.x == [] 80 | 81 | 82 | class TestTwoStrTwoNum(Harness): 83 | query = "SELECT name2, y3, name, y1 FROM manycoltbl" 84 | 85 | def test_pie(self, tbl): 86 | results = self.run_query() 87 | results.guess_pie_columns(xlabel_sep="//") 88 | assert results.ys[0].is_quantity 89 | assert results.ys == [[1.01, 2.01, 3.01]] 90 | assert results.xlabels == [ 91 | "r1-txt2//1.04//r1-txt1", 92 | "r2-txt2//2.04//r2-txt1", 93 | "r3-txt2//3.04//r3-txt1", 94 | ] 95 | assert results.xlabel == "name2, y3, name" 96 | 97 | def test_plot(self, tbl): 98 | results = self.run_query() 99 | results.guess_plot_columns() 100 | assert results.ys == [[1.01, 2.01, 3.01]] 101 | assert results.x == [1.04, 2.04, 3.04] 102 | 103 | 104 | class TestTwoStrThreeNum(Harness): 105 | query = "SELECT name, y1, name2, y2, y3 FROM manycoltbl" 106 | 107 | def test_pie(self, tbl): 108 | results = self.run_query() 109 | results.guess_pie_columns(xlabel_sep="//") 110 | assert results.ys[0].is_quantity 111 | assert results.ys == [[1.04, 2.04, 3.04]] 112 | assert results.xlabels == [ 113 | "r1-txt1//1.01//r1-txt2//1.02", 114 | "r2-txt1//2.01//r2-txt2//2.02", 115 | "r3-txt1//3.01//r3-txt2//3.02", 116 | ] 117 | 118 | def test_plot(self, tbl): 119 | results = self.run_query() 120 | results.guess_plot_columns() 121 | assert results.ys == [[1.02, 2.02, 3.02], [1.04, 2.04, 3.04]] 122 | assert results.x == [1.01, 2.01, 3.01] 123 | -------------------------------------------------------------------------------- /src/tests/test_dsn_config.ini: -------------------------------------------------------------------------------- 1 | [DB_CONFIG_1] 2 | drivername = postgres 3 | host = my.remote.host 4 | port = 5432 5 | database = pgmain 6 | username = goesto11 7 | password = seentheelephant 8 | 9 | [DB_CONFIG_2] 10 | drivername = mysql 11 | host = 127.0.0.1 12 | database = dolfin 13 | username = thefin 14 | password = fishputsfishonthetable 15 | -------------------------------------------------------------------------------- /src/tests/test_magic.py: -------------------------------------------------------------------------------- 1 | import os.path 2 | import re 3 | import tempfile 4 | from textwrap import dedent 5 | 6 | import pytest 7 | 8 | 9 | def runsql(ip_session, statements): 10 | if isinstance(statements, str): 11 | statements = [statements] 12 | for statement in statements: 13 | result = ip_session.run_line_magic("sql", "sqlite:// %s" % statement) 14 | return result # returns only last result 15 | 16 | 17 | @pytest.fixture 18 | def ip(): 19 | """Provides an IPython session in which tables have been created""" 20 | 21 | ip_session = get_ipython() 22 | runsql( 23 | ip_session, 24 | [ 25 | "CREATE TABLE test (n INT, name TEXT)", 26 | "INSERT INTO test VALUES (1, 'foo')", 27 | "INSERT INTO test VALUES (2, 'bar')", 28 | "CREATE TABLE author (first_name, last_name, year_of_death)", 29 | "INSERT INTO author VALUES ('William', 'Shakespeare', 1616)", 30 | "INSERT INTO author VALUES ('Bertold', 'Brecht', 1956)", 31 | ], 32 | ) 33 | yield ip_session 34 | runsql(ip_session, "DROP TABLE test") 35 | runsql(ip_session, "DROP TABLE author") 36 | 37 | 38 | def test_memory_db(ip): 39 | assert runsql(ip, "SELECT * FROM test;")[0][0] == 1 40 | assert runsql(ip, "SELECT * FROM test;")[1].name == "bar" 41 | 42 | 43 | def test_html(ip): 44 | result = runsql(ip, "SELECT * FROM test;") 45 | assert "foo" in result._repr_html_().lower() 46 | 47 | 48 | def test_print(ip): 49 | result = runsql(ip, "SELECT * FROM test;") 50 | assert re.search(r"1\s+\|\s+foo", str(result)) 51 | 52 | 53 | def test_plain_style(ip): 54 | ip.run_line_magic("config", "SqlMagic.style = 'PLAIN_COLUMNS'") 55 | result = runsql(ip, "SELECT * FROM test;") 56 | assert re.search(r"1\s+\|\s+foo", str(result)) 57 | 58 | 59 | @pytest.mark.skip 60 | def test_multi_sql(ip): 61 | result = ip.run_cell_magic( 62 | "sql", 63 | "", 64 | """ 65 | sqlite:// 66 | SELECT last_name FROM author; 67 | """, 68 | ) 69 | assert "Shakespeare" in str(result) and "Brecht" in str(result) 70 | 71 | 72 | def test_result_var(ip): 73 | ip.run_cell_magic( 74 | "sql", 75 | "", 76 | """ 77 | sqlite:// 78 | x << 79 | SELECT last_name FROM author; 80 | """, 81 | ) 82 | result = ip.user_global_ns["x"] 83 | assert "Shakespeare" in str(result) and "Brecht" in str(result) 84 | 85 | 86 | def test_result_var_multiline_shovel(ip): 87 | ip.run_cell_magic( 88 | "sql", 89 | "", 90 | """ 91 | sqlite:// x << SELECT last_name 92 | FROM author; 93 | """, 94 | ) 95 | result = ip.user_global_ns["x"] 96 | assert "Shakespeare" in str(result) and "Brecht" in str(result) 97 | 98 | 99 | def test_access_results_by_keys(ip): 100 | assert runsql(ip, "SELECT * FROM author;")["William"] == ( 101 | u"William", 102 | u"Shakespeare", 103 | 1616, 104 | ) 105 | 106 | 107 | def test_duplicate_column_names_accepted(ip): 108 | result = ip.run_cell_magic( 109 | "sql", 110 | "", 111 | """ 112 | sqlite:// 113 | SELECT last_name, last_name FROM author; 114 | """, 115 | ) 116 | assert (u"Brecht", u"Brecht") in result 117 | 118 | 119 | def test_autolimit(ip): 120 | ip.run_line_magic("config", "SqlMagic.autolimit = 0") 121 | result = runsql(ip, "SELECT * FROM test;") 122 | assert len(result) == 2 123 | ip.run_line_magic("config", "SqlMagic.autolimit = 1") 124 | result = runsql(ip, "SELECT * FROM test;") 125 | assert len(result) == 1 126 | 127 | 128 | def test_persist(ip): 129 | runsql(ip, "") 130 | ip.run_cell("results = %sql SELECT * FROM test;") 131 | ip.run_cell("results_dframe = results.DataFrame()") 132 | ip.run_cell("%sql --persist sqlite:// results_dframe") 133 | persisted = runsql(ip, "SELECT * FROM results_dframe") 134 | assert "foo" in str(persisted) 135 | 136 | 137 | def test_append(ip): 138 | runsql(ip, "") 139 | ip.run_cell("results = %sql SELECT * FROM test;") 140 | ip.run_cell("results_dframe = results.DataFrame()") 141 | ip.run_cell("%sql --persist sqlite:// results_dframe") 142 | persisted = runsql(ip, "SELECT COUNT(*) FROM results_dframe") 143 | ip.run_cell("%sql --append sqlite:// results_dframe") 144 | appended = runsql(ip, "SELECT COUNT(*) FROM results_dframe") 145 | assert appended[0][0] == persisted[0][0] * 2 146 | 147 | 148 | def test_persist_nonexistent_raises(ip): 149 | runsql(ip, "") 150 | result = ip.run_cell("%sql --persist sqlite:// no_such_dataframe") 151 | assert result.error_in_exec 152 | 153 | 154 | def test_persist_non_frame_raises(ip): 155 | ip.run_cell("not_a_dataframe = 22") 156 | runsql(ip, "") 157 | result = ip.run_cell("%sql --persist sqlite:// not_a_dataframe") 158 | assert result.error_in_exec 159 | 160 | 161 | def test_persist_bare(ip): 162 | result = ip.run_cell("%sql --persist sqlite://") 163 | assert result.error_in_exec 164 | 165 | 166 | def test_persist_frame_at_its_creation(ip): 167 | ip.run_cell("results = %sql SELECT * FROM author;") 168 | ip.run_cell("%sql --persist sqlite:// results.DataFrame()") 169 | persisted = runsql(ip, "SELECT * FROM results") 170 | assert "Shakespeare" in str(persisted) 171 | 172 | 173 | def test_connection_args_enforce_json(ip): 174 | result = ip.run_cell('%sql --connection_arguments {"badlyformed":true') 175 | assert result.error_in_exec 176 | 177 | 178 | def test_connection_args_in_connection(ip): 179 | ip.run_cell('%sql --connection_arguments {"timeout":10} sqlite:///:memory:') 180 | result = ip.run_cell("%sql --connections") 181 | assert "timeout" in result.result["sqlite:///:memory:"].connect_args 182 | 183 | 184 | def test_connection_args_single_quotes(ip): 185 | ip.run_cell("%sql --connection_arguments '{\"timeout\": 10}' sqlite:///:memory:") 186 | result = ip.run_cell("%sql --connections") 187 | assert "timeout" in result.result["sqlite:///:memory:"].connect_args 188 | 189 | 190 | def test_connection_args_double_quotes(ip): 191 | ip.run_cell('%sql --connection_arguments "{\\"timeout\\": 10}" sqlite:///:memory:') 192 | result = ip.run_cell("%sql --connections") 193 | assert "timeout" in result.result["sqlite:///:memory:"].connect_args 194 | 195 | 196 | # TODO: support 197 | # @with_setup(_setup_author, _teardown_author) 198 | # def test_persist_with_connection_info(): 199 | # ip.run_cell("results = %sql SELECT * FROM author;") 200 | # ip.run_line_magic('sql', 'sqlite:// PERSIST results.DataFrame()') 201 | # persisted = ip.run_line_magic('sql', 'SELECT * FROM results') 202 | # assert 'Shakespeare' in str(persisted) 203 | 204 | 205 | def test_displaylimit(ip): 206 | ip.run_line_magic("config", "SqlMagic.autolimit = None") 207 | ip.run_line_magic("config", "SqlMagic.displaylimit = None") 208 | result = runsql( 209 | ip, 210 | "SELECT * FROM (VALUES ('apple'), ('banana'), ('cherry')) AS Result ORDER BY 1;", 211 | ) 212 | assert "apple" in result._repr_html_() 213 | assert "banana" in result._repr_html_() 214 | assert "cherry" in result._repr_html_() 215 | ip.run_line_magic("config", "SqlMagic.displaylimit = 1") 216 | result = runsql( 217 | ip, 218 | "SELECT * FROM (VALUES ('apple'), ('banana'), ('cherry')) AS Result ORDER BY 1;", 219 | ) 220 | assert "apple" in result._repr_html_() 221 | assert "cherry" not in result._repr_html_() 222 | 223 | 224 | def test_column_local_vars(ip): 225 | ip.run_line_magic("config", "SqlMagic.column_local_vars = True") 226 | result = runsql(ip, "SELECT * FROM author;") 227 | assert result is None 228 | assert "William" in ip.user_global_ns["first_name"] 229 | assert "Shakespeare" in ip.user_global_ns["last_name"] 230 | assert len(ip.user_global_ns["first_name"]) == 2 231 | ip.run_line_magic("config", "SqlMagic.column_local_vars = False") 232 | 233 | 234 | def test_userns_not_changed(ip): 235 | ip.run_cell( 236 | dedent( 237 | """ 238 | def function(): 239 | local_var = 'local_val' 240 | %sql sqlite:// INSERT INTO test VALUES (2, 'bar'); 241 | function()""" 242 | ) 243 | ) 244 | assert "local_var" not in ip.user_ns 245 | 246 | 247 | def test_bind_vars(ip): 248 | ip.user_global_ns["x"] = 22 249 | result = runsql(ip, "SELECT :x") 250 | assert result[0][0] == 22 251 | 252 | 253 | def test_autopandas(ip): 254 | ip.run_line_magic("config", "SqlMagic.autopandas = True") 255 | dframe = runsql(ip, "SELECT * FROM test;") 256 | assert not dframe.empty 257 | assert dframe.ndim == 2 258 | assert dframe.name[0] == "foo" 259 | 260 | 261 | def test_csv(ip): 262 | ip.run_line_magic("config", "SqlMagic.autopandas = False") # uh-oh 263 | result = runsql(ip, "SELECT * FROM test;") 264 | result = result.csv() 265 | for row in result.splitlines(): 266 | assert row.count(",") == 1 267 | assert len(result.splitlines()) == 3 268 | 269 | 270 | def test_csv_to_file(ip): 271 | ip.run_line_magic("config", "SqlMagic.autopandas = False") # uh-oh 272 | result = runsql(ip, "SELECT * FROM test;") 273 | with tempfile.TemporaryDirectory() as tempdir: 274 | fname = os.path.join(tempdir, "test.csv") 275 | output = result.csv(fname) 276 | assert os.path.exists(output.file_path) 277 | with open(output.file_path) as csvfile: 278 | content = csvfile.read() 279 | for row in content.splitlines(): 280 | assert row.count(",") == 1 281 | assert len(content.splitlines()) == 3 282 | 283 | 284 | def test_sql_from_file(ip): 285 | ip.run_line_magic("config", "SqlMagic.autopandas = False") 286 | with tempfile.TemporaryDirectory() as tempdir: 287 | fname = os.path.join(tempdir, "test.sql") 288 | with open(fname, "w") as tempf: 289 | tempf.write("SELECT * FROM test;") 290 | result = ip.run_cell("%sql --file " + fname) 291 | assert result.result == [(1, "foo"), (2, "bar")] 292 | 293 | 294 | def test_sql_from_nonexistent_file(ip): 295 | ip.run_line_magic("config", "SqlMagic.autopandas = False") 296 | with tempfile.TemporaryDirectory() as tempdir: 297 | fname = os.path.join(tempdir, "nonexistent.sql") 298 | result = ip.run_cell("%sql --file " + fname) 299 | assert isinstance(result.error_in_exec, FileNotFoundError) 300 | 301 | 302 | def test_dict(ip): 303 | result = runsql(ip, "SELECT * FROM author;") 304 | result = result.dict() 305 | assert isinstance(result, dict) 306 | assert "first_name" in result 307 | assert "last_name" in result 308 | assert "year_of_death" in result 309 | assert len(result["last_name"]) == 2 310 | 311 | 312 | def test_dicts(ip): 313 | result = runsql(ip, "SELECT * FROM author;") 314 | for row in result.dicts(): 315 | assert isinstance(row, dict) 316 | assert "first_name" in row 317 | assert "last_name" in row 318 | assert "year_of_death" in row 319 | 320 | 321 | def test_bracket_var_substitution(ip): 322 | ip.user_global_ns["col"] = "first_name" 323 | assert runsql(ip, "SELECT * FROM author" " WHERE {col} = 'William' ")[0] == ( 324 | u"William", 325 | u"Shakespeare", 326 | 1616, 327 | ) 328 | 329 | ip.user_global_ns["col"] = "last_name" 330 | result = runsql(ip, "SELECT * FROM author" " WHERE {col} = 'William' ") 331 | assert not result 332 | 333 | 334 | def test_multiline_bracket_var_substitution(ip): 335 | ip.user_global_ns["col"] = "first_name" 336 | assert runsql(ip, "SELECT * FROM author\n" " WHERE {col} = 'William' ")[0] == ( 337 | u"William", 338 | u"Shakespeare", 339 | 1616, 340 | ) 341 | 342 | ip.user_global_ns["col"] = "last_name" 343 | result = runsql(ip, "SELECT * FROM author" " WHERE {col} = 'William' ") 344 | assert not result 345 | 346 | 347 | def test_multiline_bracket_var_substitution(ip): 348 | ip.user_global_ns["col"] = "first_name" 349 | result = ip.run_cell_magic( 350 | "sql", 351 | "", 352 | """ 353 | sqlite:// SELECT * FROM author 354 | WHERE {col} = 'William' 355 | """, 356 | ) 357 | assert (u"William", u"Shakespeare", 1616) in result 358 | 359 | ip.user_global_ns["col"] = "last_name" 360 | result = ip.run_cell_magic( 361 | "sql", 362 | "", 363 | """ 364 | sqlite:// SELECT * FROM author 365 | WHERE {col} = 'William' 366 | """, 367 | ) 368 | assert not result 369 | 370 | 371 | def test_json_in_select(ip): 372 | # Variable expansion does not work within json, but 373 | # at least the two usages of curly braces do not collide 374 | ip.user_global_ns["person"] = "prince" 375 | result = ip.run_cell_magic( 376 | "sql", 377 | "", 378 | """ 379 | sqlite:// 380 | SELECT 381 | '{"greeting": "Farewell sweet {person}"}' 382 | AS json 383 | """, 384 | ) 385 | assert ('{"greeting": "Farewell sweet {person}"}',) 386 | 387 | 388 | def test_close_connection(ip): 389 | connections = runsql(ip, "%sql -l") 390 | connection_name = list(connections)[0] 391 | runsql(ip, f"%sql -x {connection_name}") 392 | connections_afterward = runsql(ip, "%sql -l") 393 | assert connection_name not in connections_afterward 394 | -------------------------------------------------------------------------------- /src/tests/test_parse.py: -------------------------------------------------------------------------------- 1 | import os 2 | from pathlib import Path 3 | 4 | from sql.parse import connection_from_dsn_section, parse, without_sql_comment 5 | 6 | try: 7 | from traitlets.config.configurable import Configurable 8 | except ImportError: 9 | from IPython.config.configurable import Configurable 10 | 11 | empty_config = Configurable() 12 | default_connect_args = {"options": "-csearch_path=test"} 13 | 14 | 15 | def test_parse_no_sql(): 16 | assert parse("will:longliveliz@localhost/shakes", empty_config) == { 17 | "connection": "will:longliveliz@localhost/shakes", 18 | "sql": "", 19 | "result_var": None, 20 | } 21 | 22 | 23 | def test_parse_with_sql(): 24 | assert parse( 25 | "postgresql://will:longliveliz@localhost/shakes SELECT * FROM work", 26 | empty_config, 27 | ) == { 28 | "connection": "postgresql://will:longliveliz@localhost/shakes", 29 | "sql": "SELECT * FROM work", 30 | "result_var": None, 31 | } 32 | 33 | 34 | def test_parse_sql_only(): 35 | assert parse("SELECT * FROM work", empty_config) == { 36 | "connection": "", 37 | "sql": "SELECT * FROM work", 38 | "result_var": None, 39 | } 40 | 41 | 42 | def test_parse_postgresql_socket_connection(): 43 | assert parse("postgresql:///shakes SELECT * FROM work", empty_config) == { 44 | "connection": "postgresql:///shakes", 45 | "sql": "SELECT * FROM work", 46 | "result_var": None, 47 | } 48 | 49 | 50 | def test_expand_environment_variables_in_connection(): 51 | os.environ["DATABASE_URL"] = "postgresql:///shakes" 52 | assert parse("$DATABASE_URL SELECT * FROM work", empty_config) == { 53 | "connection": "postgresql:///shakes", 54 | "sql": "SELECT * FROM work", 55 | "result_var": None, 56 | } 57 | 58 | 59 | def test_parse_shovel_operator(): 60 | assert parse("dest << SELECT * FROM work", empty_config) == { 61 | "connection": "", 62 | "sql": "SELECT * FROM work", 63 | "result_var": "dest", 64 | } 65 | 66 | 67 | def test_parse_connect_plus_shovel(): 68 | assert parse("sqlite:// dest << SELECT * FROM work", empty_config) == { 69 | "connection": "sqlite://", 70 | "sql": "SELECT * FROM work", 71 | "result_var": None, 72 | } 73 | 74 | 75 | def test_parse_shovel_operator(): 76 | assert parse("dest << SELECT * FROM work", empty_config) == { 77 | "connection": "", 78 | "sql": "SELECT * FROM work", 79 | "result_var": "dest", 80 | } 81 | 82 | 83 | def test_parse_connect_plus_shovel(): 84 | assert parse("sqlite:// dest << SELECT * FROM work", empty_config) == { 85 | "connection": "sqlite://", 86 | "sql": "SELECT * FROM work", 87 | "result_var": "dest", 88 | } 89 | 90 | 91 | def test_parse_early_newlines(): 92 | assert parse("--comment\nSELECT *\n--comment\nFROM work", empty_config) == { 93 | "connection": "", 94 | "sql": "--comment\nSELECT *\n--comment\nFROM work", 95 | "result_var": None 96 | } 97 | 98 | 99 | def test_parse_connect_shovel_over_newlines(): 100 | assert parse("\nsqlite://\ndest\n<<\nSELECT *\nFROM work", empty_config) == { 101 | "connection": "sqlite://", 102 | "sql": "SELECT *\nFROM work", 103 | "result_var": "dest" 104 | } 105 | 106 | 107 | class DummyConfig: 108 | dsn_filename = Path("src/tests/test_dsn_config.ini") 109 | 110 | 111 | def test_connection_from_dsn_section(): 112 | result = connection_from_dsn_section(section="DB_CONFIG_1", config=DummyConfig()) 113 | assert str(result) == "postgres://goesto11:***@my.remote.host:5432/pgmain" 114 | result = connection_from_dsn_section(section="DB_CONFIG_2", config=DummyConfig()) 115 | assert str(result) == "mysql://thefin:***@127.0.0.1/dolfin" 116 | 117 | 118 | class Bunch: 119 | def __init__(self, **kwds): 120 | self.__dict__.update(kwds) 121 | 122 | 123 | class ParserStub: 124 | opstrs = [ 125 | [], 126 | ["-l", "--connections"], 127 | ["-x", "--close"], 128 | ["-c", "--creator"], 129 | ["-s", "--section"], 130 | ["-p", "--persist"], 131 | ["--append"], 132 | ["-a", "--connection_arguments"], 133 | ["-f", "--file"], 134 | ] 135 | _actions = [Bunch(option_strings=o) for o in opstrs] 136 | 137 | 138 | parser_stub = ParserStub() 139 | 140 | 141 | def test_without_sql_comment_plain(): 142 | line = "SELECT * FROM author" 143 | assert without_sql_comment(parser=parser_stub, line=line) == line 144 | 145 | 146 | def test_without_sql_comment_with_arg(): 147 | line = "--file moo.txt --persist SELECT * FROM author" 148 | assert without_sql_comment(parser=parser_stub, line=line) == line 149 | 150 | 151 | def test_without_sql_comment_with_comment(): 152 | line = "SELECT * FROM author -- uff da" 153 | expected = "SELECT * FROM author" 154 | assert without_sql_comment(parser=parser_stub, line=line) == expected 155 | 156 | 157 | def test_without_sql_comment_with_arg_and_comment(): 158 | line = "--file moo.txt --persist SELECT * FROM author -- uff da" 159 | expected = "--file moo.txt --persist SELECT * FROM author" 160 | assert without_sql_comment(parser=parser_stub, line=line) == expected 161 | 162 | 163 | def test_without_sql_comment_unspaced_comment(): 164 | line = "SELECT * FROM author --uff da" 165 | expected = "SELECT * FROM author" 166 | assert without_sql_comment(parser=parser_stub, line=line) == expected 167 | 168 | 169 | def test_without_sql_comment_dashes_in_string(): 170 | line = "SELECT '--very --confusing' FROM author -- uff da" 171 | expected = "SELECT '--very --confusing' FROM author" 172 | assert without_sql_comment(parser=parser_stub, line=line) == expected 173 | 174 | 175 | def test_without_sql_comment_with_arg_and_leading_comment(): 176 | line = "--file moo.txt --persist --comment, not arg" 177 | expected = "--file moo.txt --persist" 178 | assert without_sql_comment(parser=parser_stub, line=line) == expected 179 | 180 | 181 | def test_without_sql_persist(): 182 | line = "--persist my_table --uff da" 183 | expected = "--persist my_table" 184 | assert without_sql_comment(parser=parser_stub, line=line) == expected 185 | -------------------------------------------------------------------------------- /tox.ini: -------------------------------------------------------------------------------- 1 | [tox] 2 | envlist = py27,py36 3 | 4 | [testenv] 5 | deps = pytest 6 | -rrequirements.txt 7 | -rrequirements-dev.txt 8 | commands = 9 | ipython -c "import pytest; pytest.main(['.'])" 10 | --------------------------------------------------------------------------------