├── .gitignore ├── .travis.yml ├── LICENSE.md ├── README.md ├── docs └── images │ └── interactive_demo.gif ├── push_to_pypi.sh ├── query ├── __init__.py ├── core.py ├── helpers.py ├── html.py └── sample_data │ ├── CHINOOK_INFO.md │ └── Chinook_Sqlite.sqlite ├── requirements.txt ├── setup.py ├── test_requirements.txt └── tests ├── __init__.py ├── core_test.py └── helpers_test.py /.gitignore: -------------------------------------------------------------------------------- 1 | .DS_Store 2 | *.pyc 3 | venv/* 4 | *.ipynb 5 | -------------------------------------------------------------------------------- /.travis.yml: -------------------------------------------------------------------------------- 1 | language: python 2 | python: 3 | - "2.7" 4 | 5 | # command to install dependencies 6 | install: 7 | - "pip install -r requirements.txt" 8 | - "pip install -r test_requirements.txt" 9 | - "pip install ." 10 | 11 | # command to run tests 12 | script: 13 | - nosetests --with-cover --cover-package=query 14 | -------------------------------------------------------------------------------- /LICENSE.md: -------------------------------------------------------------------------------- 1 | The MIT License 2 | 3 | Copyright (c) 2014 Nick Greenfield, www.boydgreenfield.com 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 13 | all 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 21 | THE SOFTWARE. -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # query 2 | Python module for quick, interactive exploration of SQL databases. Designed especially for use with IPython. Light wrapper on top of Pandas (>= 0.16) and SQLAlchemy (>= 0.9.9). The most recent release is available on [PyPI](https://pypi.python.org/pypi/query) and can be installed via `pip install query` with the proper dependencies. 3 | 4 | [![Build Status](https://travis-ci.org/boydgreenfield/query.svg?branch=master)](https://travis-ci.org/boydgreenfield/query) 5 | 6 | 7 | ## Quickstart 8 | ```python 9 | from query import QueryDb 10 | db = QueryDb(demo=True) 11 | ``` 12 | 13 | But the real joy comes when using query interactively: 14 | 15 | ![Interactive query use demo #1 ](docs/images/interactive_demo.gif?raw=True) 16 | 17 | ## Key functionality 18 | A few key functions to remember: 19 | 20 | * `db`: The main database object. Print it in IPython to see a list of tables and their key attributes. 21 | * `db.inspect.*`: Tab-completion across the database's tables and columns. Print any table to see its columns and their types. 22 | * `db.query()`: Query the database with a raw SQL query. Returns a `pandas DataFrame` object by default, but can return a `sqlalchemy result` object if called with `return_as="result"`. 23 | 24 | 25 | ## Roadmap 26 | Further improvements are planned, including some of the below. Please feel free to open an Issue with desired features or submit a pull request. 27 | 28 | * **Plotting**: Graphical display of queried data (some of this can easily be done off the current `DataFrame` object, but it could be better integrated) 29 | * **More Convenience Methods**: Additional convenience methods, like ``.tail()`` and ``.where()`` 30 | * **DB Schemas**: Direct output of database schema diagrams 31 | * **Password Input via IPython**: Currently entering a DB password uses `getpass` in the user's terminal. Being able to enter the password directly into IPython would be ideal (while also not writing it into any history) 32 | * **More/Better Documentation**: Enough said. 33 | -------------------------------------------------------------------------------- /docs/images/interactive_demo.gif: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/boydgreenfield/query/03aa43b746b43832af3f0403265e648a5617b62b/docs/images/interactive_demo.gif -------------------------------------------------------------------------------- /push_to_pypi.sh: -------------------------------------------------------------------------------- 1 | set -e 2 | 3 | source venv/bin/activate 4 | nosetests 5 | echo "Tests successful. Pushing to PyPI..." 6 | python setup.py sdist upload 7 | -------------------------------------------------------------------------------- /query/__init__.py: -------------------------------------------------------------------------------- 1 | """ 2 | Module for conveniently exploring and fetching data from a variety of SQL databases. 3 | Is *not* primarily designed for writing to SQL databases, and should be used 4 | principally for read-only data exploration. Use with IPython Notebook is heartily 5 | recommended. 6 | """ 7 | from query.core import QueryDb # noqa 8 | 9 | # Aliases 10 | from query.core import QueryDb as QueryDB # noqa 11 | -------------------------------------------------------------------------------- /query/core.py: -------------------------------------------------------------------------------- 1 | import getpass 2 | import sqlalchemy 3 | import numpy as np 4 | import pandas as pd 5 | import os 6 | import warnings 7 | 8 | import query 9 | from query.html import df_to_html, GETPASS_USE_WARNING, QUERY_DB_ATTR_MSG 10 | 11 | 12 | # Exceptions 13 | class QueryDbError(Exception): 14 | pass 15 | 16 | 17 | class NoPrimaryKeyException(Exception): 18 | pass 19 | 20 | 21 | # Helper classes 22 | class QueryDbAttributes(object): 23 | def _repr_html_(self): 24 | return QUERY_DB_ATTR_MSG 25 | 26 | 27 | class QueryDbMeta(sqlalchemy.schema.MetaData): 28 | pass 29 | 30 | 31 | # Main classes 32 | class QueryDbOrm(object): 33 | """ 34 | A helper class for QueryDb -- allows making Sqlalchemy schema 35 | (Tables and Columns) more accessible via the QueryDb interface 36 | and exposing additional helper methods, e.g., .last(). 37 | """ 38 | def __init__(self, orm_object, db): 39 | self._db = db 40 | if orm_object.__class__ == sqlalchemy.schema.Table: 41 | self._is_table = True 42 | self.table = orm_object 43 | self.column = None 44 | 45 | # Support custom styling of Pandas dataframe by calling .to_html() over _repr_html() 46 | # incl. not displaying dimensions, for example 47 | self._column_df = pd.DataFrame( 48 | [(c.name, c.type, c.primary_key) for c in self.table.columns.values()], 49 | columns=["Column", "Type", "Primary Key"] 50 | ) 51 | self._html = df_to_html(self._column_df, ("Column Information for the %s Table" 52 | % self.table.name)) 53 | 54 | else: 55 | self._is_table = False 56 | self.table = orm_object.table 57 | self.column = orm_object 58 | self._html = ("Inspecting column %s of the %s table. " 59 | "Try the .head(), .tail(), and .where() " 60 | "methods to further explore." % 61 | (self.column.name, self.table.name)) 62 | 63 | def _repr_html_(self): 64 | return self._html 65 | 66 | def __repr__(self): 67 | if self._is_table: 68 | return self.table.__repr__() 69 | else: 70 | return self.column.__repr__() 71 | 72 | def _query_helper(self, by=None): 73 | """ 74 | Internal helper for preparing queries. 75 | """ 76 | if by is None: 77 | primary_keys = self.table.primary_key.columns.keys() 78 | 79 | if len(primary_keys) > 1: 80 | warnings.warn("WARNING: MORE THAN 1 PRIMARY KEY FOR TABLE %s. " 81 | "USING THE FIRST KEY %s." % 82 | (self.table.name, primary_keys[0])) 83 | 84 | if not primary_keys: 85 | raise NoPrimaryKeyException("Table %s needs a primary key for" 86 | "the .last() method to work properly. " 87 | "Alternatively, specify an ORDER BY " 88 | "column with the by= argument. " % 89 | self.table.name) 90 | id_col = primary_keys[0] 91 | else: 92 | id_col = by 93 | 94 | if self.column is None: 95 | col = "*" 96 | else: 97 | col = self.column.name 98 | 99 | return col, id_col 100 | 101 | def head(self, n=10, by=None, **kwargs): 102 | """ 103 | Get the first n entries for a given Table/Column. Additional keywords 104 | passed to QueryDb.query(). 105 | 106 | Requires that the given table has a primary key specified. 107 | """ 108 | col, id_col = self._query_helper(by=by) 109 | 110 | select = ("SELECT %s FROM %s ORDER BY %s ASC LIMIT %d" % 111 | (col, self.table.name, id_col, n)) 112 | 113 | return self._db.query(select, **kwargs) 114 | 115 | def tail(self, n=10, by=None, **kwargs): 116 | """ 117 | Get the last n entries for a given Table/Column. Additional keywords 118 | passed to QueryDb.query(). 119 | 120 | Requires that the given table has a primary key specified. 121 | """ 122 | col, id_col = self._query_helper(by=by) 123 | 124 | select = ("SELECT %s FROM %s ORDER BY %s DESC LIMIT %d" % 125 | (col, self.table.name, id_col, n)) 126 | 127 | return self._db.query(select, **kwargs) 128 | 129 | def first(self, n=10, by=None, **kwargs): 130 | """ 131 | Alias for .head(). 132 | """ 133 | 134 | def last(self, n=10, by=None, **kwargs): 135 | """ 136 | Alias for .tail(). 137 | """ 138 | return self.tail(n=n, by=by, **kwargs) 139 | 140 | def where(self, where_string, **kwargs): 141 | """ 142 | Select from a given Table or Column with the specified WHERE clause 143 | string. Additional keywords are passed to ExploreSqlDB.query(). For 144 | convenience, if there is no '=', '>', '<', 'like', or 'LIKE' clause 145 | in the WHERE statement .where() tries to match the input string 146 | against the primary key column of the Table. 147 | 148 | Args: 149 | where_string (str): Where clause for the query against the Table 150 | or Column 151 | 152 | Kwars: 153 | **kwargs: Optional **kwargs passed to the QueryDb.query() call 154 | 155 | Returns: 156 | result (pandas.DataFrame or sqlalchemy ResultProxy): Query result 157 | as a DataFrame (default) or sqlalchemy result. 158 | """ 159 | col, id_col = self._query_helper(by=None) 160 | 161 | where_string = str(where_string) # Coerce here, for .__contains___ 162 | where_operators = ["=", ">", "<", "LIKE", "like"] 163 | if np.any([where_string.__contains__(w) for w in where_operators]): 164 | select = ("SELECT %s FROM %s WHERE %s" % 165 | (col, self.table.name, where_string)) 166 | else: 167 | select = ("SELECT %s FROM %s WHERE %s = %s" % 168 | (col, self.table.name, id_col, where_string)) 169 | 170 | return self._db.query(select, **kwargs) 171 | 172 | 173 | class QueryDb(object): 174 | """ 175 | A database object for interactively exploring a SQL database. 176 | """ 177 | def __init__(self, drivername=None, database=None, 178 | host=None, port=None, 179 | password=None, username=None, 180 | use_env_vars=True, demo=False): 181 | """ 182 | Initialize and test the connection. 183 | 184 | Kwargs: 185 | drivername (str): Drivername passed to sqlalchemy. 186 | 187 | database (str): Name of the database. 188 | 189 | host (str): IP address for the host to connect to. 190 | 191 | port (int): Port to connect on. 192 | 193 | username (str): Username for the database. 194 | 195 | password (str): Optionally specify a password. Defaults to None, 196 | which prompts the user for a password using getpass. 197 | 198 | use_env_vars (bool): Use environmental variables if specified? 199 | 200 | Returns: 201 | engine: The sqlalchemy database engine. 202 | 203 | Raises: 204 | OperationalError 205 | """ 206 | # Demo mode w/ included dummy database 207 | if demo: 208 | drivername = "sqlite" 209 | database = os.path.join( 210 | os.path.split(os.path.abspath(query.__file__))[0], 211 | "sample_data/Chinook_Sqlite.sqlite") 212 | use_env_vars = False 213 | 214 | # Check if the host, port. or database name options are overwritten 215 | # by environmental variables 216 | environ_driver = os.environ.get('QUERY_DB_DRIVER') 217 | environ_host = os.environ.get('QUERY_DB_HOST') 218 | environ_port = os.environ.get('QUERY_DB_PORT') 219 | environ_name = os.environ.get('QUERY_DB_NAME') 220 | if environ_driver is not None and use_env_vars: 221 | drivername = environ_driver 222 | if environ_host is not None and use_env_vars: 223 | host = environ_host 224 | if environ_port is not None and use_env_vars: 225 | port = environ_port 226 | if environ_name is not None and use_env_vars: 227 | database = environ_name 228 | 229 | # Note: This will require the user's terminal to be open. In the 230 | # case of IPython QtConsole or Notebook, this will be the terminal 231 | # from which the kernel was launched 232 | if password is None and drivername != "sqlite": # sqlite does not support pwds 233 | password = os.environ.get('QUERY_DB_PASS') 234 | if password is None: 235 | if pd.core.common.in_ipnb(): 236 | # Display a somewhat obnoxious warning to the user 237 | try: 238 | from IPython.display import display, HTML 239 | display(HTML(GETPASS_USE_WARNING)) 240 | except ImportError: 241 | pass 242 | password = getpass.getpass( 243 | "Please enter the %s server password:" % drivername) 244 | 245 | # Connection 246 | url = sqlalchemy.engine.url.URL( 247 | drivername=drivername, 248 | username=username, 249 | password=password, 250 | host=host, 251 | port=port, 252 | database=database) 253 | engine = sqlalchemy.create_engine(url) 254 | 255 | # Tests the connection 256 | with engine.begin(): 257 | pass 258 | 259 | # Set the engine ane metadata 260 | self._engine = engine 261 | self._summary_info = [] 262 | self._set_metadata() 263 | 264 | # Finally, set some pretty printing params 265 | # (schema diagram setup to go here) 266 | self._db_name = database.split("/")[-1].split(":")[0] 267 | self._summary_info = pd.DataFrame(self._summary_info, 268 | columns=["Table", "Primary Key(s)", 269 | "# of Columns", "# of Column Types"]) 270 | self._html = df_to_html(self._summary_info, "%s Database Summary" % self._db_name, 271 | bold=True) 272 | 273 | def _repr_html_(self): 274 | return self._html 275 | 276 | def __repr__(self): 277 | if self.test_connection(): 278 | c = "Working connection" 279 | else: 280 | c = "Inactive connection" 281 | return ("%s to a remote %s DB: %s" % 282 | (c, self._engine.name.upper(), self._db_name)) 283 | 284 | def test_connection(self): 285 | """ 286 | Test the connection to the QueryDb. Returns True if working. 287 | 288 | Returns: 289 | test_result (bool): Did the test pass? 290 | """ 291 | try: 292 | with self._engine.begin(): 293 | pass 294 | return True 295 | except sqlalchemy.exc.OperationalError: 296 | return False 297 | 298 | def query(self, sql_query, return_as="dataframe"): 299 | """ 300 | Execute a raw SQL query against the the SQL DB. 301 | 302 | Args: 303 | sql_query (str): A raw SQL query to execute. 304 | 305 | Kwargs: 306 | return_as (str): Specify what type of object should be 307 | returned. The following are acceptable types: 308 | - "dataframe": pandas.DataFrame or None if no matching query 309 | - "result": sqlalchemy.engine.result.ResultProxy 310 | 311 | Returns: 312 | result (pandas.DataFrame or sqlalchemy ResultProxy): Query result 313 | as a DataFrame (default) or sqlalchemy result (specified with 314 | return_as="result") 315 | 316 | Raises: 317 | QueryDbError 318 | """ 319 | if isinstance(sql_query, str): 320 | pass 321 | elif isinstance(sql_query, unicode): 322 | sql_query = str(sql_query) 323 | else: 324 | raise QueryDbError("query() requires a str or unicode input.") 325 | 326 | query = sqlalchemy.sql.text(sql_query) 327 | 328 | if return_as.upper() in ["DF", "DATAFRAME"]: 329 | return self._to_df(query, self._engine) 330 | elif return_as.upper() in ["RESULT", "RESULTPROXY"]: 331 | with self._engine.connect() as conn: 332 | result = conn.execute(query) 333 | return result 334 | else: 335 | raise QueryDbError("Other return types not implemented.") 336 | 337 | def _set_metadata(self): 338 | """ 339 | Internal helper to set metadata attributes. 340 | """ 341 | meta = QueryDbMeta() 342 | with self._engine.connect() as conn: 343 | meta.bind = conn 344 | meta.reflect() 345 | self._meta = meta 346 | 347 | # Set an inspect attribute, whose subattributes 348 | # return individual tables / columns. Tables and columns 349 | # are special classes with .last() and other convenience methods 350 | self.inspect = QueryDbAttributes() 351 | for table in self._meta.tables: 352 | setattr(self.inspect, table, 353 | QueryDbOrm(self._meta.tables[table], self)) 354 | 355 | table_attr = getattr(self.inspect, table) 356 | table_cols = table_attr.table.columns 357 | 358 | for col in table_cols.keys(): 359 | setattr(table_attr, col, 360 | QueryDbOrm(table_cols[col], self)) 361 | 362 | # Finally add some summary info: 363 | # Table name 364 | # Primary Key item or list 365 | # N of Cols 366 | # Distinct Col Values (class so NVARCHAR(20) and NVARCHAR(30) are not different) 367 | primary_keys = table_attr.table.primary_key.columns.keys() 368 | self._summary_info.append(( 369 | table, 370 | primary_keys[0] if len(primary_keys) == 1 else primary_keys, 371 | len(table_cols), 372 | len(set([x.type.__class__ for x in table_cols.values()])), 373 | )) 374 | 375 | def _to_df(self, query, conn, index_col=None, coerce_float=True, params=None, 376 | parse_dates=None, columns=None): 377 | """ 378 | Internal convert-to-DataFrame convenience wrapper. 379 | """ 380 | return pd.io.sql.read_sql(str(query), conn, index_col=index_col, 381 | coerce_float=coerce_float, params=params, 382 | parse_dates=parse_dates, columns=columns) 383 | -------------------------------------------------------------------------------- /query/helpers.py: -------------------------------------------------------------------------------- 1 | import query 2 | import os 3 | 4 | 5 | def setup_demo_env(): 6 | os.environ["QUERY_DB_DRIVER"] = "sqlite" 7 | os.environ["QUERY_DB_NAME"] = os.path.join( 8 | os.path.split(os.path.abspath(query.__file__))[0], 9 | "sample_data/Chinook_Sqlite.sqlite") 10 | 11 | if os.environ.get("QUERY_DB_HOST") is not None: 12 | os.environ.pop("QUERY_DB_HOST") 13 | if os.environ.get("QUERY_DB_PORT") is not None: 14 | os.environ.pop("QUERY_DB_PORT") 15 | -------------------------------------------------------------------------------- /query/html.py: -------------------------------------------------------------------------------- 1 | # Don't have lots of HTML polluting the core.py file 2 | 3 | # Styles/constants 4 | WARNING_BACKGROUND_COLOR = "#F6DF9F" 5 | WARNING_TEXT_COLOR = "#E18300" 6 | 7 | # Messages 8 | GETPASS_USE_WARNING = ("

" 10 | "Please enter your password in the terminal from which " 11 | "this IPython notebook was launched or pass a password " 12 | "explicitly to QueryDb()

" % 13 | (WARNING_TEXT_COLOR, WARNING_BACKGROUND_COLOR, WARNING_TEXT_COLOR)) 14 | 15 | QUERY_DB_ATTR_MSG = ("You're using a rich interactive terminal. " 16 | "Great! Just hit tab to see the tables and columns " 17 | "of this database.") 18 | 19 | 20 | # Functions 21 | def df_to_html(df, title, bold=False): 22 | if bold: 23 | style = 'font-weight: bold;' 24 | else: 25 | style = 'font-style: italic;' 26 | return ('
\n' + 27 | ('

' % style) + title + '

' + 28 | df.to_html(show_dimensions=False, max_rows=None, max_cols=None) + 29 | '\n
') 30 | -------------------------------------------------------------------------------- /query/sample_data/CHINOOK_INFO.md: -------------------------------------------------------------------------------- 1 | The code for generating the Chinook Database is available under the Microsoft Public License (Ms-PL) at http://chinookdatabase.codeplex.com/license. Premade databases may be downloaded at http://chinookdatabase.codeplex.com/releases/. This database is the SQLite version dated December 2, 2012. -------------------------------------------------------------------------------- /query/sample_data/Chinook_Sqlite.sqlite: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/boydgreenfield/query/03aa43b746b43832af3f0403265e648a5617b62b/query/sample_data/Chinook_Sqlite.sqlite -------------------------------------------------------------------------------- /requirements.txt: -------------------------------------------------------------------------------- 1 | SQLAlchemy>=1.3.0 2 | numpy==1.9.2 3 | pandas==0.16.0 4 | -------------------------------------------------------------------------------- /setup.py: -------------------------------------------------------------------------------- 1 | """ 2 | ``query`` 3 | --------- 4 | 5 | ``query`` is a simple module for quickly, interactively exploring a SQL 6 | database. Together with IPython, it supports quick tab-completion of table 7 | and column names, convenience methods for quickly looking at data (e.g., 8 | ``.head()``, ``.tail()``), and the ability to get a rich interactive database 9 | connection up in only 2 lines by setting a few required environmental 10 | variables. 11 | 12 | .. image:: https://travis-ci.org/boydgreenfield/query.svg?branch=v0.1.4 13 | 14 | 15 | Demo in 2 lines 16 | ``````````````` 17 | 18 | Explore the included demo database: 19 | 20 | .. code:: python 21 | 22 | from query import QueryDb 23 | db = QueryDb(demo=True) 24 | 25 | 26 | Real-world use case in 2 lines 27 | `````````````````````````````` 28 | 29 | Or set a few environmental variables (``QUERY_DB_DRIVER``, 30 | ``QUERY_DB_HOST``, ``QUERY_DB_PORT``, ``QUERY_DB_NAME``, and 31 | ``QUERY_DB_PASS``) and get started just as quickly: 32 | 33 | .. code:: python 34 | 35 | from query import QueryDB # capital 'B' is OK too :) 36 | db = QueryDB() 37 | 38 | 39 | Interactive example 40 | ``````````````````` 41 | .. image:: https://github.com/boydgreenfield/query/raw/v0.1.2/docs/images/interactive_demo.gif?raw=True 42 | 43 | 44 | 45 | Links 46 | ````` 47 | * `Code and additional details on Github: `_ 48 | 49 | """ 50 | from setuptools import setup 51 | 52 | 53 | setup( 54 | name='query', 55 | version='0.1.4', # When incrementing, 56 | # make sure to update Travis link above as well 57 | url='http://github.com/boydgreenfield/query/', 58 | license='MIT', 59 | author='Nick Boyd Greenfield', 60 | author_email='boyd.greenfield@gmail.com', 61 | description='Quick interactive exploration of SQL databases.', 62 | long_description=__doc__, 63 | packages=['query'], 64 | package_data={'query': ['sample_data/*.sqlite', 'sample_data/*.md']}, 65 | zip_safe=True, 66 | platforms='any', 67 | install_requires=[ 68 | 'pandas>=0.16', 69 | 'sqlalchemy>=1.3.0' 70 | ], 71 | classifiers=[ 72 | 'Development Status :: 4 - Beta', 73 | 'Framework :: IPython', 74 | 'Environment :: Web Environment', 75 | 'Intended Audience :: Developers', 76 | 'License :: OSI Approved :: MIT License', 77 | 'Operating System :: OS Independent', 78 | 'Programming Language :: Python', 79 | 'Topic :: Database', 80 | 'Topic :: Database :: Front-Ends' 81 | ], 82 | test_suite='nose.collector' 83 | ) 84 | -------------------------------------------------------------------------------- /test_requirements.txt: -------------------------------------------------------------------------------- 1 | # Test only requirements 2 | coverage==3.7.1 3 | ipython==3.1.0 4 | MySQL-python==1.2.5 5 | nose==1.3.6 6 | -------------------------------------------------------------------------------- /tests/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/boydgreenfield/query/03aa43b746b43832af3f0403265e648a5617b62b/tests/__init__.py -------------------------------------------------------------------------------- /tests/core_test.py: -------------------------------------------------------------------------------- 1 | from nose.tools import * # noqa 2 | from query.core import * # noqa 3 | from query.helpers import setup_demo_env 4 | import query 5 | import os 6 | import pandas as pd 7 | import sqlalchemy 8 | import warnings 9 | 10 | 11 | # Global setup function. Download DBs if missing. 12 | def my_setup(): 13 | # Set fake environmental variables, tests this functionality implicitly througout 14 | setup_demo_env() 15 | if not os.path.exists(os.environ.get("QUERY_DB_NAME")): 16 | raise Exception("Necessary Chinook SQLite test database not found.") 17 | 18 | 19 | @with_setup(my_setup) 20 | def test_querydb_init_sqlite(): 21 | # Test that init w/ environmental variables works 22 | db = QueryDb() 23 | assert db.__class__ is QueryDb 24 | 25 | # Test normal initialization 26 | db = QueryDb(drivername="sqlite", database="sample_data/Chinook_Sqlite.sqlite") 27 | 28 | # Basic proper initialization tests 29 | assert db.__class__ is QueryDb 30 | assert db.test_connection() 31 | assert db.__repr__() == "Working connection to a remote SQLITE DB: Chinook_Sqlite.sqlite" 32 | 33 | # And test if connection is broken / or _engine.begin() chokes 34 | def _raise(exc): 35 | raise exc 36 | 37 | db._engine.begin = lambda: _raise(sqlalchemy.exc.OperationalError("Bad test conn", "", "")) 38 | assert db.__repr__() == "Inactive connection to a remote SQLITE DB: Chinook_Sqlite.sqlite" 39 | assert True 40 | 41 | 42 | # Note no setup, as called in demo=True condition 43 | def test_querydb_demo(): 44 | db = QueryDb(demo=True) 45 | assert db.test_connection() 46 | 47 | 48 | @with_setup(my_setup) 49 | def test_getpass(): 50 | os.environ["QUERY_DB_DRIVER"] = "mysql" # Trick to prompt for PW, disabled for sqlite 51 | query.core.getpass.getpass = lambda _: "Fake testing password" 52 | with assert_raises(sqlalchemy.exc.OperationalError): # can't load bc no mysql db 53 | QueryDb() 54 | 55 | # Now test interactive path 56 | pd.core.common.in_ipnb = lambda: True 57 | with assert_raises(sqlalchemy.exc.OperationalError): # still fails bc mysql 58 | QueryDb() 59 | 60 | # And finally test graceful failure w/o IPython 61 | import sys 62 | ipy_module = sys.modules['IPython'] 63 | sys.modules['IPython'] = None 64 | with assert_raises(sqlalchemy.exc.OperationalError): # still fails bc mysql 65 | QueryDb() 66 | 67 | sys.modules['IPython'] = ipy_module 68 | 69 | 70 | @with_setup(my_setup) 71 | def test_querydborm(): 72 | db = QueryDb() 73 | assert db.inspect.__class__ is QueryDbAttributes 74 | assert db.inspect._repr_html_() == query.html.QUERY_DB_ATTR_MSG 75 | assert db._meta.__class__ is QueryDbMeta 76 | assert db.inspect.Track.__class__ is QueryDbOrm 77 | assert db.inspect.Track._repr_html_() == db.inspect.Track._html 78 | 79 | # Track is a sqlalchemy Table obj. 80 | assert db.inspect.Track.__repr__() == db.inspect.Track.table.__repr__() 81 | assert db.inspect.Track.Composer.__repr__() == db.inspect.Track.Composer.column.__repr__() 82 | 83 | 84 | @with_setup(my_setup) 85 | def test_querydb_query(): 86 | db = QueryDb() 87 | 88 | with assert_raises(sqlalchemy.exc.OperationalError): 89 | db.query("SELECT * FROM genres") # table is genre 90 | df = db.query("SELECT * FROM genre") 91 | assert df.__class__ is pd.DataFrame # Default result 92 | assert df.shape == (25, 2) 93 | 94 | res = db.query("SELECT * FROM genre", return_as="result") 95 | assert res.__class__ == sqlalchemy.engine.result.ResultProxy 96 | with assert_raises(sqlalchemy.exc.ProgrammingError): 97 | res.fetchall() # Fails for sqlite 98 | 99 | # And that those are all the values and properly returned by last 100 | # Note: last takes a DESC index and so we need to reverse the df here 101 | assert (df.Name.values[::-1] == db.inspect.Genre.last(25).Name.values).all() 102 | 103 | # And test failures: 104 | # - wrong input 105 | with assert_raises(QueryDbError): 106 | db.query(1) 107 | 108 | # - bad return type 109 | with assert_raises(QueryDbError): 110 | db.query("SELECT * FROM Tracks", return_as="junk") 111 | 112 | 113 | @with_setup(my_setup) 114 | def test_querydb_multi_primary_keys(): 115 | db = QueryDb() 116 | with warnings.catch_warnings(True) as w: 117 | db.inspect.PlaylistTrack.last() 118 | assert len(w) >= 1 119 | 120 | db.inspect.Track.table.primary_key.columns = sqlalchemy.sql.base.ColumnCollection() 121 | with assert_raises(NoPrimaryKeyException): 122 | db.inspect.Track.last() 123 | 124 | 125 | @with_setup(my_setup) 126 | def test_query_db_inspect(): 127 | db = QueryDb() 128 | 129 | # Table-wise 130 | assert db.inspect.Track.__class__ == QueryDbOrm 131 | assert db.inspect.Track.table.__class__ == sqlalchemy.sql.schema.Table 132 | 133 | # Column-wise 134 | assert db.inspect.Track.Composer.__class__ == QueryDbOrm # Potentially subclass in the future 135 | assert db.inspect.Track.Composer.table.__class__ == sqlalchemy.sql.schema.Table 136 | assert db.inspect.Track.Composer.column.__class__ == sqlalchemy.sql.schema.Column 137 | assert db.inspect.Track.Composer.column.name == "Composer" 138 | 139 | 140 | @with_setup(my_setup) 141 | def test_query_db_attr_methods(): 142 | db = QueryDb() 143 | 144 | # Test last method - note this column .Composer syntax is bc everything is a DataFrame 145 | assert db.inspect.Track.Composer.last().Composer.values[0] == 'Philip Glass' 146 | assert (db.inspect.Track.Composer.last().Composer.values[0] == 147 | db.inspect.Track.Composer.tail().Composer.values[0]) # equivalency 148 | 149 | # Test head method 150 | assert db.inspect.Genre.head().Name.values[0] == 'Rock' 151 | assert db.inspect.Genre.head().Name.values[1] == 'Jazz' 152 | 153 | # Test by= keyword 154 | assert (db.inspect.Track.head(n=5, by="Milliseconds").Milliseconds.values < 10000).all() 155 | 156 | # Test where 157 | assert len(db.inspect.Track.where("composer == 'Philip Glass'")) == 1 # :( 158 | -------------------------------------------------------------------------------- /tests/helpers_test.py: -------------------------------------------------------------------------------- 1 | from nose.tools import * # noqa 2 | from query.helpers import setup_demo_env 3 | import os 4 | 5 | 6 | def test_demo_setup(): 7 | setup_demo_env() 8 | assert os.environ.get("QUERY_DB_DRIVER") is "sqlite" 9 | assert os.environ.get("QUERY_DB_HOST") is None 10 | assert os.environ.get("QUERY_DB_PORT") is None 11 | 12 | # Test override of existing host/port params 13 | os.environ["QUERY_DB_HOST"] = "bad_host" 14 | os.environ["QUERY_DB_PORT"] = "9999" 15 | setup_demo_env() 16 | assert os.environ.get("QUERY_DB_HOST") is None 17 | assert os.environ.get("QUERY_DB_PORT") is None 18 | --------------------------------------------------------------------------------