├── .gitignore ├── .travis.yml ├── LICENSE ├── README.rst ├── TODO ├── aq ├── __init__.py ├── engines.py ├── errors.py ├── formatters.py ├── logger.py ├── parsers.py ├── prompt.py ├── select_parser.py ├── sqlite_util.py └── util.py ├── setup.cfg ├── setup.py └── tests ├── test_boto_engine.py ├── test_command_line_arg.py ├── test_parsers.py └── test_sqlite_util.py /.gitignore: -------------------------------------------------------------------------------- 1 | .idea 2 | 3 | # Byte-compiled / optimized / DLL files 4 | __pycache__/ 5 | *.py[cod] 6 | *$py.class 7 | 8 | # C extensions 9 | *.so 10 | 11 | # Distribution / packaging 12 | .Python 13 | env/ 14 | build/ 15 | develop-eggs/ 16 | dist/ 17 | downloads/ 18 | eggs/ 19 | .eggs/ 20 | lib/ 21 | lib64/ 22 | parts/ 23 | sdist/ 24 | var/ 25 | *.egg-info/ 26 | .installed.cfg 27 | *.egg 28 | 29 | # PyInstaller 30 | # Usually these files are written by a python script from a template 31 | # before PyInstaller builds the exe, so as to inject date/other infos into it. 32 | *.manifest 33 | *.spec 34 | 35 | # Installer logs 36 | pip-log.txt 37 | pip-delete-this-directory.txt 38 | 39 | # Unit test / coverage reports 40 | htmlcov/ 41 | .tox/ 42 | .coverage 43 | .coverage.* 44 | .cache 45 | nosetests.xml 46 | coverage.xml 47 | *,cover 48 | .hypothesis/ 49 | 50 | # Translations 51 | *.mo 52 | *.pot 53 | 54 | # Django stuff: 55 | *.log 56 | 57 | # Sphinx documentation 58 | docs/_build/ 59 | 60 | # PyBuilder 61 | target/ 62 | 63 | #Ipython Notebook 64 | .ipynb_checkpoints 65 | 66 | # Virtual env 67 | venv 68 | -------------------------------------------------------------------------------- /.travis.yml: -------------------------------------------------------------------------------- 1 | language: python 2 | python: 3 | - "2.6" 4 | - "2.7" 5 | - "3.3" 6 | - "3.4" 7 | - "3.5" 8 | install: 9 | - pip install . 10 | script: nosetests 11 | -------------------------------------------------------------------------------- /LICENSE: -------------------------------------------------------------------------------- 1 | The MIT License (MIT) 2 | 3 | Copyright (c) 2016 Binh Le 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 | -------------------------------------------------------------------------------- /README.rst: -------------------------------------------------------------------------------- 1 | ================================= 2 | aq - Query AWS resources with SQL 3 | ================================= 4 | 5 | ``aq`` allows you to query your AWS resources (EC2 instances, S3 buckets, etc.) with plain SQL. 6 | 7 | .. image:: https://travis-ci.org/lebinh/aq.svg?branch=master 8 | :target: https://travis-ci.org/lebinh/aq 9 | 10 | .. image:: https://asciinema.org/a/79468.png 11 | :target: https://asciinema.org/a/79468 12 | 13 | *But why?* 14 | 15 | Fun, mostly fun. But see sample queries below for useful queries that can be performed with ``aq``. 16 | 17 | Usage 18 | ~~~~~ 19 | :: 20 | 21 | Usage: 22 | aq [--profile=] [--region=] [--table-cache-ttl=] [-v] [--debug] 23 | aq [--profile=] [--region=] [--table-cache-ttl=] [-v] [--debug] 24 | 25 | Options: 26 | --profile= Use a specific profile from your credential file 27 | --region= The region to use. Overrides config/env settings 28 | --table-cache-ttl= number of seconds to cache the tables 29 | before we update them from AWS again [default: 300] 30 | -v, --verbose enable verbose logging 31 | 32 | Running ``aq`` without specifying any query will start a REPL to run your queries interactively. 33 | 34 | Sample queries 35 | ~~~~~~~~~~~~~~ 36 | 37 | One of the most important benefit of being able to query which SQL is aggregation and join, 38 | which can be very complicated or even impossible to do with AWS CLI. 39 | 40 | To count how many running instances per instance type 41 | ----------------------------------------------------- 42 | :: 43 | 44 | > SELECT instance_type, count(*) count 45 | FROM ec2_instances 46 | WHERE state->'Name' = 'running' 47 | GROUP BY instance_type 48 | ORDER BY count DESC 49 | +-----------------+---------+ 50 | | instance_type | count | 51 | |-----------------+---------| 52 | | m4.2xlarge | 15 | 53 | | m4.xlarge | 6 | 54 | | r3.8xlarge | 6 | 55 | +-----------------+---------+ 56 | 57 | Find instances with largest attached EBS volumes size 58 | ----------------------------------------------------- 59 | :: 60 | 61 | > SELECT i.id, i.tags->'Name' name, count(v.id) vols, sum(v.size) size, sum(v.iops) iops 62 | FROM ec2_instances i 63 | JOIN ec2_volumes v ON v.attachments -> 0 -> 'InstanceId' = i.id 64 | GROUP BY i.id 65 | ORDER BY size DESC 66 | LIMIT 3 67 | +------------+-----------+--------+--------+--------+ 68 | | id | name | vols | size | iops | 69 | |------------+-----------+--------+--------+--------| 70 | | i-12345678 | foo | 4 | 2000 | 4500 | 71 | | i-12345679 | bar | 2 | 332 | 1000 | 72 | | i-12345687 | blah | 1 | 320 | 960 | 73 | +------------+-----------+--------+--------+--------+ 74 | 75 | Find instances that allows access to port 22 in their security groups 76 | --------------------------------------------------------------------- 77 | :: 78 | 79 | > SELECT i.id, i.tags->'Name' name, sg.group_name 80 | FROM ec2_instances i 81 | JOIN ec2_security_groups sg ON instr(i.security_groups, sg.id) 82 | WHERE instr(sg.ip_permissions, '"ToPort": 22,') 83 | +------------+-----------+---------------------+ 84 | | id | name | group_name | 85 | |------------+-----------+---------------------| 86 | | i-foobar78 | foobar | launch-wizard-1 | 87 | | i-foobar87 | blah | launch-wizard-2 | 88 | +------------+-----------+---------------------+ 89 | 90 | AWS Credential 91 | ~~~~~~~~~~~~~~ 92 | 93 | ``aq`` relies on ``boto3`` for AWS API access so all the 94 | `credential configuration mechanisms `_ 95 | of boto3 will work. If you are using the AWS CLI then you can use ``aq`` without any further configurations. 96 | 97 | Available tables 98 | ~~~~~~~~~~~~~~~~ 99 | 100 | AWS resources are specified as table names in ``_`` format with: 101 | 102 | resource 103 | one of the `resources `_ 104 | defined in boto3: ``ec2``, ``s3``, ``iam``, etc. 105 | collection 106 | one of the resource's `collections `_ 107 | defined in boto3: ``instances``, ``images``, etc. 108 | 109 | An optional schema (i.e. database) name can be used to specify the AWS region to query. 110 | If you don't specify the schema name then boto's default region will be used. 111 | 112 | :: 113 | 114 | -- to count the number of ec2 instances in AWS Singapore region 115 | SELECT count(*) FROM ap_southeast_1.ec2_instances 116 | 117 | Note that the region name is specified using underscore (``ap_southeast_1``) instead of dash (``ap-southeast-1``). 118 | 119 | At the moment the full table list for AWS ``us_east_1`` region is 120 | 121 | .. list-table:: 122 | 123 | * - cloudformation_stacks 124 | * - cloudwatch_alarms 125 | * - cloudwatch_metrics 126 | * - dynamodb_tables 127 | * - ec2_classic_addresses 128 | * - ec2_dhcp_options_sets 129 | * - ec2_images 130 | * - ec2_instances 131 | * - ec2_internet_gateways 132 | * - ec2_key_pairs 133 | * - ec2_network_acls 134 | * - ec2_network_interfaces 135 | * - ec2_placement_groups 136 | * - ec2_route_tables 137 | * - ec2_security_groups 138 | * - ec2_snapshots 139 | * - ec2_subnets 140 | * - ec2_volumes 141 | * - ec2_vpc_addresses 142 | * - ec2_vpc_peering_connections 143 | * - ec2_vpcs 144 | * - glacier_vaults 145 | * - iam_groups 146 | * - iam_instance_profiles 147 | * - iam_policies 148 | * - iam_roles 149 | * - iam_saml_providers 150 | * - iam_server_certificates 151 | * - iam_users 152 | * - iam_virtual_mfa_devices 153 | * - opsworks_stacks 154 | * - s3_buckets 155 | * - sns_platform_applications 156 | * - sns_subscriptions 157 | * - sns_topics 158 | * - sqs_queues 159 | 160 | Query with structured value 161 | ~~~~~~~~~~~~~~~~~~~~~~~~~~~ 162 | 163 | Quite a number of resource contain structured value (e.g. instance tags) that cannot be use directly in SQL. 164 | We keep and present these values as JSON serialized string and add a new operator ``->`` to make querying on them easier. 165 | The ``->`` (replaced to ``json_get`` before execution) can be used to access an object field, ``object->'fieldName'``, or access 166 | an array item, ``array->index``:: 167 | 168 | > SELECT '{"foo": "bar"}' -> 'foo' 169 | +-------------------------------------+ 170 | | json_get('{"foo": "bar"}', 'foo') | 171 | |-------------------------------------| 172 | | bar | 173 | +-------------------------------------+ 174 | > SELECT '["foo", "bar", "blah"]' -> 1 175 | +--------------+ 176 | | json_get(' | 177 | |--------------| 178 | | bar | 179 | +--------------+ 180 | 181 | Install 182 | ~~~~~~~ 183 | :: 184 | 185 | pip install aq 186 | 187 | Tests (with `nose`) 188 | ~~~~~~~~~~~~~~~~~~~ 189 | :: 190 | 191 | nosetests 192 | -------------------------------------------------------------------------------- /TODO: -------------------------------------------------------------------------------- 1 | TODO 2 | ===== 3 | 4 | [x] Chaining of json_get operator '->' 5 | i.e. foo -> bar -> blah = ((foo -> bar) -> blah) = json_get(json_get(foo, bar), blah) 6 | json_field_access = expr -> (integer | string) 7 | [x] Query with region as database name 8 | [x] REPL auto completion 9 | [x] Determine freshness of tables to avoid re-loading everytime 10 | [x] Mechanism to inspect available tables [available via auto-complete] 11 | [ ] Mechanism to inspect table schema 12 | 13 | 14 | MAYBE WILL DO 15 | ============== 16 | 17 | [ ] Constant, pre-loaded tables (e.g. ec2 instances pricing) 18 | [ ] REPL multi-lines input 19 | [ ] Better error msg on query parsing error (with viz of error position) 20 | -------------------------------------------------------------------------------- /aq/__init__.py: -------------------------------------------------------------------------------- 1 | """aq - Query AWS resources with SQL 2 | 3 | Usage: 4 | aq [--profile=] [--region=] [--table-cache-ttl=] [-v] [--debug] 5 | aq [--profile=] [--region=] [--table-cache-ttl=] [-v] [--debug] 6 | 7 | Sample queries: 8 | aq "select tags->'Name' from ec2_instances" 9 | aq "select count(*) from us_west_1.ec2_instances" 10 | 11 | Options: 12 | --profile= Use a specific profile from your credential file 13 | --region= The region to use. Overrides config/env settings 14 | --table-cache-ttl= number of seconds to cache the tables 15 | before we update them from AWS again [default: 300] 16 | -v, --verbose enable verbose logging 17 | --debug enable debug mode 18 | """ 19 | from __future__ import print_function 20 | 21 | import traceback 22 | from collections import namedtuple 23 | 24 | from docopt import docopt 25 | 26 | from aq.engines import BotoSqliteEngine 27 | from aq.errors import QueryError 28 | from aq.formatters import TableFormatter 29 | from aq.logger import initialize_logger 30 | from aq.parsers import SelectParser 31 | from aq.prompt import AqPrompt 32 | 33 | __version__ = '0.1.1' 34 | 35 | QueryResult = namedtuple('QueryResult', ('parsed_query', 'query_metadata', 'columns', 'rows')) 36 | 37 | 38 | def get_engine(options): 39 | return BotoSqliteEngine(options) 40 | 41 | 42 | def get_parser(options): 43 | return SelectParser(options) 44 | 45 | 46 | def get_formatter(options): 47 | return TableFormatter(options) 48 | 49 | 50 | def get_prompt(parser, engine, options): 51 | return AqPrompt(parser, engine, options) 52 | 53 | 54 | def main(): 55 | args = docopt(__doc__) 56 | initialize_logger(verbose=args['--verbose'], debug=args['--debug']) 57 | 58 | parser = get_parser(args) 59 | engine = get_engine(args) 60 | formatter = get_formatter(args) 61 | 62 | if args['']: 63 | query = args[''] 64 | res = execute_query(engine, formatter, parser, query) 65 | print(formatter.format(res.columns, res.rows)) 66 | else: 67 | repl = get_prompt(parser, engine, args) 68 | while True: 69 | try: 70 | query = repl.prompt() 71 | res = execute_query(engine, formatter, parser, query) 72 | print(formatter.format(res.columns, res.rows)) 73 | repl.update_with_result(res.query_metadata) 74 | except EOFError: 75 | break 76 | except QueryError as e: 77 | print('QueryError: {0}'.format(e)) 78 | except: 79 | traceback.print_exc() 80 | 81 | 82 | def execute_query(engine, formatter, parser, query): 83 | parsed_query, metadata = parser.parse_query(query) 84 | columns, rows = engine.execute(parsed_query, metadata) 85 | return QueryResult(parsed_query=parsed_query, query_metadata=metadata, 86 | columns=columns, rows=rows) 87 | -------------------------------------------------------------------------------- /aq/engines.py: -------------------------------------------------------------------------------- 1 | import itertools 2 | import os.path 3 | import pprint 4 | import sqlite3 5 | import time 6 | from collections import defaultdict 7 | from multiprocessing.dummy import Pool 8 | 9 | import boto3 10 | from boto3.resources.collection import CollectionManager 11 | from botocore.exceptions import NoCredentialsError 12 | 13 | from aq import logger, util, sqlite_util 14 | from aq.errors import QueryError 15 | 16 | DEFAULT_REGION = 'us_east_1' 17 | 18 | LOGGER = logger.get_logger() 19 | 20 | 21 | class BotoSqliteEngine(object): 22 | def __init__(self, options=None): 23 | self.options = options if options else {} 24 | self.debug = options.get('--debug', False) 25 | 26 | self.profile = options.get('--profile', None) 27 | self.region = options.get('--region', None) 28 | self.table_cache_ttl = int(options.get('--table-cache-ttl', 300)) 29 | self.last_refresh_time = defaultdict(int) 30 | 31 | self.boto3_session = boto3.Session(profile_name=self.profile) 32 | # dash (-) is not allowed in database name so we use underscore (_) instead in region name 33 | # throughout this module region name will *always* use underscore 34 | if self.region: 35 | self.default_region = self.region.replace('-', '_') 36 | elif self.boto3_session.region_name: 37 | self.default_region = self.boto3_session.region_name.replace('-', '_') 38 | else: 39 | self.default_region = DEFAULT_REGION 40 | 41 | self.boto3_session = boto3.Session(profile_name=self.profile, region_name=self.default_region.replace('_', '-')) 42 | self.db = self.init_db() 43 | # attach the default region too 44 | self.attach_region(self.default_region) 45 | 46 | def init_db(self): 47 | util.ensure_data_dir_exists() 48 | db_path = '~/.aq/{0}.db'.format(self.default_region) 49 | absolute_path = os.path.expanduser(db_path) 50 | return sqlite_util.connect(absolute_path) 51 | 52 | def execute(self, query, metadata): 53 | LOGGER.info('Executing query: %s', query) 54 | self.load_tables(query, metadata) 55 | try: 56 | cursor = self.db.execute(query) 57 | except sqlite3.OperationalError as e: 58 | raise QueryError(str(e)) 59 | columns = [d[0] for d in cursor.description] 60 | rows = cursor.fetchall() 61 | return columns, rows 62 | 63 | def load_tables(self, query, meta): 64 | """ 65 | Load necessary resources tables into db to execute given query. 66 | """ 67 | try: 68 | for table in meta.tables: 69 | self.load_table(table) 70 | except NoCredentialsError: 71 | help_link = 'http://boto3.readthedocs.io/en/latest/guide/configuration.html' 72 | raise QueryError('Unable to locate AWS credential. ' 73 | 'Please see {0} on how to configure AWS credential.'.format(help_link)) 74 | 75 | def load_table(self, table): 76 | """ 77 | Load resources as specified by given table into our db. 78 | """ 79 | region = table.database if table.database else self.default_region 80 | resource_name, collection_name = table.table.split('_', 1) 81 | # we use underscore "_" instead of dash "-" for region name but boto3 need dash 82 | boto_region_name = region.replace('_', '-') 83 | resource = self.boto3_session.resource(resource_name, region_name=boto_region_name) 84 | if not hasattr(resource, collection_name): 85 | raise QueryError( 86 | 'Unknown collection <{0}> of resource <{1}>'.format(collection_name, resource_name)) 87 | 88 | self.attach_region(region) 89 | self.refresh_table(region, table.table, resource, getattr(resource, collection_name)) 90 | 91 | def attach_region(self, region): 92 | if not self.is_attached_region(region): 93 | LOGGER.info('Attaching new database for region: %s', region) 94 | region_db_file_path = '~/.aq/{0}.db'.format(region) 95 | absolute_path = os.path.expanduser(region_db_file_path) 96 | self.db.execute('ATTACH DATABASE ? AS ?', (absolute_path, region)) 97 | 98 | def is_attached_region(self, region): 99 | databases = self.db.execute('PRAGMA database_list') 100 | db_names = (db[1] for db in databases) 101 | return region in db_names 102 | 103 | def refresh_table(self, schema_name, table_name, resource, collection): 104 | if not self.is_fresh_enough(schema_name, table_name): 105 | LOGGER.info('Refreshing table: %s.%s', schema_name, table_name) 106 | columns = get_columns_list(resource, collection) 107 | LOGGER.info('Columns list: %s', columns) 108 | with self.db: 109 | sqlite_util.create_table(self.db, schema_name, table_name, columns) 110 | items = collection.all() 111 | # special treatment for tags field 112 | items = [convert_tags_to_dict(item) for item in items] 113 | sqlite_util.insert_all(self.db, schema_name, table_name, columns, items) 114 | self.last_refresh_time[(schema_name, table_name)] = time.time() 115 | 116 | def is_fresh_enough(self, schema_name, table_name): 117 | last_refresh = self.last_refresh_time[(schema_name, table_name)] 118 | age = time.time() - last_refresh 119 | return age < self.table_cache_ttl 120 | 121 | @property 122 | def available_schemas(self): 123 | # we want to return all regions if possible so ec2 is a good enough guess 124 | regions = self.boto3_session.get_available_regions(service_name='ec2') 125 | return [r.replace('-', '_') for r in regions] 126 | 127 | @property 128 | def available_tables(self): 129 | resources = self.boto3_session.get_available_resources() 130 | tables = Pool(processes=len(resources)).map(self._get_table_names_for_resource, resources) 131 | return itertools.chain.from_iterable(tables) 132 | 133 | def _get_table_names_for_resource(self, resource_name): 134 | resource = self.boto3_session.resource(resource_name) 135 | for attr in dir(resource): 136 | if isinstance(getattr(resource, attr), CollectionManager): 137 | yield '{0}_{1}'.format(resource_name, attr) 138 | 139 | 140 | class ObjectProxy(object): 141 | def __init__(self, source, **replaced_fields): 142 | self.source = source 143 | self.replaced_fields = replaced_fields 144 | 145 | def __getattr__(self, item): 146 | if item in self.replaced_fields: 147 | return self.replaced_fields[item] 148 | return getattr(self.source, item) 149 | 150 | 151 | def convert_tags_to_dict(item): 152 | """ 153 | Convert AWS inconvenient tags model of a list of {"Key": , "Value": } pairs 154 | to a dict of {: } for easier querying. 155 | 156 | This returns a proxied object over given item to return a different tags format as the tags 157 | attribute is read-only and we cannot modify it directly. 158 | """ 159 | if hasattr(item, 'tags'): 160 | tags = item.tags 161 | if isinstance(tags, list): 162 | tags_dict = {} 163 | for kv_dict in tags: 164 | if isinstance(kv_dict, dict) and 'Key' in kv_dict and 'Value' in kv_dict: 165 | tags_dict[kv_dict['Key']] = kv_dict['Value'] 166 | return ObjectProxy(item, tags=tags_dict) 167 | return item 168 | 169 | 170 | def get_resource_model_attributes(resource, collection): 171 | service_model = resource.meta.client.meta.service_model 172 | resource_model = get_resource_model(collection) 173 | shape_name = resource_model.shape 174 | shape = service_model.shape_for(shape_name) 175 | return resource_model.get_attributes(shape) 176 | 177 | 178 | def get_columns_list(resource, collection): 179 | resource_model = get_resource_model(collection) 180 | LOGGER.debug('Resource model: %s', resource_model) 181 | 182 | identifiers = sorted(i.name for i in resource_model.identifiers) 183 | LOGGER.debug('Model identifiers: %s', identifiers) 184 | 185 | attributes = get_resource_model_attributes(resource, collection) 186 | LOGGER.debug('Model attributes: %s', pprint.pformat(attributes)) 187 | 188 | return list(itertools.chain(identifiers, attributes)) 189 | 190 | 191 | def get_resource_model(collection): 192 | return collection._model.resource.model 193 | -------------------------------------------------------------------------------- /aq/errors.py: -------------------------------------------------------------------------------- 1 | class AQError(Exception): 2 | pass 3 | 4 | 5 | class QueryError(AQError): 6 | pass 7 | 8 | 9 | class QueryParsingError(AQError): 10 | pass 11 | -------------------------------------------------------------------------------- /aq/formatters.py: -------------------------------------------------------------------------------- 1 | from tabulate import tabulate 2 | 3 | 4 | class TableFormatter(object): 5 | def __init__(self, options=None): 6 | self.options = options if options else {} 7 | 8 | @staticmethod 9 | def format(columns, rows): 10 | return tabulate(rows, headers=columns, tablefmt='psql', missingval='NULL') 11 | -------------------------------------------------------------------------------- /aq/logger.py: -------------------------------------------------------------------------------- 1 | import logging 2 | 3 | import sys 4 | 5 | _logger = None 6 | 7 | 8 | def get_logger(): 9 | global _logger 10 | if _logger: 11 | return _logger 12 | 13 | formatter = logging.Formatter(fmt='%(levelname)s: %(message)s') 14 | handler = logging.StreamHandler(sys.stderr) 15 | handler.setFormatter(formatter) 16 | _logger = logging.getLogger('aq') 17 | _logger.addHandler(handler) 18 | return _logger 19 | 20 | 21 | def initialize_logger(verbose=False, debug=False): 22 | level = logging.INFO if verbose else logging.WARNING 23 | if debug: 24 | level = logging.DEBUG 25 | get_logger().setLevel(level) 26 | -------------------------------------------------------------------------------- /aq/parsers.py: -------------------------------------------------------------------------------- 1 | import collections 2 | from collections import namedtuple 3 | 4 | from six import string_types 5 | 6 | from aq.errors import QueryParsingError 7 | from aq.select_parser import select_stmt, ParseException 8 | 9 | TableId = namedtuple('TableId', ('database', 'table', 'alias')) 10 | QueryMetadata = namedtuple('QueryMetadata', ('tables',)) 11 | 12 | 13 | class SelectParser(object): 14 | def __init__(self, options): 15 | self.options = options 16 | 17 | @staticmethod 18 | def parse_query(query): 19 | try: 20 | parse_result = select_stmt.parseString(query, parseAll=True) 21 | except ParseException as e: 22 | raise QueryParsingError(e) 23 | 24 | tables = [parse_table_id(tid) for tid in parse_result.table_ids] 25 | parsed_query = concat(parse_result) 26 | return parsed_query, QueryMetadata(tables=tables) 27 | 28 | 29 | def parse_table_id(table_id): 30 | database = table_id.database[0] if table_id.database else None 31 | table = table_id.table[0] if table_id.table else None 32 | alias = table_id.alias[0] if table_id.alias else None 33 | return TableId(database, table, alias) 34 | 35 | 36 | def flatten(nested_list): 37 | for item in nested_list: 38 | if isinstance(item, collections.Iterable) and not isinstance(item, string_types): 39 | for nested_item in flatten(item): 40 | yield nested_item 41 | else: 42 | yield item 43 | 44 | 45 | def concat(tokens): 46 | return ' '.join(flatten(tokens)) 47 | -------------------------------------------------------------------------------- /aq/prompt.py: -------------------------------------------------------------------------------- 1 | from __future__ import unicode_literals 2 | 3 | import itertools 4 | import os 5 | 6 | from prompt_toolkit import AbortAction, CommandLineInterface 7 | from prompt_toolkit.auto_suggest import AutoSuggestFromHistory 8 | from prompt_toolkit.completion import Completion, Completer 9 | from prompt_toolkit.history import FileHistory 10 | from prompt_toolkit.layout.lexers import PygmentsLexer 11 | from prompt_toolkit.shortcuts import create_prompt_application, create_eventloop 12 | from prompt_toolkit.validation import Validator, ValidationError 13 | from pygments.lexers.sql import SqlLexer 14 | 15 | from aq import util 16 | from aq.errors import QueryParsingError 17 | from aq.logger import get_logger 18 | 19 | LOGGER = get_logger() 20 | 21 | 22 | class AqPrompt(object): 23 | def __init__(self, parser, engine, options=None): 24 | self.parser = parser 25 | self.engine = engine 26 | self.options = options if options is not None else {} 27 | util.ensure_data_dir_exists() 28 | application = create_prompt_application( 29 | message='> ', 30 | lexer=PygmentsLexer(SqlLexer), 31 | history=FileHistory(os.path.expanduser('~/.aq/history')), 32 | completer=AqCompleter(schemas=engine.available_schemas, tables=engine.available_tables), 33 | auto_suggest=AutoSuggestFromHistory(), 34 | validator=QueryValidator(parser), 35 | on_abort=AbortAction.RETRY, 36 | ) 37 | loop = create_eventloop() 38 | self.cli = CommandLineInterface(application=application, eventloop=loop) 39 | self.patch_context = self.cli.patch_stdout_context() 40 | 41 | def prompt(self): 42 | with self.patch_context: 43 | return self.cli.run(reset_current_buffer=True).text 44 | 45 | def update_with_result(self, query_metadata): 46 | # TODO 47 | pass 48 | 49 | 50 | class AqCompleter(Completer): 51 | keywords = '''UNION, ALL, AND, INTERSECT, EXCEPT, COLLATE, ASC, DESC, ON, USING, 52 | NATURAL, INNER, CROSS, LEFT, OUTER, JOIN, AS, INDEXED, NOT, SELECT, DISTINCT, 53 | FROM, WHERE, GROUP, BY, HAVING, ORDER, BY, LIMIT, OFFSET CAST, ISNULL, NOTNULL, 54 | NULL, IS, BETWEEN, ELSE, END, CASE, WHEN, THEN, EXISTS, COLLATE, IN, LIKE, GLOB, 55 | REGEXP, MATCH, ESCAPE, CURRENT_TIME, CURRENT_DATE, CURRENT_TIMESTAMP 56 | '''.replace(',', '').split() 57 | 58 | functions = '''avg, count, max, min, sum, json_get'''.replace(',', '').split() 59 | 60 | starters = ['SELECT'] 61 | 62 | def __init__(self, schemas=None, tables=None): 63 | self.schemas = schemas if schemas else [] 64 | self.tables = tables if tables else [] 65 | self.tables_and_schemas = itertools.chain(self.schemas, self.tables) 66 | self.all_completions = itertools.chain(self.keywords, self.functions, 67 | self.tables_and_schemas) 68 | 69 | def get_completions(self, document, complete_event): 70 | start_of_current_word = document.find_start_of_previous_word(1) 71 | current_word = document.text_before_cursor[start_of_current_word:].strip().lower() 72 | 73 | start_of_previous_2_words = document.find_start_of_previous_word(2) 74 | previous_word = document.text_before_cursor[ 75 | start_of_previous_2_words:start_of_current_word].strip().lower() 76 | 77 | if document.text_before_cursor[-1:].isspace(): 78 | previous_word = current_word 79 | current_word = '' 80 | 81 | candidates = self.get_completion_candidates(current_word, previous_word, document) 82 | for candidate in candidates: 83 | if candidate.lower().startswith(current_word): 84 | yield Completion(candidate, -len(current_word)) 85 | 86 | def get_completion_candidates(self, current_word, previous_word, document): 87 | if not previous_word: 88 | return self.starters 89 | 90 | if current_word == ',' or previous_word in [',', 'from', 'join']: 91 | # we delay the materialize of table list until here for faster startup time 92 | if not isinstance(self.tables_and_schemas, list): 93 | self.tables_and_schemas = list(self.tables_and_schemas) 94 | return self.tables_and_schemas 95 | 96 | if not isinstance(self.all_completions, list): 97 | # ensure we materialized the table list first 98 | if not isinstance(self.tables_and_schemas, list): 99 | self.tables_and_schemas = list(self.tables_and_schemas) 100 | self.all_completions = list(self.all_completions) 101 | return self.all_completions 102 | 103 | 104 | class QueryValidator(Validator): 105 | def __init__(self, parser): 106 | self.parser = parser 107 | 108 | def validate(self, document): 109 | try: 110 | self.parser.parse_query(document.text) 111 | except QueryParsingError as e: 112 | raise ValidationError(message='Invalid SQL query. {0}'.format(e), 113 | cursor_position=document.cursor_position) 114 | -------------------------------------------------------------------------------- /aq/select_parser.py: -------------------------------------------------------------------------------- 1 | # Adapted from select_parser.py by Paul McGuire 2 | # http://pyparsing.wikispaces.com/file/view/select_parser.py/158651233/select_parser.py 3 | # 4 | # a simple SELECT statement parser, taken from SQLite's SELECT statement 5 | # definition at http://www.sqlite.org/lang_select.html 6 | # 7 | from pyparsing import * 8 | 9 | ParserElement.enablePackrat() 10 | 11 | 12 | def no_suppress_delimited_list(expression, delimiter=','): 13 | return expression + ZeroOrMore(delimiter + expression) 14 | 15 | 16 | def concat(tokens): 17 | return ''.join(tokens) 18 | 19 | 20 | def build_json_get_expr(terms): 21 | if len(terms) < 2: 22 | raise ValueError('Not enough terms') 23 | if len(terms) == 2: 24 | return 'json_get({0}, {1})'.format(terms[0], terms[1]) 25 | return 'json_get({0}, {1})'.format(build_json_get_expr(terms[:-1]), terms[-1]) 26 | 27 | 28 | def replace_json_get(tokens): 29 | terms = [t for t in tokens[0] if t != '->'] 30 | return build_json_get_expr(terms) 31 | 32 | 33 | # keywords 34 | (UNION, ALL, AND, INTERSECT, EXCEPT, COLLATE, ASC, DESC, ON, USING, NATURAL, INNER, 35 | CROSS, LEFT, OUTER, JOIN, AS, INDEXED, NOT, SELECT, DISTINCT, FROM, WHERE, GROUP, BY, 36 | HAVING, ORDER, BY, LIMIT, OFFSET, OR) = map(CaselessKeyword, """UNION, ALL, AND, INTERSECT, 37 | EXCEPT, COLLATE, ASC, DESC, ON, USING, NATURAL, INNER, CROSS, LEFT, OUTER, JOIN, AS, INDEXED, 38 | NOT, SELECT, DISTINCT, FROM, WHERE, GROUP, BY, HAVING, ORDER, BY, LIMIT, OFFSET, OR 39 | """.replace(",", "").split()) 40 | 41 | (CAST, ISNULL, NOTNULL, NULL, IS, BETWEEN, ELSE, END, CASE, WHEN, THEN, EXISTS, 42 | COLLATE, IN, LIKE, GLOB, REGEXP, MATCH, ESCAPE, CURRENT_TIME, CURRENT_DATE, 43 | CURRENT_TIMESTAMP) = map(CaselessKeyword, """CAST, ISNULL, NOTNULL, NULL, IS, BETWEEN, ELSE, 44 | END, CASE, WHEN, THEN, EXISTS, COLLATE, IN, LIKE, GLOB, REGEXP, MATCH, ESCAPE, 45 | CURRENT_TIME, CURRENT_DATE, CURRENT_TIMESTAMP""".replace(",", "").split()) 46 | 47 | keyword = MatchFirst((UNION, ALL, INTERSECT, EXCEPT, COLLATE, ASC, DESC, ON, USING, NATURAL, INNER, 48 | CROSS, LEFT, OUTER, JOIN, AS, INDEXED, NOT, SELECT, DISTINCT, FROM, WHERE, 49 | GROUP, BY, 50 | HAVING, ORDER, BY, LIMIT, OFFSET, CAST, ISNULL, NOTNULL, NULL, IS, BETWEEN, 51 | ELSE, END, CASE, WHEN, THEN, EXISTS, 52 | COLLATE, IN, LIKE, GLOB, REGEXP, MATCH, ESCAPE, CURRENT_TIME, CURRENT_DATE, 53 | CURRENT_TIMESTAMP)) 54 | 55 | identifier = ~keyword + Word(alphas, alphanums + "_") 56 | collation_name = identifier.copy() 57 | column_name = identifier.copy() 58 | column_alias = identifier.copy() 59 | table_name = identifier.copy() 60 | table_alias = identifier.copy() 61 | index_name = identifier.copy() 62 | function_name = identifier.copy() 63 | parameter_name = identifier.copy() 64 | database_name = identifier.copy() 65 | 66 | # expression 67 | LPAR, RPAR, COMMA = map(Word, "(),") 68 | select_stmt = Forward() 69 | expr = Forward() 70 | 71 | integer = Regex(r"[+-]?\d+") 72 | numeric_literal = Regex(r"\d+(\.\d*)?([eE][+-]?\d+)?") 73 | string_literal = QuotedString("'", unquoteResults=False) 74 | blob_literal = Combine(oneOf("x X") + "'" + Word(hexnums) + "'") 75 | literal_value = (numeric_literal | string_literal | blob_literal | 76 | NULL | CURRENT_TIME | CURRENT_DATE | CURRENT_TIMESTAMP) 77 | bind_parameter = ( 78 | Word("?", nums) | 79 | Combine(oneOf(": @ $") + parameter_name) 80 | ) 81 | type_name = oneOf("TEXT REAL INTEGER BLOB NULL") 82 | 83 | expr_term = ( 84 | CAST + LPAR + expr + AS + type_name + RPAR | 85 | EXISTS + LPAR + select_stmt + RPAR | 86 | function_name + LPAR + Optional(no_suppress_delimited_list(expr) | "*") + RPAR | 87 | literal_value | 88 | bind_parameter | 89 | (database_name + "." + table_name + "." + identifier) | 90 | (table_name + "." + identifier) | 91 | identifier 92 | ).setParseAction(concat) 93 | 94 | UNARY, BINARY, TERNARY = 1, 2, 3 95 | 96 | expr << operatorPrecedence(expr_term, 97 | [ 98 | ('->', BINARY, opAssoc.LEFT, replace_json_get), 99 | (oneOf('- + ~') | NOT, UNARY, opAssoc.LEFT), 100 | (ISNULL | NOTNULL | (NOT + NULL), UNARY, opAssoc.LEFT), 101 | (IS + NOT, BINARY, opAssoc.LEFT), 102 | ('||', BINARY, opAssoc.LEFT), 103 | (oneOf('* / %'), BINARY, opAssoc.LEFT), 104 | (oneOf('+ -'), BINARY, opAssoc.LEFT), 105 | (oneOf('<< >> & |'), BINARY, opAssoc.LEFT), 106 | (oneOf('< <= > >='), BINARY, opAssoc.LEFT), 107 | ( 108 | oneOf('= == != <>') | IS | IN | LIKE | GLOB | MATCH | REGEXP, 109 | BINARY, 110 | opAssoc.LEFT), 111 | (AND, BINARY, opAssoc.LEFT), 112 | (OR, BINARY, opAssoc.LEFT), 113 | ((BETWEEN, AND), TERNARY, opAssoc.LEFT), 114 | ]) 115 | 116 | compound_operator = (UNION + Optional(ALL) | INTERSECT | EXCEPT) 117 | 118 | ordering_term = expr + Optional(COLLATE + collation_name) + Optional(ASC | DESC) 119 | 120 | join_constraint = Optional( 121 | ON + expr | USING + LPAR + Group(no_suppress_delimited_list(column_name)) + RPAR) 122 | 123 | join_op = COMMA | (Optional(NATURAL) + Optional(INNER | CROSS | LEFT + OUTER | LEFT | OUTER) + JOIN) 124 | 125 | table_reference = ( 126 | (database_name("database") + "." + table_name("table") | table_name("table")) + 127 | Optional(Optional(AS) + table_alias("alias")) 128 | ).setResultsName("table_ids", listAllMatches=True) 129 | 130 | join_source = Forward() 131 | single_source = ( 132 | table_reference + 133 | Optional(INDEXED + BY + index_name | NOT + INDEXED) | 134 | (LPAR + select_stmt + RPAR + Optional(Optional(AS) + table_alias)) | 135 | (LPAR + join_source + RPAR)) 136 | 137 | join_source << single_source + ZeroOrMore(join_op + single_source + join_constraint) 138 | 139 | result_column = table_name + "." + "*" | (expr + Optional(Optional(AS) + column_alias)) | "*" 140 | select_core = ( 141 | SELECT + Optional(DISTINCT | ALL) + Group(no_suppress_delimited_list(result_column)) + 142 | Optional(FROM + join_source) + 143 | Optional(WHERE + expr) + 144 | Optional(GROUP + BY + Group(no_suppress_delimited_list(ordering_term)) + 145 | Optional(HAVING + expr))) 146 | 147 | select_stmt << (select_core + ZeroOrMore(compound_operator + select_core) + 148 | Optional(ORDER + BY + Group(no_suppress_delimited_list(ordering_term))) + 149 | Optional( 150 | LIMIT + (integer | integer + OFFSET + integer | integer + COMMA + integer))) 151 | -------------------------------------------------------------------------------- /aq/sqlite_util.py: -------------------------------------------------------------------------------- 1 | import json 2 | import sqlite3 3 | from datetime import datetime 4 | 5 | from six import string_types 6 | 7 | 8 | def connect(path): 9 | sqlite3.register_adapter(dict, jsonify) 10 | sqlite3.register_adapter(list, jsonify) 11 | db = sqlite3.connect(path) 12 | db.create_function('json_get', 2, json_get) 13 | return db 14 | 15 | 16 | def jsonify(obj): 17 | return json.dumps(obj, default=json_serialize) 18 | 19 | 20 | def json_serialize(obj): 21 | """ 22 | Simple generic JSON serializer for common objects. 23 | """ 24 | if isinstance(obj, datetime): 25 | return obj.isoformat() 26 | 27 | if hasattr(obj, 'id'): 28 | return jsonify(obj.id) 29 | 30 | if hasattr(obj, 'name'): 31 | return jsonify(obj.name) 32 | 33 | raise TypeError('{0} is not JSON serializable'.format(obj)) 34 | 35 | 36 | def json_get(serialized_object, field): 37 | """ 38 | This emulates the HSTORE `->` get value operation. 39 | It get value from JSON serialized column by given key and return `null` if not present. 40 | Key can be either an integer for array index access or a string for object field access. 41 | 42 | :return: JSON serialized value of key in object 43 | """ 44 | # return null if serialized_object is null or "serialized null" 45 | if serialized_object is None: 46 | return None 47 | obj = json.loads(serialized_object) 48 | if obj is None: 49 | return None 50 | 51 | if isinstance(field, int): 52 | # array index access 53 | res = obj[field] if 0 <= field < len(obj) else None 54 | else: 55 | # object field access 56 | res = obj.get(field) 57 | 58 | if not isinstance(res, (int, float, string_types)): 59 | res = json.dumps(res) 60 | 61 | return res 62 | 63 | 64 | def create_table(db, schema_name, table_name, columns): 65 | """ 66 | Create a table, schema_name.table_name, in given database with given list of column names. 67 | """ 68 | table = '{0}.{1}'.format(schema_name, table_name) if schema_name else table_name 69 | db.execute('DROP TABLE IF EXISTS {0}'.format(table)) 70 | columns_list = ', '.join(columns) 71 | db.execute('CREATE TABLE {0} ({1})'.format(table, columns_list)) 72 | 73 | 74 | def insert_all(db, schema_name, table_name, columns, items): 75 | """ 76 | Insert all item in given items list into the specified table, schema_name.table_name. 77 | """ 78 | table = '{0}.{1}'.format(schema_name, table_name) if schema_name else table_name 79 | columns_list = ', '.join(columns) 80 | values_list = ', '.join(['?'] * len(columns)) 81 | query = 'INSERT INTO {table} ({columns}) VALUES ({values})'.format( 82 | table=table, columns=columns_list, values=values_list) 83 | for item in items: 84 | values = [getattr(item, col) for col in columns] 85 | db.execute(query, values) 86 | -------------------------------------------------------------------------------- /aq/util.py: -------------------------------------------------------------------------------- 1 | import os 2 | 3 | from aq.errors import AQError 4 | 5 | 6 | def ensure_data_dir_exists(): 7 | data_dir = os.path.expanduser('~/.aq') 8 | if not os.path.exists(data_dir): 9 | try: 10 | os.mkdir(data_dir) 11 | except OSError as e: 12 | raise AQError('Cannot create data dir at "{0}" because of: {1}.' 13 | 'aq need a working dir to store the temporary tables before querying.' 14 | ''.format(data_dir, e)) 15 | 16 | 17 | -------------------------------------------------------------------------------- /setup.cfg: -------------------------------------------------------------------------------- 1 | [bdist_wheel] 2 | universal=1 -------------------------------------------------------------------------------- /setup.py: -------------------------------------------------------------------------------- 1 | # Always prefer setuptools over distutils 2 | from setuptools import setup, find_packages 3 | # To use a consistent encoding 4 | from codecs import open 5 | from os import path 6 | 7 | here = path.abspath(path.dirname(__file__)) 8 | 9 | # Get the long description from the README file 10 | with open(path.join(here, 'README.rst'), encoding='utf-8') as f: 11 | long_description = f.read() 12 | 13 | setup( 14 | name='aq', 15 | 16 | # Versions should comply with PEP440. For a discussion on single-sourcing 17 | # the version across setup.py and the project code, see 18 | # https://packaging.python.org/en/latest/single_source_version.html 19 | version='0.1.1', 20 | 21 | description='Query AWS resources with SQL', 22 | long_description=long_description, 23 | 24 | # The project's main homepage. 25 | url='https://github.com/lebinh/aq', 26 | 27 | # Author details 28 | author='Binh Le', 29 | author_email='lebinh.it@gmail.com', 30 | 31 | # Choose your license 32 | license='MIT', 33 | 34 | # See https://pypi.python.org/pypi?%3Aaction=list_classifiers 35 | classifiers=[ 36 | # How mature is this project? Common values are 37 | # 3 - Alpha 38 | # 4 - Beta 39 | # 5 - Production/Stable 40 | 'Development Status :: 3 - Alpha', 41 | 42 | # Indicate who your project is intended for 43 | 'Intended Audience :: Developers', 44 | 45 | # Pick your license as you wish (should match "license" above) 46 | 'License :: OSI Approved :: MIT License', 47 | 48 | # Specify the Python versions you support here. In particular, ensure 49 | # that you indicate whether you support Python 2, Python 3 or both. 50 | 'Programming Language :: Python :: 2', 51 | 'Programming Language :: Python :: 2.6', 52 | 'Programming Language :: Python :: 2.7', 53 | 'Programming Language :: Python :: 3', 54 | 'Programming Language :: Python :: 3.3', 55 | 'Programming Language :: Python :: 3.4', 56 | 'Programming Language :: Python :: 3.5', 57 | ], 58 | 59 | # What does your project relate to? 60 | keywords='aws sql query', 61 | 62 | # You can just specify the packages manually here if your project is 63 | # simple. Or you can use find_packages(). 64 | packages=find_packages(exclude=['contrib', 'docs', 'tests']), 65 | 66 | # List run-time dependencies here. These will be installed by pip when 67 | # your project is installed. For an analysis of "install_requires" vs pip's 68 | # requirements files see: 69 | # https://packaging.python.org/en/latest/requirements.html 70 | install_requires=[ 71 | 'boto3', 72 | 'docopt', 73 | 'pyparsing', 74 | 'tabulate', 75 | 'prompt_toolkit', 76 | 'pygments', 77 | 'six', 78 | ], 79 | 80 | # List additional groups of dependencies here (e.g. development 81 | # dependencies). You can install these using the following syntax, 82 | # for example: 83 | # $ pip install -e .[dev,test] 84 | extras_require={ 85 | 'dev': [], 86 | 'test': [], 87 | }, 88 | 89 | entry_points={ 90 | 'console_scripts': [ 91 | 'aq = aq:main' 92 | ] 93 | }, 94 | ) 95 | -------------------------------------------------------------------------------- /tests/test_boto_engine.py: -------------------------------------------------------------------------------- 1 | from unittest import TestCase 2 | 3 | import boto3 4 | from botocore.exceptions import NoRegionError 5 | 6 | from aq import BotoSqliteEngine 7 | from aq.engines import get_resource_model_attributes 8 | 9 | 10 | class TestBotoEngine(TestCase): 11 | engine = BotoSqliteEngine({}) 12 | 13 | def test_is_attached_region(self): 14 | # main is always attached 15 | assert self.engine.is_attached_region('main') 16 | assert not self.engine.is_attached_region('foobar') 17 | 18 | def test_attach_region(self): 19 | assert not self.engine.is_attached_region('us_west_1') 20 | self.engine.attach_region('us_west_1') 21 | assert self.engine.is_attached_region('us_west_1') 22 | self.engine.db.execute('DETACH DATABASE us_west_1') 23 | 24 | def test_get_resource_model_attributes(self): 25 | try: 26 | resource = boto3.resource('ec2') 27 | except NoRegionError: 28 | # skip for environment that doesn't have boto config like CI 29 | pass 30 | else: 31 | collection = resource.instances.all() 32 | attributes = get_resource_model_attributes(resource, collection) 33 | assert attributes 34 | assert 'instance_id' in attributes 35 | assert 'image_id' in attributes 36 | -------------------------------------------------------------------------------- /tests/test_command_line_arg.py: -------------------------------------------------------------------------------- 1 | from unittest import TestCase 2 | 3 | import boto3 4 | from botocore.exceptions import NoRegionError 5 | 6 | from aq import BotoSqliteEngine 7 | from aq.engines import get_resource_model_attributes 8 | 9 | import os, tempfile 10 | from nose.tools import eq_ 11 | 12 | class TestCommandLineArg(TestCase): 13 | 14 | def setUp(self): 15 | try: 16 | del os.environ['AWS_PROFILE'] 17 | del os.environ['AWS_DEFAULT_REGION'] 18 | del os.environ['AWS_CONFIG_FILE'] 19 | del os.environ['AWS_SHARED_CREDENTIALS_FILE'] 20 | except: 21 | pass 22 | 23 | self.credential_file = tempfile.NamedTemporaryFile() 24 | os.environ['AWS_SHARED_CREDENTIALS_FILE'] = self.credential_file.name 25 | self.credential_file.write( 26 | b'[profile_env]\n' 27 | b'region=region-profile-env\n' 28 | b'\n' 29 | b'[profile_arg]\n' 30 | b'region=region-profile-arg\n' 31 | ) 32 | self.credential_file.flush() 33 | 34 | self.config_file = tempfile.NamedTemporaryFile() 35 | os.environ['AWS_CONFIG_FILE'] = self.config_file.name 36 | self.config_file.write( 37 | b'[default]\n' 38 | b'region=region-config-default\n' 39 | b'\n' 40 | b'[config_env]\n' 41 | b'region=region-config-env\n' 42 | ) 43 | self.config_file.flush() 44 | 45 | def test_command_line_arg_profile(self): 46 | os.environ['AWS_PROFILE'] = 'profile_env' 47 | os.environ['AWS_CONFIG_FILE'] = 'config_env' 48 | os.environ['AWS_DEFAULT_REGION'] = 'region-env' 49 | engine = BotoSqliteEngine({ '--profile': 'profile_arg' }) 50 | 51 | eq_(engine.boto3_session.profile_name, 'profile_arg') 52 | 53 | def test_command_line_arg_region(self): 54 | os.environ['AWS_PROFILE'] = 'profile_env' 55 | os.environ['AWS_CONFIG_FILE'] = 'config_env' 56 | os.environ['AWS_DEFAULT_REGION'] = 'region-env' 57 | engine = BotoSqliteEngine({ '--region': 'region-arg' }) 58 | 59 | eq_(engine.boto3_session.region_name, 'region-arg') 60 | 61 | def test_command_line_arg_none(self): 62 | os.environ['AWS_PROFILE'] = 'profile_env' 63 | os.environ['AWS_CONFIG_FILE'] = 'config_env' 64 | os.environ['AWS_DEFAULT_REGION'] = 'region-env' 65 | engine = BotoSqliteEngine({}) 66 | 67 | eq_(engine.boto3_session.profile_name, 'profile_env') 68 | eq_(engine.boto3_session.region_name, 'region-env') 69 | 70 | def test_command_line_arg_and_env_file_none(self): 71 | del os.environ['AWS_CONFIG_FILE'] 72 | del os.environ['AWS_SHARED_CREDENTIALS_FILE'] 73 | 74 | engine = BotoSqliteEngine({}) 75 | 76 | eq_(engine.boto3_session.profile_name, 'default') 77 | eq_(engine.boto3_session.region_name, 'us-east-1') 78 | -------------------------------------------------------------------------------- /tests/test_parsers.py: -------------------------------------------------------------------------------- 1 | from unittest import TestCase 2 | 3 | from aq.errors import QueryParsingError 4 | from aq.parsers import SelectParser, TableId 5 | 6 | 7 | class TestSelectParser(TestCase): 8 | parser = SelectParser({}) 9 | 10 | def test_parse_query_simplest(self): 11 | query, meta = self.parser.parse_query('select * from foo') 12 | self.assertEqual(query, 'SELECT * FROM foo') 13 | table = TableId(None, 'foo', None) 14 | self.assertEqual(meta.tables, [table]) 15 | 16 | def test_parse_query_db(self): 17 | query, meta = self.parser.parse_query('select * from foo.bar') 18 | self.assertEqual(query, 'SELECT * FROM foo . bar') 19 | table = TableId('foo', 'bar', None) 20 | self.assertEqual(meta.tables, [table]) 21 | 22 | def test_parse_query_table_alias(self): 23 | query, meta = self.parser.parse_query('select * from foo.bar a') 24 | self.assertEqual(query, 'SELECT * FROM foo . bar a') 25 | table = TableId('foo', 'bar', 'a') 26 | self.assertEqual(meta.tables, [table]) 27 | 28 | def test_parse_query_table_as_alias(self): 29 | query, meta = self.parser.parse_query('select * from foo.bar as a') 30 | self.assertEqual(query, 'SELECT * FROM foo . bar AS a') 31 | table = TableId('foo', 'bar', 'a') 32 | self.assertEqual(meta.tables, [table]) 33 | 34 | def test_parse_query_join_simple(self): 35 | query, meta = self.parser.parse_query('select * from foo, bar') 36 | self.assertEqual(query, 'SELECT * FROM foo , bar') 37 | table1 = TableId(None, 'foo', None) 38 | table2 = TableId(None, 'bar', None) 39 | self.assertEqual(meta.tables, [table1, table2]) 40 | 41 | def test_parse_query_join_expr(self): 42 | query, meta = self.parser.parse_query('select * from foo join bar on foo.a = bar.b') 43 | self.assertEqual(query, 'SELECT * FROM foo JOIN bar ON foo.a = bar.b') 44 | table1 = TableId(None, 'foo', None) 45 | table2 = TableId(None, 'bar', None) 46 | self.assertEqual(meta.tables, [table1, table2]) 47 | 48 | def test_parse_query_join_table_with_using(self): 49 | query, meta = self.parser.parse_query('select * from foo join foo.bar using (name)') 50 | self.assertEqual(query, 'SELECT * FROM foo JOIN foo . bar USING ( name )') 51 | table1 = TableId(None, 'foo', None) 52 | table2 = TableId('foo', 'bar', None) 53 | self.assertEqual(meta.tables, [table1, table2]) 54 | 55 | def test_parse_query_sub_select(self): 56 | query, meta = self.parser.parse_query('select * from (select * from foo)') 57 | self.assertEqual(query, 'SELECT * FROM ( SELECT * FROM foo )') 58 | table = TableId(None, 'foo', None) 59 | self.assertEqual(meta.tables, [table]) 60 | 61 | def test_parse_query_sub_select_and_join(self): 62 | query, meta = self.parser.parse_query('select * from (select * from foo.bar) left join blah') 63 | self.assertEqual(query, 'SELECT * FROM ( SELECT * FROM foo . bar ) LEFT JOIN blah') 64 | table1 = TableId('foo', 'bar', None) 65 | table2 = TableId(None, 'blah', None) 66 | self.assertEqual(meta.tables, [table1, table2]) 67 | 68 | def test_parse_query_invalid(self): 69 | try: 70 | self.parser.parse_query('foo') 71 | except QueryParsingError: 72 | pass 73 | else: 74 | self.fail() 75 | 76 | def test_parse_query_expand_json_get(self): 77 | query, _ = self.parser.parse_query("select foo->1") 78 | self.assertEqual(query, 'SELECT json_get(foo, 1)') 79 | 80 | query, _ = self.parser.parse_query("select foo.bar -> 'blah'") 81 | self.assertEqual(query, "SELECT json_get(foo.bar, 'blah')") 82 | 83 | query, _ = self.parser.parse_query("select foo->bar->blah") 84 | self.assertEqual(query, 'SELECT json_get(json_get(foo, bar), blah)') 85 | 86 | def test_parse_query_expand_not_json_get(self): 87 | query, _ = self.parser.parse_query("select * from foo where x = 'bar -> 1'") 88 | self.assertEqual(query, "SELECT * FROM foo WHERE x = 'bar -> 1'") 89 | 90 | query, _ = self.parser.parse_query("select * from foo where x -> 'bar -> 1'") 91 | self.assertEqual(query, "SELECT * FROM foo WHERE json_get(x, 'bar -> 1')") 92 | 93 | def test_parse_query_with_and(self): 94 | query, _ = self.parser.parse_query("select * from foo where x = 'foo' and y = 'bar'") 95 | self.assertEqual(query, "SELECT * FROM foo WHERE x = 'foo' AND y = 'bar'") 96 | 97 | def test_parse_query_with_or(self): 98 | query, _ = self.parser.parse_query("select * from foo where x = 'foo' or y = 'bar'") 99 | self.assertEqual(query, "SELECT * FROM foo WHERE x = 'foo' OR y = 'bar'") 100 | -------------------------------------------------------------------------------- /tests/test_sqlite_util.py: -------------------------------------------------------------------------------- 1 | from unittest import TestCase 2 | 3 | from aq.sqlite_util import connect, create_table, insert_all 4 | 5 | 6 | class TestSqliteUtil(TestCase): 7 | def test_dict_adapter(self): 8 | with connect(':memory:') as conn: 9 | conn.execute('CREATE TABLE foo (foo)') 10 | conn.execute('INSERT INTO foo (foo) VALUES (?)', ({'bar': 'blah'},)) 11 | values = conn.execute('SELECT * FROM foo').fetchone() 12 | self.assertEqual(len(values), 1) 13 | self.assertEqual(values[0], '{"bar": "blah"}') 14 | 15 | def test_create_table(self): 16 | with connect(':memory:') as conn: 17 | create_table(conn, None, 'foo', ('col1', 'col2')) 18 | tables = conn.execute("PRAGMA table_info(\'foo\')").fetchall() 19 | self.assertEqual(len(tables), 2) 20 | self.assertEqual(tables[0][1], 'col1') 21 | self.assertEqual(tables[1][1], 'col2') 22 | 23 | def test_insert_all(self): 24 | class Foo(object): 25 | def __init__(self, c1, c2): 26 | self.c1 = c1 27 | self.c2 = c2 28 | 29 | columns = ('c1', 'c2') 30 | values = (Foo(1, 2), Foo(3, 4)) 31 | with connect(':memory:') as conn: 32 | create_table(conn, None, 'foo', columns) 33 | insert_all(conn, None, 'foo', columns, values) 34 | rows = conn.execute('SELECT * FROM foo').fetchall() 35 | self.assertTrue((1, 2) in rows, '(1, 2) in rows') 36 | self.assertTrue((3, 4) in rows, '(3, 4) in rows') 37 | 38 | def test_json_get_field(self): 39 | with connect(':memory:') as conn: 40 | json_obj = '{"foo": "bar"}' 41 | query = "select json_get('{0}', 'foo')".format(json_obj) 42 | self.assertEqual(conn.execute(query).fetchone()[0], 'bar') 43 | 44 | def test_json_get_index(self): 45 | with connect(':memory:') as conn: 46 | json_obj = '[1, 2, 3]' 47 | query = "select json_get('{0}', 1)".format(json_obj) 48 | self.assertEqual(conn.execute(query).fetchone()[0], 2) 49 | 50 | def test_json_get_field_nested(self): 51 | with connect(':memory:') as conn: 52 | json_obj = '{"foo": {"bar": "blah"}}' 53 | query = "select json_get('{0}', 'foo')".format(json_obj) 54 | self.assertEqual(conn.execute(query).fetchone()[0], '{"bar": "blah"}') 55 | query = "select json_get(json_get('{0}', 'foo'), 'bar')".format(json_obj) 56 | self.assertEqual(conn.execute(query).fetchone()[0], 'blah') 57 | 58 | def test_json_get_field_of_null(self): 59 | with connect(':memory:') as conn: 60 | query = "select json_get(NULL, 'foo')" 61 | self.assertEqual(conn.execute(query).fetchone()[0], None) 62 | 63 | def test_json_get_field_of_serialized_null(self): 64 | with connect(':memory:') as conn: 65 | json_obj = 'null' 66 | query = "select json_get('{0}', 'foo')".format(json_obj) 67 | self.assertEqual(conn.execute(query).fetchone()[0], None) 68 | --------------------------------------------------------------------------------