From 51e2fe48dbf39de0f2bec76364e9c6b04ffeece1 Mon Sep 17 00:00:00 2001
From: Alexandr <ab@sun.hexor.ru>
Date: Mon, 26 Feb 2018 19:13:04 +0100
Subject: [PATCH] Add sql query feature.

---
 database.py |  3 +++
 worker.py   | 12 ++++++++++++
 2 files changed, 15 insertions(+)

diff --git a/database.py b/database.py
index df11487..ff99253 100755
--- a/database.py
+++ b/database.py
@@ -3,6 +3,7 @@ import datetime as dt
 class DataBase:
     def __init__(self, basefile, scheme):
         import sqlite3
+        self.scheme = ''
         try:
             self.conn = sqlite3.connect(basefile, check_same_thread=False)
         except:
@@ -10,6 +11,7 @@ class DataBase:
             return None
         with open(scheme, 'r') as scheme_sql:
             sql = scheme_sql.read()
+            self.scheme = sql
             if self.conn is not None:
                 try:
                     cursor = self.conn.cursor()
@@ -129,6 +131,7 @@ class DataBase:
         if 'DELETE' in sql.upper() \
         or 'INSERT' in sql.upper() \
         or 'UPDATE' in sql.upper() \
+        or 'DROP' in sql.upper() \
         or 'CREATE' in sql.upper() \
         or 'ALTER' in sql.upper():
             return('gtfo')
diff --git a/worker.py b/worker.py
index 239eeb5..6156a74 100755
--- a/worker.py
+++ b/worker.py
@@ -17,6 +17,14 @@ class MessageWorker:
 
     def handleUpdate(self, msg):
         try:
+            if msg['message']['text'] == '/scheme':
+                conf_id = msg['message']['chat']['id']
+                user_id = msg['message']['from']['id']
+                chat_title = msg['message']['chat']['title']
+                self.db.add_conf(conf_id, chat_title)
+                self.send(id=conf_id, msg='```\n' + self.db.scheme + '\n```')
+                return True
+
             if msg['message']['text'] == '/stat':
                 conf_id = msg['message']['chat']['id']
                 user_id = msg['message']['from']['id']
@@ -53,6 +61,10 @@ class MessageWorker:
                 sql = msg['message']['text'][5:]
 
                 res = self.db.command(sql)
+                if 'syntax' in str(res) \
+                or 'ambiguous' in str(res):
+                  self.send(id=conf_id, msg=str(res))
+                  return False
                 try:
                   msg = '```\n'
                   for z in res: