sqlnotion

#!/usr/bin/env python3
import argparse
import datetime
import os
import re
import sys

import requests

def main():
    args = get_args()
    executor = get_executor(args)
    reporter = Reporter(args.api_token, args.notion_host)

    query_file_paths = get_queries(args.queries_folder)
    for query_file_path in query_file_paths:
        with open(query_file_path) as query_file:
            ingredient_key = query_name(query_file_path, executor.extension)
            value = executor.execute_query(query_file.read())

            print(f'{ingredient_key} : {value}')

            reporter.report(ingredient_key, value)

def query_name(query_file_path, extension):
    basename = os.path.basename(query_file_path)
    return re.sub(extension + '$', '', basename)

def get_executor(args):
    return db_types[args.dbtype](args)

def get_queries(folder_path):
    for file_path in os.listdir(folder_path):
        if file_path.endswith('.sql'):
            yield os.path.join(folder_path, file_path)

def get_args():
    parser = argparse.ArgumentParser()

    optional = parser
    required = parser.add_argument_group('required arguments')

    required.add_argument('--dbtype', choices=db_types.keys(), default=None, required=True,
        help='the type of database ')

    required.add_argument('--dbname', type=str, default=None, required=True,
        help='the name of the database to connect to')

    required.add_argument('--user', type=str, default=None, required=True,
        help='the database username to connect with')

    optional.add_argument('--password', type=str, default=None, required=False,
        help="the database username's password")

    optional.add_argument('--host', type=str, default=None, required=False,
        help='the hostname the database is located at')

    optional.add_argument('--port', type=str, default=None, required=False,
        help='the port to connect to the database with')

    required.add_argument('--queries-folder', type=str, default=None, required=True,
        help='the path to the folder of queries')

    required.add_argument('--api-token', type=str, default=None, required=True,
        help='the api token of a notion account https://app.usenotion.com/api_control_panel')

    parser.add_argument('--notion-host', type=str, default='https://analyze.jamacloud.com/',
        help=argparse.SUPPRESS)

    try:
        return parser.parse_args()
    except SystemExit as e:
        sys.exit(1)


class ResultFormatError(Exception):
    pass

class PostgresExecutor:

    def __init__(self, args):
        try:
            import psycopg2
        except ImportError:
            print('psycopg2 is required but not installed', file=sys.stderr)
            sys.exit(1)

        self.extension = '.sql'

        self.connection = psycopg2.connect(
            dbname=args.dbname,
            user=args.user,
            password=args.password,
            host=args.host,
            port=args.port,
        )

    def _execute_query(self, query_string):
        cursor = self.connection.cursor()
        cursor.execute(query_string)
        return cursor.fetchall()

    def _get_number_result(self, query_string):
        result = self._execute_query(query_string)

        if len(result) != 1:
            count = len(result)
            raise ResultFormatError(
                f'Expected 1 row to be returned, got {count}: ' + str(result))

        if len(result[0]) != 1:
            count = len(result[0])
            raise ResultFormatError(
                f'Expected 1 column to be returned, got {count}: ' + str(result))

        return result[0][0]

    def execute_query(self, query_string):
        return self._get_number_result(query_string)

class MysqlExecutor:
    def __init__(self, args):
        try:
            import _mysql
        except ImportError:
            print('mysqlclient is required but not installed', file=sys.stderr)
            sys.exit(1)

        self.extension = '.sql'

        self.connection = _mysql.connect(
            args.host,
            args.user,
            args.password,
            args.dbname,
        )

    def _execute_query(self, query_string):
        self.connection.query(query_string)
        result = self.connection.store_result()
        return result.fetch_row()

    def execute_query(self, query_string):
        result = self._execute_query(query_string)

        if len(result) != 1:
            count = len(result)
            raise ResultFormatError(
                f'Expected 1 row to be returned, got {count}: ' + str(result))

        if len(result[0]) != 1:
            count = len(result[0])
            raise ResultFormatError(
                f'Expected 1 column to be returned, got {count}: ' + str(result))

        return result[0][0]


class Reporter:
    def __init__(self, api_token, notion_host):
        self.api_token = api_token
        self.notion_host = notion_host
        self.headers = {
            'Authorization': self.api_token,
            'Content-Type': 'application/json',
            'Accept': 'application/json',
        }

    def report(self, ingredient_key, value):
        url = self.notion_host + '/api/v1/report'

        json_data = {
            'ingredient_key': ingredient_key,
            'value': value,
            'date': datetime.datetime.today().strftime('%Y-%m-%d'),
        }

        response = requests.post(url, json=json_data, headers=self.headers)
        response.raise_for_status()

        return response

db_types = {
    'postgres': PostgresExecutor,
    'mysql': MysqlExecutor,
}

if __name__ == '__main__':
    main()