From a38beca61c16969ed3162c5d5e082077252ce8e6 Mon Sep 17 00:00:00 2001 From: Chris Clark Date: Mon, 13 May 2024 20:41:51 -0400 Subject: [PATCH] Userland connections --- .gitignore | 2 + LICENSE | 11 +- docs/features.rst | 38 +-- docs/settings.rst | 11 +- explorer/admin.py | 3 +- explorer/app_settings.py | 18 +- explorer/apps.py | 5 +- explorer/assistant/tests.py | 108 ------- explorer/assistant/utils.py | 71 ++++- explorer/assistant/views.py | 66 ++-- explorer/connections.py | 54 +++- explorer/ee/LICENSE | 21 ++ explorer/ee/__init__.py | 1 + explorer/ee/db_connections/__init__.py | 0 explorer/ee/db_connections/admin.py | 9 + explorer/ee/db_connections/forms.py | 17 ++ explorer/ee/db_connections/models.py | 63 ++++ explorer/ee/db_connections/utils.py | 170 +++++++++++ explorer/ee/db_connections/views.py | 127 ++++++++ explorer/ee/urls.py | 21 ++ explorer/forms.py | 5 +- explorer/locale/ru/LC_MESSAGES/django.po | 2 +- .../migrations/0017_databaseconnection.py | 26 ++ ..._alter_databaseconnection_host_and_more.py | 29 ++ explorer/models.py | 9 +- explorer/permissions.py | 4 + explorer/src/js/assistant.js | 32 +- explorer/src/js/codemirror-config.js | 3 +- explorer/src/js/explorer.js | 27 +- explorer/src/js/main.js | 4 +- explorer/src/js/query-list.js | 2 +- explorer/src/js/schemaService.js | 8 +- explorer/src/js/uploads.js | 83 +++++ explorer/tasks.py | 15 + explorer/telemetry.py | 5 +- .../templates/connections/connections.html | 43 +++ .../database_connection_confirm_delete.html | 12 + .../database_connection_detail.html | 39 +++ .../connections/database_connection_form.html | 23 ++ explorer/templates/explorer/assistant.html | 2 + explorer/templates/explorer/base.html | 12 +- explorer/templates/explorer/query.html | 2 +- explorer/templates/explorer/query_list.html | 35 ++- explorer/tests/csvs/all_types.csv | 8 + explorer/tests/csvs/dates.csv | 13 + explorer/tests/csvs/floats.csv | 8 + explorer/tests/csvs/integers.csv | 5 + explorer/tests/csvs/mixed.csv | 5 + explorer/tests/csvs/test_case1.csv | 31 ++ explorer/tests/settings.py | 5 +- explorer/tests/test_assistant.py | 283 ++++++++++++++++++ explorer/tests/test_connections.py | 71 +++++ explorer/tests/test_models.py | 18 +- explorer/tests/test_tasks.py | 108 ++++++- explorer/tests/test_telemetry.py | 13 + explorer/tests/test_user_connection_utils.py | 192 ++++++++++++ explorer/tests/test_utils.py | 4 +- explorer/tests/test_views.py | 192 +++++++++++- explorer/urls.py | 10 +- explorer/utils.py | 31 +- explorer/views/__init__.py | 1 - explorer/views/list.py | 4 +- explorer/views/mixins.py | 7 +- explorer/views/schema.py | 4 +- public_key.pem | 9 + requirements/base.txt | 4 +- requirements/extra/assistant.txt | 2 +- requirements/extra/user_connections.txt | 3 + test_project/celery_config.py | 23 +- test_project/settings.py | 17 +- tox.ini | 2 +- 71 files changed, 2028 insertions(+), 283 deletions(-) delete mode 100644 explorer/assistant/tests.py create mode 100644 explorer/ee/LICENSE create mode 100644 explorer/ee/__init__.py create mode 100644 explorer/ee/db_connections/__init__.py create mode 100644 explorer/ee/db_connections/admin.py create mode 100644 explorer/ee/db_connections/forms.py create mode 100644 explorer/ee/db_connections/models.py create mode 100644 explorer/ee/db_connections/utils.py create mode 100644 explorer/ee/db_connections/views.py create mode 100644 explorer/ee/urls.py create mode 100644 explorer/migrations/0017_databaseconnection.py create mode 100644 explorer/migrations/0018_alter_databaseconnection_host_and_more.py create mode 100644 explorer/src/js/uploads.js create mode 100644 explorer/templates/connections/connections.html create mode 100644 explorer/templates/connections/database_connection_confirm_delete.html create mode 100644 explorer/templates/connections/database_connection_detail.html create mode 100644 explorer/templates/connections/database_connection_form.html create mode 100644 explorer/tests/csvs/all_types.csv create mode 100644 explorer/tests/csvs/dates.csv create mode 100644 explorer/tests/csvs/floats.csv create mode 100644 explorer/tests/csvs/integers.csv create mode 100644 explorer/tests/csvs/mixed.csv create mode 100644 explorer/tests/csvs/test_case1.csv create mode 100644 explorer/tests/test_assistant.py create mode 100644 explorer/tests/test_connections.py create mode 100644 explorer/tests/test_user_connection_utils.py create mode 100644 public_key.pem create mode 100644 requirements/extra/user_connections.txt diff --git a/.gitignore b/.gitignore index 45ec76b9..1ba5acb7 100644 --- a/.gitignore +++ b/.gitignore @@ -22,3 +22,5 @@ docs/_build/ .env tst tst2 +user_dbs/* +tmp2 diff --git a/LICENSE b/LICENSE index 4de004d8..dab1e28f 100644 --- a/LICENSE +++ b/LICENSE @@ -1,6 +1,11 @@ -The MIT License (MIT) +* All content that resides under the "explorer/ee/" directory of this repository is licensed under the license defined +in "explorer/ee/LICENSE". -Copyright (c) 2013 Chris Clark, ePantry LLC +* Content outside of the above mentioned directory is provided under the "MIT" license as defined below. + +** The MIT License (MIT) ** + +Copyright (c) 2024, SQL Explorer, Inc Permission is hereby granted, free of charge, to any person obtaining a copy of this software and associated documentation files (the "Software"), to deal @@ -18,4 +23,4 @@ FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN -THE SOFTWARE. \ No newline at end of file +THE SOFTWARE. diff --git a/docs/features.rst b/docs/features.rst index 98ffb7e1..457497b1 100644 --- a/docs/features.rst +++ b/docs/features.rst @@ -3,19 +3,16 @@ Features Security -------- -- Let's not kid ourselves - this tool is all about giving people - access to running SQL in production. So if that makes you - nervous (**and it should**) - you've been warned. Explorer makes an - effort to not allow terrible things to happen, but be careful! - It's recommended you setup read-only roles for each of your database +- It's recommended you setup read-only roles for each of your database connections and only use these particular connections for your queries through the - ``EXPLORER_CONNECTIONS`` setting. -- Explorer supports two different permission checks for users of + ``EXPLORER_CONNECTIONS`` setting -- or set up userland connections via DatabaseConnections in + the Django admin, or the SQL Explorer front-end. +- SQL Explorer supports two different permission checks for users of the tool. Users passing the ``EXPLORER_PERMISSION_CHANGE`` test can create, edit, delete, and execute queries. Users who do not pass this test but pass the ``EXPLORER_PERMISSION_VIEW`` test can only execute queries. Other users cannot access any part of - Explorer. Both permission groups are set to is_staff by default + SQL Explorer. Both permission groups are set to is_staff by default and can be overridden in your settings file. - Enforces a SQL blacklist so destructive queries don't get executed (delete, drop, alter, update etc). This is not @@ -37,19 +34,25 @@ SQL Assistant to quickly get help with your query, with relevant schema automatically injected into the prompt. +Configurable Connections +------------------------ +- Configure connections via the settings.py file, or via the SQL Explorer UI. +- Supports drag-and-drop uploading of CSV files or SQLite databases for instant SQL access to your data. + Snapshots --------- - Tick the 'snapshot' box on a query, and Explorer will upload a .csv snapshot of the query results to S3. Configure the snapshot - frequency via a celery cron task, e.g. for daily at 1am: + frequency via a celery cron task, e.g. for daily at 1am + (see test_project/celery_config.py for an example of this, along with test_project/__init__.py): .. code-block:: python app.conf.beat_schedule = { - 'explorer.tasks.snapshot_queries': { - 'task': 'explorer.tasks.snapshot_queries', - 'schedule': crontab(hour=1, minute=0) - } + "explorer.tasks.snapshot_queries": { + "task": "explorer.tasks.snapshot_queries", + "schedule": crontab(hour="1", minute="0") + }, } - Requires celery, obviously. Also uses boto3. All @@ -64,6 +67,7 @@ Email query results - Click the email icon in the query listing view, enter an email address, and the query results (zipped .csv) will be sent to you asynchronously. Very handy for long-running queries. +- You must also have the setting ``EXPLORER_TASKS_ENABLED`` enabled. Parameterized Queries --------------------- @@ -168,10 +172,10 @@ Query Logs .. code-block:: python app.conf.beat_schedule = { - 'explorer.tasks.truncate_querylogs': { - 'task': 'explorer.tasks.truncate_querylogs', - 'schedule': crontab(hour=1, minute=0), - 'kwargs': {'days': 30} + "explorer.tasks.truncate_querylogs": { + "task": "explorer.tasks.truncate_querylogs", + "schedule": crontab(hour="1", minute="10"), + "kwargs": {"days": 30} } } diff --git a/docs/settings.rst b/docs/settings.rst index fa164e2a..653b002c 100644 --- a/docs/settings.rst +++ b/docs/settings.rst @@ -357,4 +357,13 @@ The following three settings control the SQL Assistant. More information is avai EXPLORER_AI_API_KEY = getattr(settings, "EXPLORER_AI_API_KEY", None) EXPLORER_ASSISTANT_BASE_URL = getattr(settings, "EXPLORER_ASSISTANT_BASE_URL", "https://api.openai.com/v1") - EXPLORER_ASSISTANT_MODEL + EXPLORER_ASSISTANT_MODEL = getattr(settings, "EXPLORER_ASSISTANT_MODEL", + # Return the model name and max_tokens it supports + {"name": "gpt-4o", + "max_tokens": 128000}) + + +User-Configured DB Connections +****************************** +Set `EXPLORER_DB_CONNECTIONS_ENABLED` to `True` to enable DB connections to get configured in the browser (e.g. not +just in settings.py). This also allows uploading of CSV or SQLite files for instant querying. diff --git a/explorer/admin.py b/explorer/admin.py index aee8a3c7..0cad25ac 100644 --- a/explorer/admin.py +++ b/explorer/admin.py @@ -2,6 +2,7 @@ from explorer.actions import generate_report_action from explorer.models import Query, ExplorerValue +from explorer.ee.db_connections.admin import DatabaseConnectionAdmin # noqa @admin.register(Query) @@ -16,7 +17,6 @@ class QueryAdmin(admin.ModelAdmin): class ExplorerValueAdmin(admin.ModelAdmin): list_display = ("key", "value", "display_key") list_filter = ("key",) - readonly_fields = ("key",) search_fields = ("key", "value") def display_key(self, obj): @@ -24,4 +24,3 @@ def display_key(self, obj): return dict(ExplorerValue.EXPLORER_SETTINGS_CHOICES).get(obj.key, "") display_key.short_description = "Setting Name" - diff --git a/explorer/app_settings.py b/explorer/app_settings.py index 5f52b97a..72156cf5 100644 --- a/explorer/app_settings.py +++ b/explorer/app_settings.py @@ -69,8 +69,11 @@ EXPLORER_PERMISSION_CHANGE = getattr( settings, "EXPLORER_PERMISSION_CHANGE", lambda r: r.user.is_staff ) +EXPLORER_PERMISSION_CONNECTIONS = getattr( + settings, "EXPLORER_PERMISSION_CONNECTIONS", lambda r: r.user.is_staff +) EXPLORER_RECENT_QUERY_COUNT = getattr( - settings, "EXPLORER_RECENT_QUERY_COUNT", 10 + settings, "EXPLORER_RECENT_QUERY_COUNT", 5 ) EXPLORER_ASYNC_SCHEMA = getattr(settings, "EXPLORER_ASYNC_SCHEMA", False) @@ -125,7 +128,7 @@ S3_REGION = getattr(settings, "EXPLORER_S3_REGION", "us-east-1") S3_ENDPOINT_URL = getattr(settings, "EXPLORER_S3_ENDPOINT_URL", None) S3_DESTINATION = getattr(settings, "EXPLORER_S3_DESTINATION", "") -S3_SIGNATURE_VERSION = getattr(settings, "EXPLORER_S3_SIGNATURE_VERSION", "v2") +S3_SIGNATURE_VERSION = getattr(settings, "EXPLORER_S3_SIGNATURE_VERSION", "v4") UNSAFE_RENDERING = getattr(settings, "EXPLORER_UNSAFE_RENDERING", False) @@ -146,8 +149,17 @@ # AI Assistant settings. Setting the first to an OpenAI key is the simplest way to enable the assistant EXPLORER_AI_API_KEY = getattr(settings, "EXPLORER_AI_API_KEY", None) + EXPLORER_ASSISTANT_BASE_URL = getattr(settings, "EXPLORER_ASSISTANT_BASE_URL", "https://api.openai.com/v1") EXPLORER_ASSISTANT_MODEL = getattr(settings, "EXPLORER_ASSISTANT_MODEL", # Return the model name and max_tokens it supports - {"name": "gpt-4-0125-preview", + {"name": "gpt-4o", "max_tokens": 128000}) + +EXPLORER_DB_CONNECTIONS_ENABLED = getattr(settings, "EXPLORER_DB_CONNECTIONS_ENABLED", False) +EXPLORER_PRUNE_LOCAL_UPLOAD_COPY_DAYS_INACTIVITY = getattr(settings, + "EXPLORER_PRUNE_LOCAL_UPLOAD_COPY_DAYS_INACTIVITY", 7) + + +def has_assistant(): return EXPLORER_AI_API_KEY is not None +def db_connections_enabled(): return EXPLORER_DB_CONNECTIONS_ENABLED diff --git a/explorer/apps.py b/explorer/apps.py index 6797abfe..522d9ada 100644 --- a/explorer/apps.py +++ b/explorer/apps.py @@ -28,8 +28,9 @@ def _get_explorer_connections(): def _validate_connections(): - # Validate connections - if _get_default() not in _get_explorer_connections().values(): + # Validate connections, when using settings.EXPLORER_CONNECTIONS + # Skip if none are configured, as the app will use user-configured connections (DatabaseConnection models) + if _get_explorer_connections().values() and _get_default() not in _get_explorer_connections().values(): raise ImproperlyConfigured( f"EXPLORER_DEFAULT_CONNECTION is {_get_default()}, " f"but that alias is not present in the values of " diff --git a/explorer/assistant/tests.py b/explorer/assistant/tests.py deleted file mode 100644 index a92dacf9..00000000 --- a/explorer/assistant/tests.py +++ /dev/null @@ -1,108 +0,0 @@ -from explorer.tests.factories import SimpleQueryFactory -from unittest.mock import patch, Mock -import unittest - -import json -from django.test import TestCase -from django.urls import reverse -from django.contrib.auth.models import User -from explorer.app_settings import EXPLORER_DEFAULT_CONNECTION as CONN -from explorer import app_settings - - -class TestAssistantViews(TestCase): - - def setUp(self): - self.user = User.objects.create_superuser( - "admin", "admin@admin.com", "pwd" - ) - self.client.login(username="admin", password="pwd") - self.request_data = { - "sql": "SELECT * FROM explorer_query", - "connection": CONN, - "assistant_request": "Test Request" - } - - @unittest.skipIf(not app_settings.EXPLORER_AI_API_KEY, "assistant not enabled") - @patch("explorer.assistant.utils.openai_client") - def test_do_modify_query(self, mocked_openai_client): - from explorer.assistant.views import run_assistant - - # create.return_value should match: resp.choices[0].message - mocked_openai_client.return_value.chat.completions.create.return_value = Mock( - choices=[Mock(message=Mock(content="smart computer"))]) - resp = run_assistant(self.request_data, None) - self.assertEqual(resp, "smart computer") - - @unittest.skipIf(not app_settings.EXPLORER_AI_API_KEY, "assistant not enabled") - def test_assistant_help(self): - resp = self.client.post(reverse("assistant"), - data=json.dumps(self.request_data), - content_type="application/json") - self.assertIsNone(json.loads(resp.content)["message"]) - - -class TestPromptContext(TestCase): - - def test_retrieves_sample_rows(self): - from explorer.assistant.utils import sample_rows_from_table, ROW_SAMPLE_SIZE - SimpleQueryFactory(title="First Query") - SimpleQueryFactory(title="Second Query") - SimpleQueryFactory(title="Third Query") - SimpleQueryFactory(title="Fourth Query") - ret = sample_rows_from_table(CONN, "explorer_query") - self.assertEqual(len(ret), ROW_SAMPLE_SIZE+1) # includes header row - - def test_format_rows_from_table(self): - from explorer.assistant.utils import format_rows_from_table - d = [ - ["col1", "col2"], - ["val1", "val2"], - ] - ret = format_rows_from_table(d) - self.assertEqual(ret, "col1 | col2\n" + "-" * 50 + "\nval1 | val2\n") - - def test_parsing_tables_from_query(self): - from explorer.assistant.utils import get_table_names_from_query - sql = "SELECT * FROM explorer_query" - ret = get_table_names_from_query(sql) - self.assertEqual(ret, ["explorer_query"]) - - def test_parsing_tables_from_query_bad_sql(self): - from explorer.assistant.utils import get_table_names_from_query - sql = "foo" - ret = get_table_names_from_query(sql) - self.assertEqual(ret, []) - - def test_schema_info_from_table_names(self): - from explorer.assistant.utils import tables_from_schema_info - ret = tables_from_schema_info(CONN, ["explorer_query"]) - expected = [("explorer_query", [ - ("id", "AutoField"), - ("title", "CharField"), - ("sql", "TextField"), - ("description", "TextField"), - ("created_at", "DateTimeField"), - ("last_run_date", "DateTimeField"), - ("created_by_user_id", "IntegerField"), - ("snapshot", "BooleanField"), - ("connection", "CharField")])] - self.assertEqual(ret, expected) - - -class TestAssistantUtils(TestCase): - - def test_sample_rows_from_tables(self): - from explorer.assistant.utils import sample_rows_from_tables - SimpleQueryFactory(title="First Query") - SimpleQueryFactory(title="Second Query") - ret = sample_rows_from_tables(CONN, ["explorer_query"]) - self.assertTrue("First Query" in ret) - self.assertTrue("Second Query" in ret) - - def test_sample_rows_from_tables_no_tables(self): - from explorer.assistant.utils import sample_rows_from_tables - SimpleQueryFactory(title="First Query") - SimpleQueryFactory(title="Second Query") - ret = sample_rows_from_tables(CONN, []) - self.assertEqual(ret, "") diff --git a/explorer/assistant/utils.py b/explorer/assistant/utils.py index 7f444219..448d78fe 100644 --- a/explorer/assistant/utils.py +++ b/explorer/assistant/utils.py @@ -1,11 +1,13 @@ from explorer import app_settings from explorer.schema import schema_info +from explorer.models import ExplorerValue from explorer.utils import get_valid_connection from django.db.utils import OperationalError OPENAI_MODEL = app_settings.EXPLORER_ASSISTANT_MODEL["name"] ROW_SAMPLE_SIZE = 2 +MAX_FIELD_SAMPLE_SIZE = 500 # characters def openai_client(): @@ -41,20 +43,46 @@ def tables_from_schema_info(connection, table_names): def sample_rows_from_tables(connection, table_names): ret = "" for table_name in table_names: - ret = f"SAMPLE FROM TABLE {table_name}:\n" - ret = ret + format_rows_from_table( + ret += f"SAMPLE FROM TABLE {table_name}:\n" + ret += format_rows_from_table( sample_rows_from_table(connection, table_name) ) + "\n\n" return ret def sample_rows_from_table(connection, table_name): + """ + Fetches a sample of rows from the specified table and ensures that any field values + exceeding 500 characters (or bytes) are truncated. This is useful for handling fields + like "description" that might contain very long strings of text or binary data. + Truncating these fields prevents issues with displaying or processing overly large values. + An ellipsis ("...") is appended to indicate that the data has been truncated. + + Args: + connection: The database connection. + table_name: The name of the table to sample rows from. + + Returns: + A list of rows with field values truncated if they exceed 500 characters/bytes. + """ conn = get_valid_connection(connection) cursor = conn.cursor() try: cursor.execute(f"SELECT * FROM {table_name} LIMIT {ROW_SAMPLE_SIZE}") ret = [[header[0] for header in cursor.description]] - ret = ret + cursor.fetchall() + rows = cursor.fetchall() + + for row in rows: + processed_row = [] + for field in row: + new_val = field + if isinstance(field, str) and len(field) > MAX_FIELD_SAMPLE_SIZE: + new_val = field[:MAX_FIELD_SAMPLE_SIZE] + "..." # Truncate and add ellipsis + elif isinstance(field, (bytes, bytearray)) and len(field) > MAX_FIELD_SAMPLE_SIZE: + new_val = field[:MAX_FIELD_SAMPLE_SIZE] + b"..." # Truncate binary data + processed_row.append(new_val) + ret.append(processed_row) + return ret except OperationalError as e: return [[str(e)]] @@ -83,7 +111,10 @@ def get_table_names_from_query(sql): def num_tokens_from_string(string: str) -> int: """Returns the number of tokens in a text string.""" import tiktoken - encoding = tiktoken.encoding_for_model(OPENAI_MODEL) + try: + encoding = tiktoken.encoding_for_model(OPENAI_MODEL) + except KeyError: + encoding = tiktoken.get_encoding("cl100k_base") num_tokens = len(encoding.encode(string)) return num_tokens @@ -92,3 +123,35 @@ def fits_in_window(string: str) -> bool: # Ratchet down by 5% to account for other boilerplate and system prompt # TODO make this better by actually looking at the token count of the system prompt return num_tokens_from_string(string) < (app_settings.EXPLORER_ASSISTANT_MODEL["max_tokens"] * 0.95) + + +def build_prompt(request_data, included_tables): + user_prompt = "" + + db_vendor = get_valid_connection(request_data.get("connection")).vendor + user_prompt += f"## Database Vendor / SQL Flavor is {db_vendor}\n\n" + + db_error = request_data.get("db_error") + if db_error: + user_prompt += f"## Query Error ##\n\n{db_error}\n\n" + + sql = request_data.get("sql") + if sql: + user_prompt += f"## Existing SQL ##\n\n{sql}\n\n" + + results_sample = sample_rows_from_tables(request_data["connection"], + included_tables) + if fits_in_window(user_prompt + results_sample): + user_prompt += f"## Table Structure with Sampled Data ##\n\n{results_sample}\n\n" + else: # If it's too large with sampling, then provide *just* the structure + table_struct = tables_from_schema_info(request_data["connection"], + included_tables) + user_prompt += f"## Table Structure ##\n\n{table_struct}\n\n" + + user_prompt += f"## User's Request to Assistant ##\n\n{request_data['assistant_request']}\n\n" + + prompt = { + "system": ExplorerValue.objects.get_item(ExplorerValue.ASSISTANT_SYSTEM_PROMPT).value, + "user": user_prompt + } + return prompt diff --git a/explorer/assistant/views.py b/explorer/assistant/views.py index 05598ab6..ff557619 100644 --- a/explorer/assistant/views.py +++ b/explorer/assistant/views.py @@ -1,53 +1,24 @@ from django.http import JsonResponse +from django.views import View from django.utils import timezone -from django.views.decorators.http import require_POST import json from explorer.telemetry import Stat, StatNames -from explorer.utils import get_valid_connection -from explorer.models import ExplorerValue from explorer.assistant.models import PromptLog from explorer.assistant.utils import ( - do_req, extract_response, tables_from_schema_info, - get_table_names_from_query, sample_rows_from_tables, - fits_in_window + do_req, extract_response, + get_table_names_from_query, + build_prompt ) def run_assistant(request_data, user): - user_prompt = "" - - db_vendor = get_valid_connection(request_data.get("connection")).vendor - user_prompt += f"## Database Vendor / SQL Flavor is {db_vendor}\n\n" - - db_error = request_data.get("db_error") - if db_error: - user_prompt += f"## Query Error ##\n\n{db_error}\n\n" - sql = request_data.get("sql") - if sql: - user_prompt += f"## Existing SQL ##\n\n{sql}\n\n" - extra_tables = request_data.get("selected_tables", []) included_tables = get_table_names_from_query(sql) + extra_tables - if included_tables: - table_struct = tables_from_schema_info(request_data["connection"], - included_tables) - user_prompt += f"## Relevant Tables' Structure ##\n\n{table_struct}\n\n" - - if fits_in_window(user_prompt): - results_sample = sample_rows_from_tables(request_data["connection"], - included_tables) - if fits_in_window(user_prompt + results_sample): - user_prompt += f"## Sample Results from Tables ##\n\n{results_sample}\n\n" - user_prompt += f"## User's Request to Assistant ##\n\n{request_data['assistant_request']}\n\n" - - prompt = { - "system": ExplorerValue.objects.get_item(ExplorerValue.ASSISTANT_SYSTEM_PROMPT).value, - "user": user_prompt - } + prompt = build_prompt(request_data, included_tables) start = timezone.now() pl = PromptLog( @@ -75,19 +46,16 @@ def run_assistant(request_data, user): return response_text -@require_POST -def assistant_help(request): - try: - data = json.loads(request.body) - - resp = run_assistant(data, request.user) - - response_data = { - "status": "success", - "message": resp - } - - return JsonResponse(response_data) +class AssistantHelpView(View): - except json.JSONDecodeError: - return JsonResponse({"status": "error", "message": "Invalid JSON"}, status=400) + def post(self, request, *args, **kwargs): + try: + data = json.loads(request.body) + resp = run_assistant(data, request.user) + response_data = { + "status": "success", + "message": resp + } + return JsonResponse(response_data) + except json.JSONDecodeError: + return JsonResponse({"status": "error", "message": "Invalid JSON"}, status=400) diff --git a/explorer/connections.py b/explorer/connections.py index 03427e7a..e1921bf9 100644 --- a/explorer/connections.py +++ b/explorer/connections.py @@ -1,6 +1,40 @@ from django.db import connections as djcs +from django.db import transaction, DEFAULT_DB_ALIAS -from explorer.app_settings import EXPLORER_CONNECTIONS +from explorer.ee.db_connections.utils import create_django_style_connection +from explorer import app_settings +from explorer.models import DatabaseConnection + +# To support user-configured database connections that can be managed through the Explorer UI, *as well* as the +# 'legacy' connections that are configured in Django's normal settings.DATABASES config, we stitch together the two. + +# We allow queries to be associated with either type of connection, seamlessly. + +# The approach is to allow users to create connections with approximately the same parameters that a settings.DATABASE +# would expect. We then stitch them together into one list. When Explorer needs to access a connection, it coughs up a +# Django DatabaseWrapper connection in either case (natively, if it's coming from settings.DATABASES, or by taking the +# user-created connection and running it through the create_django_style_connection() function in this file). + +# In general, amazingly, this "just works" and the entire application is perfectly happy to use either type as a +# connection. The exception to this is that there are a few bits of code that ultimately (or directly) use the +# django.db.transaction.atomic context manager. For some reason that particular Django innard takes an *alias*, not a +# proper connection. Then it retrieves the connection based on that alias. But of course if we are providing a +# user-created connection alias, Django doesn't find it (because it is looking in settings.DATABASES). + +# The solution is to monkey-patch the get_connection function that transaction.atomic uses, to make it aware of the +# user-created connections. + + +def new_get_connection(using=None): + if using is None: + using = DEFAULT_DB_ALIAS + if using in djcs: + return djcs[using] + return create_django_style_connection(DatabaseConnection.objects.get(alias=using)) + + +# Monkey patch +transaction.get_connection = new_get_connection # We export valid SQL connections here so that consuming code never has to @@ -10,15 +44,19 @@ # Django insists that connections that are created in a thread are only accessed # by that thread, so here we create a dictionary-like collection of the valid # connections, but does a 'live' lookup of the connection on each item access. - - -_connections = {c: c for c in djcs if c in EXPLORER_CONNECTIONS.values()} - - class ExplorerConnections(dict): def __getitem__(self, item): - return djcs[item] + if item in djcs: + return djcs[item] + else: + return create_django_style_connection(DatabaseConnection.objects.get(alias=item)) + + +def connections(): + _connections = [c for c in djcs if c in app_settings.EXPLORER_CONNECTIONS.values()] + db_connections = DatabaseConnection.objects.all() + _connections += [c.alias for c in db_connections] + return ExplorerConnections(zip(_connections, _connections)) -connections = ExplorerConnections(_connections) diff --git a/explorer/ee/LICENSE b/explorer/ee/LICENSE new file mode 100644 index 00000000..4d007618 --- /dev/null +++ b/explorer/ee/LICENSE @@ -0,0 +1,21 @@ +** Additional License for "explorer/ee/" Directory ** + +All content that resides under the "explorer/ee/" directory of this repository is provided under the MIT License, +except that the following additional rights are reserved: + +** "Commons Clause" License Condition v1.0 ** + +The Software may not be used as part of any commercial offering or service, such as hosting, service, or consulting +offerings. For purposes of the foregoing, "commercial offering" means the provision of the Software to third parties +for a fee or other consideration, including without limitation, providing the Software as part of a managed service, +as part of a platform as a service, or providing it to third parties on a software as a service basis. + +This restriction does not apply to non-commercial use by any individual, nor does it prevent such individual from +providing services to third parties using the Software. + +** Exceptions ** + +SQL Explorer, Inc. may grant a Commercial License Agreement to provide exceptions to the "Commons Clause" License +Condition. If you wish to obtain such a license, please contact SQL Explorer, Inc. at support@sqlexplorer.io. No +exception is granted implicitly or explicitly without a written and signed Commercial License Agreement from SQL +Explorer, Inc. diff --git a/explorer/ee/__init__.py b/explorer/ee/__init__.py new file mode 100644 index 00000000..8b137891 --- /dev/null +++ b/explorer/ee/__init__.py @@ -0,0 +1 @@ + diff --git a/explorer/ee/db_connections/__init__.py b/explorer/ee/db_connections/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/explorer/ee/db_connections/admin.py b/explorer/ee/db_connections/admin.py new file mode 100644 index 00000000..f8b26789 --- /dev/null +++ b/explorer/ee/db_connections/admin.py @@ -0,0 +1,9 @@ +from django.contrib import admin + +from explorer.models import DatabaseConnection + + +@admin.register(DatabaseConnection) +class DatabaseConnectionAdmin(admin.ModelAdmin): + pass + diff --git a/explorer/ee/db_connections/forms.py b/explorer/ee/db_connections/forms.py new file mode 100644 index 00000000..04401efe --- /dev/null +++ b/explorer/ee/db_connections/forms.py @@ -0,0 +1,17 @@ +from django import forms +from explorer.ee.db_connections.models import DatabaseConnection + + +class DatabaseConnectionForm(forms.ModelForm): + class Meta: + model = DatabaseConnection + fields = "__all__" + widgets = { + "alias": forms.TextInput(attrs={"class": "form-control"}), + "engine": forms.Select(attrs={"class": "form-select"}), + "name": forms.TextInput(attrs={"class": "form-control"}), + "user": forms.TextInput(attrs={"class": "form-control"}), + "password": forms.PasswordInput(attrs={"class": "form-control"}), + "host": forms.TextInput(attrs={"class": "form-control"}), + "port": forms.TextInput(attrs={"class": "form-control"}), + } diff --git a/explorer/ee/db_connections/models.py b/explorer/ee/db_connections/models.py new file mode 100644 index 00000000..627a436e --- /dev/null +++ b/explorer/ee/db_connections/models.py @@ -0,0 +1,63 @@ +import os + +from django.conf import settings +from django.core.exceptions import ValidationError +from django.db import models +from django.db.models.signals import pre_save +from django.dispatch import receiver + +from django_cryptography.fields import encrypt + + +class DatabaseConnection(models.Model): + + SQLITE = "django.db.backends.sqlite3" + + DATABASE_ENGINES = ( + (SQLITE, "SQLite3"), + ("django.db.backends.postgresql", "PostgreSQL"), + ("django.db.backends.mysql", "MySQL"), + ("django.db.backends.oracle", "Oracle"), + ("django.db.backends.mysql", "MariaDB"), + ("django_cockroachdb", "CockroachDB"), + ("django.db.backends.sqlserver", "SQL Server (mssql-django)"), + ) + + alias = models.CharField(max_length=255, unique=True) + engine = models.CharField(max_length=255, choices=DATABASE_ENGINES) + name = models.CharField(max_length=255) + user = encrypt(models.CharField(max_length=255, blank=True)) + password = encrypt(models.CharField(max_length=255, blank=True)) + host = encrypt(models.CharField(max_length=255, blank=True)) + port = models.CharField(max_length=255, blank=True) + + def __str__(self): + return f"{self.name} ({self.alias})" + + @property + def is_upload(self): + return self.engine == self.SQLITE and self.host + + @property + def local_name(self): + if self.is_upload: + return os.path.normpath(os.path.join(os.getcwd(), f"user_dbs/{self.name}")) + + @classmethod + def from_django_connection(cls, connection_alias): + conn = settings.DATABASES.get(connection_alias) + if conn: + return DatabaseConnection( + alias=connection_alias, + engine=conn.get("ENGINE"), + name=conn.get("NAME"), + user=conn.get("USER"), + password=conn.get("PASSWORD"), + host=conn.get("HOST"), + port=conn.get("PORT"), + ) + +@receiver(pre_save, sender=DatabaseConnection) +def validate_database_connection(sender, instance, **kwargs): + if instance.name in settings.DATABASES.keys(): + raise ValidationError(f"Database name '{instance.name}' already exists.") diff --git a/explorer/ee/db_connections/utils.py b/explorer/ee/db_connections/utils.py new file mode 100644 index 00000000..7eae7e30 --- /dev/null +++ b/explorer/ee/db_connections/utils.py @@ -0,0 +1,170 @@ +from django.db import DatabaseError +from django.db.utils import load_backend +import os + +import locale +from dateutil import parser + +import pandas as pd +import sqlite3 +import io + + +# TODO deal with uploading the same file / conflicting file again +def upload_sqlite(db_bytes, path): + from explorer.utils import get_s3_bucket + bucket = get_s3_bucket() + bucket.put_object(Key=path, Body=db_bytes, ServerSideEncryption="AES256") + + +def create_connection_for_uploaded_sqlite(filename, s3_path): + from explorer.models import DatabaseConnection + return DatabaseConnection.objects.create( + alias=filename, + engine=DatabaseConnection.SQLITE, + name=filename, + host=s3_path + ) + + +def get_sqlite_for_connection(explorer_connection): + from explorer.utils import get_s3_bucket + + # Get the database from s3, then modify the connection to work with the downloaded file. + # E.g. "host" should not be set, and we need to get the full path to the file + local_name = explorer_connection.local_name + if not os.path.exists(local_name): + s3 = get_s3_bucket() + s3.download_file(explorer_connection.host, local_name) + explorer_connection.host = None + explorer_connection.name = local_name + return explorer_connection + + +def create_django_style_connection(explorer_connection): + + if explorer_connection.is_upload: + explorer_connection = get_sqlite_for_connection(explorer_connection) + + connection_settings = { + "ENGINE": explorer_connection.engine, + "NAME": explorer_connection.name, + "USER": explorer_connection.user, + "PASSWORD": explorer_connection.password, + "HOST": explorer_connection.host, + "PORT": explorer_connection.port, + "TIME_ZONE": None, + "CONN_MAX_AGE": 0, + "CONN_HEALTH_CHECKS": False, + "OPTIONS": {}, + "TEST": {}, + "AUTOCOMMIT": True, + "ATOMIC_REQUESTS": False, + } + + try: + backend = load_backend(explorer_connection.engine) + return backend.DatabaseWrapper(connection_settings, explorer_connection.alias) + except DatabaseError as e: + raise DatabaseError(f"Failed to create explorer connection: {e}") from e + + +def pandas_to_sqlite(df, local_path="local_database.db"): + # Write the DataFrame to a local SQLite database + # In theory, it would be nice to write the dataframe to an in-memory SQLite DB, and then dump the bytes from that + # but there is no way to get to the underlying bytes from an in-memory SQLite DB + con = sqlite3.connect(local_path) + try: + df.to_sql(name="data", con=con, if_exists="replace", index=False) + finally: + con.close() + + # Read the local SQLite database file into a BytesIO buffer + try: + db_file = io.BytesIO() + with open(local_path, "rb") as f: + db_file.write(f.read()) + db_file.seek(0) + return db_file + finally: + # Delete the local SQLite database file + # Finally block to ensure we don't litter files around + os.remove(local_path) + + +MAX_TYPING_SAMPLE_SIZE = 10000 +SHORTEST_PLAUSIBLE_DATE_STRING = 5 + + +def csv_to_typed_df(csv_bytes, delimiter=",", has_headers=True): # noqa + try: + + csv_file = io.BytesIO(csv_bytes) + df = pd.read_csv(csv_file, sep=delimiter, header=0 if has_headers else None) + locale.setlocale(locale.LC_NUMERIC, "en_US.UTF-8") + + for column in df.columns: + values = df[column].dropna().unique() + if len(values) > MAX_TYPING_SAMPLE_SIZE: + values = pd.Series(values).sample(MAX_TYPING_SAMPLE_SIZE, random_state=42).to_numpy() + + is_date = False + is_integer = True + is_float = True + + for value in values: + try: + float_val = locale.atof(str(value)) + if float_val == int(float_val): + continue # This is effectively an integer + else: + is_integer = False + except ValueError: + is_integer = False + is_float = False + break + + if is_integer: + is_float = False + + if not is_integer and not is_float: + is_date = True + + # The dateutil parser is very aggressive and will interpret many short strings as dates. + # For example "12a" will be interpreted as 12:00 AM on the current date. + # That is not the behavior anyone wants. The shortest plausible date string is e.g. 1-1-23 + try_parse = [v for v in values if len(str(v)) > SHORTEST_PLAUSIBLE_DATE_STRING] + if len(try_parse) > 0: + for value in try_parse: + try: + parser.parse(str(value)) + except (ValueError, TypeError, OverflowError): + is_date = False + break + else: + is_date = False + + if is_date: + df[column] = pd.to_datetime(df[column], errors="coerce", utc=True) + elif is_integer: + df[column] = df[column].apply(lambda x: int(locale.atof(str(x))) if pd.notna(x) else x) + # If there are NaN / blank values, the column will be converted to float + # Convert it back to integer + df[column] = df[column].astype("Int64") + elif is_float: + df[column] = df[column].apply(lambda x: locale.atof(str(x)) if pd.notna(x) else x) + else: + inferred_type = pd.api.types.infer_dtype(values) + if inferred_type == "integer": + df[column] = pd.to_numeric(df[column], errors="coerce", downcast="integer") + elif inferred_type == "floating": + df[column] = pd.to_numeric(df[column], errors="coerce") + + return df + + except pd.errors.ParserError as e: + return str(e) + + +def is_csv(file): + return file.content_type == "text/csv" diff --git a/explorer/ee/db_connections/views.py b/explorer/ee/db_connections/views.py new file mode 100644 index 00000000..d4b1383b --- /dev/null +++ b/explorer/ee/db_connections/views.py @@ -0,0 +1,127 @@ +from django.views.generic import ListView, DetailView, CreateView, UpdateView, DeleteView +from django.views import View +from django.http import JsonResponse +from django.urls import reverse_lazy +from explorer.models import DatabaseConnection +from explorer.ee.db_connections.utils import ( + upload_sqlite, + create_connection_for_uploaded_sqlite, + is_csv, + csv_to_typed_df, + pandas_to_sqlite +) +from explorer import app_settings +from explorer.ee.db_connections.forms import DatabaseConnectionForm +from explorer.utils import delete_from_s3 +from explorer.views.auth import PermissionRequiredMixin +from explorer.views.mixins import ExplorerContextMixin +from explorer.ee.db_connections.utils import create_django_style_connection + + +class UploadDbView(PermissionRequiredMixin, View): + + permission_required = "connections_permission" + + def post(self, request): + try: + file = request.FILES.get("file") + if file: + f_name = file.name + f_bytes = file.read() + + if is_csv(file): + df = csv_to_typed_df(f_bytes) + f_bytes = pandas_to_sqlite(df) + f_name = f_name.replace("csv", "db") + s3_path = f"user_dbs/user_{request.user.id}/{f_name}" + upload_sqlite(f_bytes, s3_path) + create_connection_for_uploaded_sqlite(f_name, s3_path) + return JsonResponse({"success": True}) + else: + # TODO handle this error client side. + return JsonResponse({"error": "No file provided"}, status=400) + except Exception as e: + return JsonResponse({"error": str(e)}, status=400) + + +class DatabaseConnectionsListView(PermissionRequiredMixin, ExplorerContextMixin, ListView): + + context_object_name = "sqlite_uploads" + permission_required = "connections_permission" + template_name = "connections/connections.html" + model = DatabaseConnection + + def get_queryset(self): + qs = list(DatabaseConnection.objects.all()) + for _, alias in app_settings.EXPLORER_CONNECTIONS.items(): + django_conn = DatabaseConnection.from_django_connection(alias) + if django_conn: + qs.append(django_conn) + return qs + + +class DatabaseConnectionDetailView(PermissionRequiredMixin, DetailView): + permission_required = "connections_permission" + model = DatabaseConnection + template_name = "connections/database_connection_detail.html" + + +class DatabaseConnectionCreateView(PermissionRequiredMixin, ExplorerContextMixin, CreateView): + permission_required = "connections_permission" + model = DatabaseConnection + form_class = DatabaseConnectionForm + template_name = "connections/database_connection_form.html" + success_url = reverse_lazy("explorer_connections") + + +class DatabaseConnectionUpdateView(PermissionRequiredMixin, UpdateView): + permission_required = "connections_permission" + model = DatabaseConnection + form_class = DatabaseConnectionForm + template_name = "connections/database_connection_form.html" + success_url = reverse_lazy("explorer_connections") + + +class DatabaseConnectionDeleteView(PermissionRequiredMixin, DeleteView): + permission_required = "connections_permission" + model = DatabaseConnection + template_name = "connections/database_connection_confirm_delete.html" + success_url = reverse_lazy("explorer_connections") + + def delete(self, request, *args, **kwargs): + connection = self.get_object() + if connection.is_upload: + delete_from_s3(connection.host) + return super().delete(request, *args, **kwargs) + + +class DatabaseConnectionValidateView(PermissionRequiredMixin, View): + + permission_required = "connections_permission" + + def post(self, request): + form = DatabaseConnectionForm(request.POST) + + instance = DatabaseConnection.objects.filter(alias=request.POST["alias"]).first() + if instance: + form = DatabaseConnectionForm(request.POST, instance=instance) + if form.is_valid(): + connection_data = form.cleaned_data + explorer_connection = DatabaseConnection( + alias=connection_data["alias"], + engine=connection_data["engine"], + name=connection_data["name"], + user=connection_data["user"], + password=connection_data["password"], + host=connection_data["host"], + port=connection_data["port"] + ) + try: + conn = create_django_style_connection(explorer_connection) + with conn.cursor() as cursor: + cursor.execute("SELECT 1") + return JsonResponse({"success": True}) + except Exception as e: + return JsonResponse({"success": False, "error": str(e)}) + else: + return JsonResponse({"success": False, "error": "Invalid form data"}) diff --git a/explorer/ee/urls.py b/explorer/ee/urls.py new file mode 100644 index 00000000..cf0ba16b --- /dev/null +++ b/explorer/ee/urls.py @@ -0,0 +1,21 @@ +from django.urls import path + +from explorer.ee.db_connections.views import ( + UploadDbView, + DatabaseConnectionsListView, + DatabaseConnectionCreateView, + DatabaseConnectionDetailView, + DatabaseConnectionUpdateView, + DatabaseConnectionDeleteView, + DatabaseConnectionValidateView +) + +ee_urls = [ + path("connections/", DatabaseConnectionsListView.as_view(), name="explorer_connections"), + path("connections/upload/", UploadDbView.as_view(), name="explorer_upload"), + path("connections//", DatabaseConnectionDetailView.as_view(), name="explorer_connection_detail"), + path("connections/new/", DatabaseConnectionCreateView.as_view(), name="explorer_connection_create"), + path("connections//edit/", DatabaseConnectionUpdateView.as_view(), name="explorer_connection_update"), + path("connections//delete/", DatabaseConnectionDeleteView.as_view(), name="explorer_connection_delete"), + path("connections/validate/", DatabaseConnectionValidateView.as_view(), name="explorer_connection_validate"), +] diff --git a/explorer/forms.py b/explorer/forms.py index db631adc..cd9dd2f7 100644 --- a/explorer/forms.py +++ b/explorer/forms.py @@ -1,7 +1,7 @@ from django.forms import BooleanField, CharField, ModelForm, ValidationError from django.forms.widgets import CheckboxInput, Select -from explorer.app_settings import EXPLORER_CONNECTIONS, EXPLORER_DEFAULT_CONNECTION +from explorer.app_settings import EXPLORER_DEFAULT_CONNECTION from explorer.models import MSG_FAILED_BLACKLIST, Query @@ -54,7 +54,8 @@ def created_at_time(self): @property def connections(self): - return [(v, k) for k, v in EXPLORER_CONNECTIONS.items()] + from explorer.connections import connections + return [(c, c) for c in connections()] class Meta: model = Query diff --git a/explorer/locale/ru/LC_MESSAGES/django.po b/explorer/locale/ru/LC_MESSAGES/django.po index 97baaaa1..4776d52c 100644 --- a/explorer/locale/ru/LC_MESSAGES/django.po +++ b/explorer/locale/ru/LC_MESSAGES/django.po @@ -211,7 +211,7 @@ msgstr "Вы уверены в удалении «%(title)s»?" #: explorer/templates/explorer/query_list.html:15 #, python-format -msgid "Your %(qlen)s Most Recently Run" +msgid "Recently Run by You" msgstr "Ваши последние запуски запросов, их %(qlen)s" #: explorer/templates/explorer/query_list.html:23 diff --git a/explorer/migrations/0017_databaseconnection.py b/explorer/migrations/0017_databaseconnection.py new file mode 100644 index 00000000..2d5c1e7c --- /dev/null +++ b/explorer/migrations/0017_databaseconnection.py @@ -0,0 +1,26 @@ +# Generated by Django 5.0.4 on 2024-05-07 18:41 + +from django.db import migrations, models + + +class Migration(migrations.Migration): + + dependencies = [ + ('explorer', '0016_alter_explorervalue_key'), + ] + + operations = [ + migrations.CreateModel( + name='DatabaseConnection', + fields=[ + ('id', models.AutoField(auto_created=True, primary_key=True, serialize=False, verbose_name='ID')), + ('alias', models.CharField(max_length=255, unique=True)), + ('engine', models.CharField(choices=[('django.db.backends.sqlite3', 'SQLite3'), ('django.db.backends.postgresql_psycopg2', 'PostgreSQL'), ('django.db.backends.mysql', 'MySQL'), ('django.db.backends.oracle', 'Oracle')], max_length=255)), + ('name', models.CharField(max_length=255)), + ('user', models.CharField(blank=True, max_length=255)), + ('password', models.CharField(blank=True, max_length=255)), + ('host', models.CharField(blank=True, max_length=255)), + ('port', models.CharField(blank=True, max_length=255)), + ], + ), + ] diff --git a/explorer/migrations/0018_alter_databaseconnection_host_and_more.py b/explorer/migrations/0018_alter_databaseconnection_host_and_more.py new file mode 100644 index 00000000..329afaa7 --- /dev/null +++ b/explorer/migrations/0018_alter_databaseconnection_host_and_more.py @@ -0,0 +1,29 @@ +# Generated by Django 5.0.4 on 2024-05-14 15:55 + +import django_cryptography.fields +from django.db import migrations, models + + +class Migration(migrations.Migration): + + dependencies = [ + ('explorer', '0017_databaseconnection'), + ] + + operations = [ + migrations.AlterField( + model_name='databaseconnection', + name='host', + field=django_cryptography.fields.encrypt(models.CharField(blank=True, max_length=255)), + ), + migrations.AlterField( + model_name='databaseconnection', + name='password', + field=django_cryptography.fields.encrypt(models.CharField(blank=True, max_length=255)), + ), + migrations.AlterField( + model_name='databaseconnection', + name='user', + field=django_cryptography.fields.encrypt(models.CharField(blank=True, max_length=255)), + ), + ] diff --git a/explorer/models.py b/explorer/models.py index f5785747..d2d5cba0 100644 --- a/explorer/models.py +++ b/explorer/models.py @@ -15,6 +15,9 @@ shared_dict_update, swap_params, ) +# Issue #618. All models must be imported so that Django understands how to manage migrations for the app +from explorer.ee.db_connections.models import DatabaseConnection # noqa +from explorer.assistant.models import PromptLog # noqa MSG_FAILED_BLACKLIST = "Query failed the SQL blacklist: %s" @@ -64,6 +67,10 @@ def __str__(self): def get_run_count(self): return self.querylog_set.count() + def last_run_log(self): + ql = self.querylog_set.first() + return ql or QueryLog(success=True, run_at=self.created_at) + def avg_duration_display(self): d = self.avg_duration() if d: @@ -439,7 +446,7 @@ class ExplorerValue(models.Model): EXPLORER_SETTINGS_CHOICES = [ (INSTALL_UUID, "Install Unique ID"), (STARTUP_METRIC_LAST_SEND, "Startup metric last send"), - (ASSISTANT_SYSTEM_PROMPT, "System prompt for SQL Assistant") + (ASSISTANT_SYSTEM_PROMPT, "System prompt for SQL Assistant"), ] key = models.CharField(max_length=5, choices=EXPLORER_SETTINGS_CHOICES, unique=True) diff --git a/explorer/permissions.py b/explorer/permissions.py index c9d7ff38..8704decf 100644 --- a/explorer/permissions.py +++ b/explorer/permissions.py @@ -23,3 +23,7 @@ def view_permission_list(request): def change_permission(request, *args, **kwargs): return app_settings.EXPLORER_PERMISSION_CHANGE(request) + + +def connections_permission(request, **kwargs): + return app_settings.EXPLORER_PERMISSION_CONNECTIONS(request) diff --git a/explorer/src/js/assistant.js b/explorer/src/js/assistant.js index 9b85c8f7..e424b6b4 100644 --- a/explorer/src/js/assistant.js +++ b/explorer/src/js/assistant.js @@ -3,7 +3,7 @@ import { marked } from "marked"; import DOMPurify from "dompurify"; import * as bootstrap from 'bootstrap'; import List from "list.js"; -import { SchemaSvc } from "./schemaService" +import { SchemaSvc, getConnElement } from "./schemaService" function getErrorMessage() { const errorElement = document.querySelector('.alert-danger.db-error'); @@ -11,8 +11,7 @@ function getErrorMessage() { } function setupTableList() { - const conn = document.querySelector('#id_connection').value; - SchemaSvc.get(conn).then(schema => { + SchemaSvc.get().then(schema => { const keys = Object.keys(schema); const tableList = document.getElementById('table-list'); tableList.innerHTML = ''; @@ -42,6 +41,23 @@ function setupTableList() { }; new List('additional_table_container', options); + + const selectAllButton = document.getElementById('select_all_button'); + const checkboxes = document.querySelectorAll('.table-checkbox'); + + let selectState = 'all'; + + selectAllButton.innerHTML = 'Select All'; + + selectAllButton.addEventListener('click', (e) => { + e.preventDefault(); + const isSelectingAll = selectState === 'all'; + checkboxes.forEach((checkbox) => { + checkbox.checked = isSelectingAll; + }); + selectState = isSelectingAll ? 'none' : 'all'; + selectAllButton.innerHTML = isSelectingAll ? 'Deselect All' : 'Select All'; + }); }) .catch(error => { console.error('Error retrieving JSON schema:', error); @@ -50,8 +66,7 @@ function setupTableList() { export function setUpAssistant(expand = false) { - const connEl = document.querySelector('#id_connection'); - connEl.addEventListener('change', setupTableList); + getConnElement().addEventListener('change', setupTableList); setupTableList(); const error = getErrorMessage(); @@ -129,6 +144,13 @@ function submitAssistantAsk() { }); } + // Similarly, if there is no description, copy the prompt into the description + const prompt = document.getElementById("id_assistant_input")?.value; + const description = document.getElementById("id_description"); + if (description?.value === "") { + description.value = prompt; + } + setUpCopyButtons(); }) .catch(error => { diff --git a/explorer/src/js/codemirror-config.js b/explorer/src/js/codemirror-config.js index 157794c9..4343eb9c 100644 --- a/explorer/src/js/codemirror-config.js +++ b/explorer/src/js/codemirror-config.js @@ -50,8 +50,7 @@ function fetchAndShowSchema(view) { if (wordRange) { const tableName = state.doc.sliceString(wordRange.from, wordRange.to); - const conn = document.querySelector('#id_connection').value; - SchemaSvc.get(conn).then(schema => { + SchemaSvc.get().then(schema => { let formattedSchema; if (schema.hasOwnProperty(tableName)) { formattedSchema = JSON.stringify(schema[tableName], null, 2); diff --git a/explorer/src/js/explorer.js b/explorer/src/js/explorer.js index c0a0da01..b76d6207 100644 --- a/explorer/src/js/explorer.js +++ b/explorer/src/js/explorer.js @@ -11,14 +11,11 @@ import { toggleFavorite } from "./favorites"; import {schemaCompletionSource, StandardSQL} from "@codemirror/lang-sql"; import {StateEffect} from "@codemirror/state"; -import {SchemaSvc} from "./schemaService"; +import {getConnElement, SchemaSvc} from "./schemaService"; function updateSchema() { - - const conn = document.querySelector('#id_connection').value; - - SchemaSvc.get(conn).then(schema => { + SchemaSvc.get().then(schema => { window.editor.dispatch({ effects: StateEffect.appendConfig.of( StandardSQL.language.data.of({ @@ -28,7 +25,7 @@ function updateSchema() { }); }); - $("#schema_frame").attr("src", `../schema/${conn}`); + $("#schema_frame").attr("src", `../schema/${getConnElement().value}`); } @@ -46,9 +43,24 @@ function editorFromTextArea(textarea) { return view } + +function selectConnection() { + var urlParams = new URLSearchParams(window.location.search); + var connectionId = urlParams.get('connection'); + + if (connectionId) { + var connectionSelect = document.getElementById('id_connection'); + if (connectionSelect) { + connectionSelect.value = connectionId; + } + } +} + export class ExplorerEditor { constructor(queryId) { + selectConnection(); + const aa = document.getElementById('assistant_accordion'); const pa = document.getElementById('nav-preview'); if (aa) { @@ -365,8 +377,7 @@ export class ExplorerEditor { }.bind(this)); // Set up schema autocomplete in the editor. When the connection changes, load new schema. - const connEl = document.querySelector('#id_connection'); - connEl.addEventListener('change', updateSchema); + getConnElement().addEventListener('change', updateSchema); updateSchema(); } } diff --git a/explorer/src/js/main.js b/explorer/src/js/main.js index 2ab7188b..419b1b04 100644 --- a/explorer/src/js/main.js +++ b/explorer/src/js/main.js @@ -15,7 +15,9 @@ const route_initializers = { new ExplorerEditor(document.getElementById('queryIdGlobal').value)), query_create: () => import('./explorer').then(({ExplorerEditor}) => new ExplorerEditor('new')), explorer_playground: () => import('./explorer').then(({ExplorerEditor}) => new ExplorerEditor('new')), - explorer_schema: () => import('./schema').then(({setupSchema}) => setupSchema()) + explorer_schema: () => import('./schema').then(({setupSchema}) => setupSchema()), + explorer_connection_create: () => import('./uploads').then(({setupUploads}) => setupUploads()), + explorer_connection_update: () => import('./uploads').then(({setupUploads}) => setupUploads()) }; document.addEventListener('DOMContentLoaded', function() { diff --git a/explorer/src/js/query-list.js b/explorer/src/js/query-list.js index abe8ca5e..cd936441 100644 --- a/explorer/src/js/query-list.js +++ b/explorer/src/js/query-list.js @@ -16,7 +16,7 @@ export function setupQueryList() { }); let options = { - valueNames: ['name'], + valueNames: ['sort-name', 'sort-created', 'sort-created', 'sort-last-run', 'sort-run-count', 'sort-connection'], handlers: {'updated': [searchFocus]} }; new List('queries', options); diff --git a/explorer/src/js/schemaService.js b/explorer/src/js/schemaService.js index f0b075ec..7f5e49bf 100644 --- a/explorer/src/js/schemaService.js +++ b/explorer/src/js/schemaService.js @@ -1,6 +1,8 @@ const schemaCache = {}; -const fetchSchema = async (conn) => { +const fetchSchema = async () => { + + const conn = getConnElement().value; if (schemaCache[conn]) { return schemaCache[conn]; @@ -23,3 +25,7 @@ const fetchSchema = async (conn) => { export const SchemaSvc = { get: fetchSchema }; + +export function getConnElement() { + return document.querySelector('#id_connection'); +} diff --git a/explorer/src/js/uploads.js b/explorer/src/js/uploads.js new file mode 100644 index 00000000..c1b9a16e --- /dev/null +++ b/explorer/src/js/uploads.js @@ -0,0 +1,83 @@ +import { getCsrfToken } from "./csrf"; + +export function setupUploads() { + var dropArea = document.getElementById('drop-area'); + var fileElem = document.getElementById('fileElem'); + + if (dropArea) { + dropArea.onclick = function() { + fileElem.click(); + }; + + dropArea.addEventListener('dragover', function(e) { + e.preventDefault(); // Prevent default behavior (Prevent file from being opened) + dropArea.classList.add('bg-info'); // Optional: add a style when dragging over + }); + + dropArea.addEventListener('dragleave', function(e) { + dropArea.classList.remove('bg-info'); // Optional: remove style when not dragging over + }); + + dropArea.addEventListener('drop', function(e) { + e.preventDefault(); + dropArea.classList.remove('bg-info'); // Optional: remove style after dropping + + let files = e.dataTransfer.files; + if (files.length) { + handleFiles(files[0]); // Assuming only one file is dropped + } + }); + + fileElem.onchange = function() { + if (this.files.length) { + handleFiles(this.files[0]); + } + }; + } + + + + function handleFiles(file) { + uploadFile(file); + } + + function uploadFile(file) { + let formData = new FormData(); + formData.append('file', file); + + fetch("../upload/", { + method: 'POST', + headers: { + 'X-CSRFToken': getCsrfToken() + }, + body: formData + }).then(response => response.json()) + .then(() => { + window.location.reload(); + }) + .catch(error => console.error('Error:', error)); + } + + document.getElementById("test-connection-btn").addEventListener("click", function() { + var form = document.getElementById("db-connection-form"); + var formData = new FormData(form); + + fetch("../../validate/", { + method: "POST", + body: formData, + headers: { + "X-CSRFToken": getCsrfToken() + } + }) + .then(response => response.json()) + .then(data => { + if (data.success) { + alert("Connection successful!"); + } else { + alert("Connection failed: " + data.error); + } + }) + .catch(error => console.error("Error:", error)); + }); + +} diff --git a/explorer/tasks.py b/explorer/tasks.py index ea8f2f05..66ad2dc5 100644 --- a/explorer/tasks.py +++ b/explorer/tasks.py @@ -1,6 +1,7 @@ import io import random import string +import os from datetime import date, datetime, timedelta from django.core.cache import cache @@ -9,6 +10,7 @@ from explorer import app_settings from explorer.exporters import get_exporter_class from explorer.models import Query, QueryLog +from explorer.ee.db_connections.models import DatabaseConnection if app_settings.ENABLE_TASKS: @@ -114,3 +116,16 @@ def build_schema_cache_async(connection_alias): cache.set(connection_schema_json_cache_key(connection_alias), transform_to_json_schema(ret)) return ret + + +# TODO tests +@shared_task +def remove_unused_sqlite_dbs(): + uploaded_dbs = DatabaseConnection.objects.filter(engine=DatabaseConnection.SQLITE, + host__isnull=False) + for db in uploaded_dbs: + if os.path.exists(db.local_name): + recent_run = QueryLog.objects.filter(connection=db.alias).first() + days = app_settings.EXPLORER_PRUNE_LOCAL_UPLOAD_COPY_DAYS_INACTIVITY + if recent_run and (datetime.now() - timedelta(days=days)) > recent_run.run_at: + os.remove(db.local_name) diff --git a/explorer/telemetry.py b/explorer/telemetry.py index b8947b28..fc460b2e 100644 --- a/explorer/telemetry.py +++ b/explorer/telemetry.py @@ -72,6 +72,7 @@ def send_summary_stats(self): def track(self): from explorer import app_settings + if not app_settings.EXPLORER_ENABLE_ANONYMOUS_STATS: return @@ -119,6 +120,7 @@ def _gather_summary_stats(): from explorer import app_settings from explorer.models import Query, QueryLog + from explorer.ee.db_connections.models import DatabaseConnection import explorer try: @@ -144,7 +146,8 @@ def _gather_summary_stats(): "tasks_enabled": app_settings.ENABLE_TASKS, "unsafe_rendering": app_settings.UNSAFE_RENDERING, "transform_count": len(app_settings.EXPLORER_TRANSFORMS), - "assistant_enabled": app_settings.EXPLORER_AI_API_KEY is not None, + "assistant_enabled": app_settings.has_assistant(), + "user_dbs": DatabaseConnection.objects.count(), "version": explorer.get_version(), "charts_enabled": app_settings.EXPLORER_CHARTS_ENABLED } diff --git a/explorer/templates/connections/connections.html b/explorer/templates/connections/connections.html new file mode 100644 index 00000000..a38153cd --- /dev/null +++ b/explorer/templates/connections/connections.html @@ -0,0 +1,43 @@ +{% extends "explorer/base.html" %} +{% load explorer_tags i18n %} + +{% block sql_explorer_content %} +
+

Connections

+
+ Add New Connection + + + + + + + + + + + {% for connection in object_list %} + + + + + + + {% endfor %} + +
AliasNameEngineActions
+ {% if connection.id %} + {{ connection.alias }} + {% else %} + {{ connection.alias }} + {% endif %} + {{ connection.name }}{{ connection.get_engine_display }} + Query + {% if connection.id %} + + + {% endif %} +
+
+
+{% endblock %} diff --git a/explorer/templates/connections/database_connection_confirm_delete.html b/explorer/templates/connections/database_connection_confirm_delete.html new file mode 100644 index 00000000..bcd221f2 --- /dev/null +++ b/explorer/templates/connections/database_connection_confirm_delete.html @@ -0,0 +1,12 @@ +{% extends 'explorer/base.html' %} +{% block sql_explorer_content %} +
+

Delete Database Connection

+

Are you sure you want to delete "{{ object }}"?

+
+ {% csrf_token %} + + Cancel +
+
+{% endblock %} diff --git a/explorer/templates/connections/database_connection_detail.html b/explorer/templates/connections/database_connection_detail.html new file mode 100644 index 00000000..162904ff --- /dev/null +++ b/explorer/templates/connections/database_connection_detail.html @@ -0,0 +1,39 @@ +{% extends "explorer/base.html" %} +{% block sql_explorer_content %} +
+

Connection Details

+ + + + + + + + + + + + + + {% if not object.is_upload %} + + + + + + + + + + + + + {% endif %} +
Alias{{ object.alias }}
Name{{ object.name }}
Engine{{ object.get_engine_display }}
User{{ object.user }}
Host{{ object.host }}
Port{{ object.port }}
+ {% if not object.is_upload %} + Edit + {% else %} + The source of this connection is an uploaded file. + {% endif %} +
+{% endblock %} diff --git a/explorer/templates/connections/database_connection_form.html b/explorer/templates/connections/database_connection_form.html new file mode 100644 index 00000000..cce67638 --- /dev/null +++ b/explorer/templates/connections/database_connection_form.html @@ -0,0 +1,23 @@ +{% extends 'explorer/base.html' %} +{% block sql_explorer_content %} +
+

{% if object %}Edit{% else %}Create New{% endif %} Connection

+
+ {% csrf_token %} + {{ form.as_p }} + + Cancel + +
+ {% if not object and s3_enabled %} +
+

...or upload a SQLite DB or CSV File

+

CSV files will get parsed and typed automatically. SQLite databases must not be password protected.

+
+

Drag and drop .csv or SQLite .db or click to upload.

+ +
+
+ {% endif %} +
+{% endblock %} diff --git a/explorer/templates/explorer/assistant.html b/explorer/templates/explorer/assistant.html index 077f28f5..dce8ed0e 100644 --- a/explorer/templates/explorer/assistant.html +++ b/explorer/templates/explorer/assistant.html @@ -21,10 +21,12 @@
+ (?) +
diff --git a/explorer/templates/explorer/base.html b/explorer/templates/explorer/base.html index c014767b..0e878608 100644 --- a/explorer/templates/explorer/base.html +++ b/explorer/templates/explorer/base.html @@ -48,16 +48,22 @@

This is easy to fix, I promise!

{% if can_change %}
  • {% trans "New Query" %} + href="{% url 'query_create' %}">{% trans "New Query" %}
  • {% trans "Playground" %} + href="{% url 'explorer_playground' %}">{% trans "Playground" %}
    • + {% if db_connections_enabled and can_manage_connections %} +
  • + {% trans "Connections" %} +
    • + {% endif %} {% endif %}
  • {% trans "Logs" %} + href="{% url 'explorer_logs' %}">{% trans "Logs" %}
  • - +
    diff --git a/explorer/templates/explorer/query_list.html b/explorer/templates/explorer/query_list.html index 79eb76f0..3f4d8979 100644 --- a/explorer/templates/explorer/query_list.html +++ b/explorer/templates/explorer/query_list.html @@ -7,7 +7,7 @@

    {% blocktrans trimmed with qlen=recent_queries|length %} - Your {{qlen}} Most Recently Run + Recently Run by You {% endblocktrans %}

    @@ -49,10 +49,13 @@

    {% trans "All Queries" %}

    - - + + {% if tasks_enabled %} - + {% endif %} {% if can_change %} @@ -60,8 +63,12 @@

    {% trans "All Queries" %}

    {% endif %} - - + + + @@ -76,10 +83,10 @@

    {% trans "All Queries" %}

    {% else %} - - {% endif %} - - + + {% endif %} {% endfor %} diff --git a/explorer/tests/csvs/all_types.csv b/explorer/tests/csvs/all_types.csv new file mode 100644 index 00000000..0f028f8a --- /dev/null +++ b/explorer/tests/csvs/all_types.csv @@ -0,0 +1,8 @@ +Dates,Integers,Floats,Strings +2020-01-31,0,42.952198961732414,THNVT +2020-02-29,1,27.66862453654746,JXPSY +2020-03-31,2,79.028965687494,FSTNN +2020-04-30,3,97.00288969016145,BVMNF +,,, +"","","" +2020-05-31,4,74.09328128351054,XUMUJ diff --git a/explorer/tests/csvs/dates.csv b/explorer/tests/csvs/dates.csv new file mode 100644 index 00000000..6df08bcd --- /dev/null +++ b/explorer/tests/csvs/dates.csv @@ -0,0 +1,13 @@ +Dates,Values +2024-01-24,0 +2024-01-24T18:45:00Z,1 +01/24/2024,2 +01/24/2024 6:45 PM,3 +24/01/2024,4 +24/01/2024 18:45,5 +24/01/2024 18:45:00,6 +"January 24, 2024",7 +"Jan 24, 2024",8 +2024-01-24T18:45:00-05:00,9 +"Thu, 24 Jan 2024 18:45:00 +0000",10 +"Thu, 24 Jan 2024 18:45:00 +0000 (GMT)",11 \ No newline at end of file diff --git a/explorer/tests/csvs/floats.csv b/explorer/tests/csvs/floats.csv new file mode 100644 index 00000000..cc06c5c4 --- /dev/null +++ b/explorer/tests/csvs/floats.csv @@ -0,0 +1,8 @@ +Floats,Values +61.3410231760637,0 +79.2367071213973,1 +96.93217099083482,2 +69.50523191870069,3 +,4 +"69.42719764194946",4 +"2,000.0128",5 diff --git a/explorer/tests/csvs/integers.csv b/explorer/tests/csvs/integers.csv new file mode 100644 index 00000000..384da892 --- /dev/null +++ b/explorer/tests/csvs/integers.csv @@ -0,0 +1,5 @@ +Integers,More_integers +"5,000",0 +"6000",1 +"",0 +"7,200,200",2 \ No newline at end of file diff --git a/explorer/tests/csvs/mixed.csv b/explorer/tests/csvs/mixed.csv new file mode 100644 index 00000000..eb1ee5e2 --- /dev/null +++ b/explorer/tests/csvs/mixed.csv @@ -0,0 +1,5 @@ +Value1,Value2,Value3 +2020-01-32,abc,123 +Variety of other dates,def,123 +,, +Another,123,12a \ No newline at end of file diff --git a/explorer/tests/csvs/test_case1.csv b/explorer/tests/csvs/test_case1.csv new file mode 100644 index 00000000..1d6dd6c8 --- /dev/null +++ b/explorer/tests/csvs/test_case1.csv @@ -0,0 +1,31 @@ +STORE,CONTENT_TYPE,EMAIL,CUSTOMER_ID,CREATED_DATE,SHIP_MONTH,SHIP_DATE,SHOPIFY_ORDER_NUMBER,SKU,SHOPIFY_PRODUCT_ID,SHOPIFY_VARIANT_ID,PRODUCT_TITLE,VARIANT_TITLE,PRICE,QUANTITY,SUB_TOTAL_PRICE,SHIPPING_FEE,TAX,DISCOUNT,SHOPIFY_ORDER_ID,GIFTCARD_AMOUNT,REFUNDS,REFUND_TAX,REFUND_SALE,REFUND_DATE,SHIPPING_ADDRESS_PROVINCE_CODE,ORDER_TAG,SHIPPING_ADDRESS_FIRST_NAME,SHIPPING_ADDRESS_LAST_NAME,SHIPPING_ADDRESS_ADDRESS_1,SHIPPING_ADDRESS_ADDRESS_2,SHIPPING_ADDRESS_CITY,SHIPPING_ADDRESS_PROVINCE,SHIPPING_ADDRESS_COUNTRY,SHIPPING_ADDRESS_ZIP,ORDER_TAGS, +556516,HBUS,Vitamin,fdf@yahoo.com,3887814443067,2023-03-31T00:00:00.000Z,2023-03-01T00:00:00.000Z,2023-03-31T00:00:00.000Z,3948556,HBVITAkidsleep75,4907580457019,33723289010235,(Vitamin Bundle) Kid's Sleep,Default Title,0,1,0,0,0,0,4857684885563,0,0,0,0,,FL,,Clark,4075493288,Petunia Terrace,113,De Pere,Wisconsin,United States,32771,"Subscription Recurring Order, sent-to-d365, Subscription, FRAUD_APPROVED" +556517,HBUS,Vitamin,dsfs@yahoo.com,3887814443067,2023-03-31T00:00:00.000Z,2023-03-01T00:00:00.000Z,2023-03-31T00:00:00.000Z,3948556,HBVITAmenim60,4331305107515,31085422018619,(Vitamin Bundle) Organic Men's Multi + Super Immune,Default Title,0,1,0,0,0,0,4857684885563,0,0,0,0,,FL,,Winter,4075493288,20 Union St,113,Warren,Michigan,United States,32771,"Subscription Recurring Order, sent-to-d365, Subscription, FRAUD_APPROVED" +556518,HBUS,Vitamin,fdsdf@comcast.net,2444580454459,2023-03-31T00:00:00.000Z,2023-03-01T00:00:00.000Z,2023-03-31T00:00:00.000Z,3946315,HBVITAkidsleep75,4907580457019,33723289010235,(Vitamin Bundle) Kid's Sleep,Default Title,0,3,0,0,0,0,4857557188667,0,0,0,0,,PA,,Berns,2532085200,Eisenhower Dr,,Port Jefferson,New York,United States,15065,"Subscription, Subscription Recurring Order, FRAUD_APPROVED, sent-to-d365" +556519,HBUS,Vitamin,fdsdf@yahoo.com,3887814443067,2023-03-31T00:00:00.000Z,2023-03-01T00:00:00.000Z,2023-03-31T00:00:00.000Z,3948556,HBVITAHSNCol60NV,6830033633339,40285923606587,"(Vitamin Bundle) Hair, Skin, Nails + Collagen",Default Title,0,1,0,0,0,0,4857684885563,0,0,0,0,,FL,,Gantt,4075493288,227 Crooked Oak Rd,113,Port Jefferson,New York,United States,32771,"Subscription Recurring Order, sent-to-d365, Subscription, FRAUD_APPROVED" +556520,HBUS,Vitamin,asd@comcast.net,2444580454459,2023-03-31T00:00:00.000Z,2023-03-01T00:00:00.000Z,2023-03-31T00:00:00.000Z,3946315,HBVitaBundle7,4331298521147,31085392724027,Vitamin Bundle,7,65.99,1,65.99,4.99,0,13.2,4857557188667,0,0,0,0,,PA,,Johnson,2532085200,227 Eisenhower Dr,,Salt Lake City,Utah,United States,15065,"Subscription, Subscription Recurring Order, FRAUD_APPROVED, sent-to-d365" +556521,HBUS,Vitamin,fda@comcast.net,2444580454459,2023-03-31T00:00:00.000Z,2023-03-01T00:00:00.000Z,2023-03-31T00:00:00.000Z,3946315,HBVITAkid60,4331303469115,31085415006267,(Vitamin Bundle) Organic Kid's Multi,60-count,0,2,0,0,0,0,4857557188667,0,0,0,0,,PA,,Moody,2532085200,Eisenhower Dr,,De Pere,Wisconsin,United States,15065,"Subscription, Subscription Recurring Order, FRAUD_APPROVED, sent-to-d365" +556522,HBUS,Vitamin,fdsa.family@gmail.com,5976811143227,2023-03-31T00:00:00.000Z,2023-03-01T00:00:00.000Z,2023-03-31T00:00:00.000Z,3948458,HBVITAimm75,4331307499579,31085430014011,(Vitamin Bundle) Immunity,Default Title,0,1,0,0,0,0,4857645498427,0,0,0,0,,MI,,Clark,6164811698,Petunia Terrace,Nw,De Pere,Wisconsin,United States,49544,"FRAUD_APPROVED, Subscription, Subscription Recurring Order, sent-to-d365" +556523,HBUS,Vitamin,asdf.ahlf@gmail.com,5765744361531,2023-03-31T00:00:00.000Z,2023-03-01T00:00:00.000Z,2023-03-31T00:00:00.000Z,3946630,HBVITAwom60,4331304681531,31085420281915,(Vitamin Bundle) Organic Women's Multi,Default Title,0,1,0,0,0,0,4857570000955,0,0,0,0,,IL,,Winter,3863347803,20 Union St,,Warren,Wisconsin,United States,60657,"FRAUD_APPROVED, Subscription, sent-to-d365, Subscription Recurring Order" +556524,HBUS,Vitamin,asdf@yahoo.com,3887814443067,2023-03-31T00:00:00.000Z,2023-03-01T00:00:00.000Z,2023-03-31T00:00:00.000Z,3948556,HBVITAadultstr,4907584782395,33723310178363,(Vitamin Bundle) Adult Stress,Default Title,0,1,0,0,0,0,4857684885563,0,0,0,0,,FL,,Berns,4075493288,Eisenhower Dr,113,Port Jefferson,Michigan,United States,32771,"Subscription Recurring Order, sent-to-d365, Subscription, FRAUD_APPROVED" +556525,HBUS,Vitamin,dfgh5@yahoo.com,6074075873339,2023-03-30T00:00:00.000Z,2023-03-01T00:00:00.000Z,2023-03-31T00:00:00.000Z,3945418,HBVitaBundle5,4331298521147,31085392691259,Vitamin Bundle,5,49.99,1,49.99,4.99,1.92,15,4856756699195,0,0,0,0,,WI,,Gantt,9204717027,227 Crooked Oak Rd,,Port Jefferson,New York,United States,54115,"FRAUD_APPROVED, Subscription First Order, Subscription, sent-to-d365" +556526,HBUS,Vitamin,gdfg5@gmail.com,5845340389435,2023-03-31T00:00:00.000Z,2023-03-01T00:00:00.000Z,2023-03-31T00:00:00.000Z,3946834,HBVITAHSNC60,7110465355835,41164903678011,"(Vitamin Bundle) Hair, Skin + Nails",Default Title,0,1,0,0,0,0,4857578029115,0,0,0,0,,MI,,Johnson,3134422087,227 Eisenhower Dr,,Salt Lake City,New York,United States,48091,"Subscription Recurring Order, FRAUD_APPROVED, Subscription, sent-to-d365" +556527,HBUS,Vitamin,ghgh@gf.com,6086223102011,2023-03-31T00:00:00.000Z,2023-03-01T00:00:00.000Z,2023-03-31T00:00:00.000Z,3946172,HBVITAPRENATDHA60,4331306811451,31085428015163,(Vitamin Bundle) Prenatal + DHA,Default Title,0,3,0,0,0,0,4857464619067,0,0,0,0,,NY,,Moody,4153162973,Eisenhower Dr,,De Pere,Utah,United States,11777,"sent-to-d365, Subscription First Order, Subscription, FRAUD_APPROVED" +556528,HBUS,Vitamin,sdfg5@fg.com,6086223102011,2023-03-31T00:00:00.000Z,2023-03-01T00:00:00.000Z,2023-03-31T00:00:00.000Z,3946172,HBVitaBundle3,4331298521147,31085392658491,Vitamin Bundle,3,32.99,1,32.99,9.99,0,0,4857464619067,0,0,0,0,,NY,,Clark,4153162973,Petunia Terrace,,De Pere,Wisconsin,United States,11777,"sent-to-d365, Subscription First Order, Subscription, FRAUD_APPROVED" +556529,HBUS,Vitamin,sdfgs@gmail.com,5815550279739,2023-03-30T00:00:00.000Z,2023-03-01T00:00:00.000Z,2023-03-31T00:00:00.000Z,3945751,HBVITAwom60,4331304681531,31085420281915,(Vitamin Bundle) Organic Women's Multi,Default Title,0,1,0,0,0,0,4857046335547,0,0,0,0,,UT,,Winter,8016081564,20 Union St,635,Warren,Wisconsin,United States,84108,"sent-to-d365, FRAUD_APPROVED, Subscription Recurring Order, Subscription" +556530,HBUS,Vitamin,sdfb@yahoo.com,6074075873339,2023-03-30T00:00:00.000Z,2023-03-01T00:00:00.000Z,2023-03-31T00:00:00.000Z,3945418,HBVITAwom60,4331304681531,31085420281915,(Vitamin Bundle) Organic Women's Multi,Default Title,0,1,0,0,0,0,4856756699195,0,0,0,0,,WI,,Berns,9204717027,Eisenhower Dr,,Port Jefferson,Wisconsin,United States,54115,"FRAUD_APPROVED, Subscription First Order, Subscription, sent-to-d365" +556531,HBUS,Vitamin,brookesdfb4colwell@yahoo.com,6074075873339,2023-03-30T00:00:00.000Z,2023-03-01T00:00:00.000Z,2023-03-31T00:00:00.000Z,3945418,HBVITAKIDMOODB60,7012437327931,40867363422267,(Vitamin Bundle) Kid's Mood Boost,Default Title,0,2,0,0,0,0,4856756699195,0,0,0,0,,WI,,Gantt,9204717027,227 Crooked Oak Rd,,Port Jefferson,Michigan,United States,54115,"FRAUD_APPROVED, Subscription First Order, Subscription, sent-to-d365" +556532,HBUS,Vitamin,34g@gmail.com,5747147309115,2023-03-31T00:00:00.000Z,2023-03-01T00:00:00.000Z,2023-03-31T00:00:00.000Z,3947011,HBVITAwom60,4331304681531,31085420281915,(Vitamin Bundle) Organic Women's Multi,Default Title,0,2,0,0,0,0,4857585467451,0,0,0,0,,SC,,Johnson,6313126335,227 Eisenhower Dr,,Salt Lake City,New York,United States,29455,"Subscription Recurring Order, sent-to-d365, Subscription, FRAUD_APPROVED" +556533,HBUS,Vitamin,dbfd@yahoo.com,6074075873339,2023-03-30T00:00:00.000Z,2023-03-01T00:00:00.000Z,2023-03-31T00:00:00.000Z,3945418,HBVITAkid60,4331303469115,31085415006267,(Vitamin Bundle) Organic Kid's Multi,60-count,0,1,0,0,0,0,4856756699195,0,0,0,0,,WI,,Moody,9204717027,Eisenhower Dr,,De Pere,New York,United States,54115,"FRAUD_APPROVED, Subscription First Order, Subscription, sent-to-d365" +556534,HBUS,Vitamin,sdbsdf@gmail.com,5815550279739,2023-03-30T00:00:00.000Z,2023-03-01T00:00:00.000Z,2023-03-31T00:00:00.000Z,3945751,HBVitaBundle3,4331298521147,31085392658491,Vitamin Bundle,3,32.99,1,32.99,9.99,0.99,0,4857046335547,0,0,0,0,,UT,,Clark,8016081564,Petunia Terrace,635,De Pere,Utah,United States,84108,"sent-to-d365, FRAUD_APPROVED, Subscription Recurring Order, Subscription" +556535,HBUS,Vitamin,sdfg3@yahoo.com,6074075873339,2023-03-30T00:00:00.000Z,2023-03-01T00:00:00.000Z,2023-03-31T00:00:00.000Z,3945418,HBVITAimm75,4331307499579,31085430014011,(Vitamin Bundle) Immunity,Default Title,0,1,0,0,0,0,4856756699195,0,0,0,0,,WI,,Winter,9204717027,20 Union St,,Warren,Wisconsin,United States,54115,"FRAUD_APPROVED, Subscription First Order, Subscription, sent-to-d365" +556536,HBUS,Vitamin,sfg@gmail.com,3058267258939,2023-03-31T00:00:00.000Z,2023-03-01T00:00:00.000Z,2023-03-31T00:00:00.000Z,3946977,HBVITAsle75,4331307204667,31085429358651,(Vitamin Bundle) Sleep Well,Default Title,0,2,0,0,0,0,4857584156731,0,0,0,0,,PA,,Berns,2012596182,Eisenhower Dr,,Port Jefferson,Wisconsin,United States,18414,"Subscription, FRAUD_APPROVED, sent-to-d365, Subscription Recurring Order" +556537,HBUS,Vitamin,sdfg@gmail.com,2971195867195,2023-03-30T00:00:00.000Z,2023-03-01T00:00:00.000Z,2023-03-31T00:00:00.000Z,3945648,HBVITAHSNCol60NV,6830033633339,40285923606587,"(Vitamin Bundle) Hair, Skin, Nails + Collagen",Default Title,0,2,0,0,0,0,4856948785211,0,0,0,0,,CA,,Gantt,4044097130,227 Crooked Oak Rd,,Port Jefferson,Wisconsin,United States,91505,"FRAUD_APPROVED, Subscription Recurring Order, sent-to-d365, Subscription" +556538,HBUS,Vitamin,vfsw@gmail.com,2971195867195,2023-03-30T00:00:00.000Z,2023-03-01T00:00:00.000Z,2023-03-31T00:00:00.000Z,3945648,HBVITACHL60,6841717850171,40321076789307,(Vitamin Bundle) Organic Chlorophyll,Default Title,0,1,0,0,0,0,4856948785211,0,0,0,0,,CA,,Johnson,4044097130,227 Eisenhower Dr,,Salt Lake City,Michigan,United States,91505,"FRAUD_APPROVED, Subscription Recurring Order, sent-to-d365, Subscription" +556539,HBUS,Vitamin,sdfgg@gmail.com,2971195867195,2023-03-30T00:00:00.000Z,2023-03-01T00:00:00.000Z,2023-03-31T00:00:00.000Z,3945648,HBVITAimm75,4331307499579,31085430014011,(Vitamin Bundle) Immunity,Default Title,0,2,0,0,0,0,4856948785211,0,0,0,0,,CA,,Moody,4044097130,Eisenhower Dr,,De Pere,New York,United States,91505,"FRAUD_APPROVED, Subscription Recurring Order, sent-to-d365, Subscription" +556540,HBUS,Vitamin,sdfgbv@gmail.com,5810735513659,2023-03-31T00:00:00.000Z,2023-03-01T00:00:00.000Z,2023-03-31T00:00:00.000Z,3947165,HBVitaBundle3,4331298521147,31085392658491,Vitamin Bundle,3,32.99,1,32.99,9.99,0,0,4857591169083,0,0,0,0,,IL,,Clark,8159092206,Petunia Terrace,,De Pere,New York,United States,61108,"Subscription Recurring Order, Subscription, FRAUD_APPROVED, sent-to-d365" +556541,HBUS,Vitamin,sdvf3@gmail.com,3379719831611,2023-03-31T00:00:00.000Z,2023-03-01T00:00:00.000Z,2023-03-31T00:00:00.000Z,3947396,HBVitaBundle3,4331298521147,31085392658491,Vitamin Bundle,3,32.99,1,32.99,9.99,0,0,4857600639035,0,0,0,0,,NH,,Winter,5174491774,20 Union St,,Warren,Utah,United States,3833,"FRAUD_APPROVED, Subscription Recurring Order, sent-to-d365, Subscription" +556542,HBUS,Vitamin,34f@gmail.com,3379719831611,2023-03-31T00:00:00.000Z,2023-03-01T00:00:00.000Z,2023-03-31T00:00:00.000Z,3947396,HBVITAHSNCol60NV,6830033633339,40285923606587,"(Vitamin Bundle) Hair, Skin, Nails + Collagen",Default Title,0,1,0,0,0,0,4857600639035,0,0,0,0,,NH,,Berns,5174491774,Eisenhower Dr,,Port Jefferson,Wisconsin,United States,3833,"FRAUD_APPROVED, Subscription Recurring Order, sent-to-d365, Subscription" +556543,HBUS,Vitamin,dsfv@gmail.com,3379719831611,2023-03-31T00:00:00.000Z,2023-03-01T00:00:00.000Z,2023-03-31T00:00:00.000Z,3947396,HBVITAmenim60,4331305107515,31085422018619,(Vitamin Bundle) Organic Men's Multi + Super Immune,Default Title,0,1,0,0,0,0,4857600639035,0,0,0,0,,NH,,Gantt,5174491774,227 Crooked Oak Rd,,Port Jefferson,Wisconsin,United States,3833,"FRAUD_APPROVED, Subscription Recurring Order, sent-to-d365, Subscription" +556544,HBUS,Vitamin,dsfv35@gmail.com,3379719831611,2023-03-31T00:00:00.000Z,2023-03-01T00:00:00.000Z,2023-03-31T00:00:00.000Z,3947396,HBVITAPRENATDHA60,4331306811451,31085428015163,(Vitamin Bundle) Prenatal + DHA,Default Title,0,1,0,0,0,0,4857600639035,0,0,0,0,,NH,,Johnson,5174491774,227 Eisenhower Dr,,Salt Lake City,Wisconsin,United States,3833,"FRAUD_APPROVED, Subscription Recurring Order, sent-to-d365, Subscription" +556545,HBUS,Vitamin,sdfv@gmail.com,5747147309115,2023-03-31T00:00:00.000Z,2023-03-01T00:00:00.000Z,2023-03-31T00:00:00.000Z,3947011,HBVITAACVGW,6785724022843,40072796504123,(Vitamin Bundle) Apple Cider Vinegar,Default Title,0,1,0,0,0,0,4857585467451,0,0,0,0,,SC,,Moody,6313126335,Eisenhower Dr,,De Pere,Michigan,United States,29455,"Subscription Recurring Order, sent-to-d365, Subscription, FRAUD_APPROVED" \ No newline at end of file diff --git a/explorer/tests/settings.py b/explorer/tests/settings.py index d956466d..206839ac 100644 --- a/explorer/tests/settings.py +++ b/explorer/tests/settings.py @@ -5,13 +5,14 @@ EXPLORER_AI_API_KEY = None # set to any value to enable assistant CELERY_BROKER_URL = "redis://localhost:6379/0" CELERY_TASK_ALWAYS_EAGER = True +TEST_MODE = True DATABASES = { "default": { "ENGINE": "django.db.backends.sqlite3", - "NAME": "tst", + "NAME": "tst1", "TEST": { - "NAME": "tst" + "NAME": "tst1" } }, "alt": { diff --git a/explorer/tests/test_assistant.py b/explorer/tests/test_assistant.py new file mode 100644 index 00000000..7ac60f7c --- /dev/null +++ b/explorer/tests/test_assistant.py @@ -0,0 +1,283 @@ +from explorer.tests.factories import SimpleQueryFactory, QueryLogFactory +from unittest.mock import patch, Mock, MagicMock + +import json +from django.test import TestCase +from django.urls import reverse +from django.contrib.auth.models import User +from django.db import OperationalError +from explorer.app_settings import EXPLORER_DEFAULT_CONNECTION as CONN +from explorer.assistant.utils import sample_rows_from_table, ROW_SAMPLE_SIZE, build_prompt + + +class TestAssistantViews(TestCase): + + def setUp(self): + self.user = User.objects.create_superuser( + "admin", "admin@admin.com", "pwd" + ) + self.client.login(username="admin", password="pwd") + self.request_data = { + "sql": "SELECT * FROM explorer_query", + "connection": CONN, + "assistant_request": "Test Request" + } + + @patch("explorer.assistant.utils.openai_client") + def test_do_modify_query(self, mocked_openai_client): + from explorer.assistant.views import run_assistant + + # create.return_value should match: resp.choices[0].message + mocked_openai_client.return_value.chat.completions.create.return_value = Mock( + choices=[Mock(message=Mock(content="smart computer"))]) + resp = run_assistant(self.request_data, None) + self.assertEqual(resp, "smart computer") + + @patch("explorer.assistant.utils.openai_client") + def test_assistant_help(self, mocked_openai_client): + mocked_openai_client.return_value.chat.completions.create.return_value = Mock( + choices=[Mock(message=Mock(content="smart computer"))]) + resp = self.client.post(reverse("assistant"), + data=json.dumps(self.request_data), + content_type="application/json") + self.assertEqual(json.loads(resp.content)["message"], "smart computer") + + +class TestBuildPrompt(TestCase): + + @patch("explorer.assistant.utils.sample_rows_from_tables", return_value="sample data") + @patch("explorer.assistant.utils.fits_in_window", return_value=True) + @patch("explorer.models.ExplorerValue.objects.get_item") + def test_build_prompt_with_vendor_only(self, mock_get_item, mock_fits_in_window, mock_sample_rows): + mock_get_item.return_value.value = "system prompt" + + request_data = { + "connection": "default", + "assistant_request": "Help me with SQL" + } + included_tables = [] + + result = build_prompt(request_data, included_tables) + self.assertIn("## Database Vendor / SQL Flavor is sqlite", result["user"]) + self.assertIn("## User's Request to Assistant ##\n\nHelp me with SQL\n\n", result["user"]) + self.assertEqual(result["system"], "system prompt") + + @patch("explorer.assistant.utils.sample_rows_from_tables", return_value="sample data") + @patch("explorer.assistant.utils.fits_in_window", return_value=True) + @patch("explorer.models.ExplorerValue.objects.get_item") + def test_build_prompt_with_sql(self, mock_get_item, mock_fits_in_window, mock_sample_rows): + mock_get_item.return_value.value = "system prompt" + + request_data = { + "connection": "default", + "sql": "SELECT * FROM table;", + "assistant_request": "Help me with SQL" + } + included_tables = [] + + result = build_prompt(request_data, included_tables) + self.assertIn("## Database Vendor / SQL Flavor is sqlite", result["user"]) + self.assertIn("## Existing SQL ##\n\nSELECT * FROM table;\n\n", result["user"]) + self.assertIn("## User's Request to Assistant ##\n\nHelp me with SQL\n\n", result["user"]) + self.assertEqual(result["system"], "system prompt") + + @patch("explorer.assistant.utils.sample_rows_from_tables", return_value="sample data") + @patch("explorer.assistant.utils.fits_in_window", return_value=True) + @patch("explorer.models.ExplorerValue.objects.get_item") + def test_build_prompt_with_sql_and_error(self, mock_get_item, mock_fits_in_window, mock_sample_rows): + mock_get_item.return_value.value = "system prompt" + + request_data = { + "connection": "default", + "sql": "SELECT * FROM table;", + "db_error": "Syntax error", + "assistant_request": "Help me with SQL" + } + included_tables = [] + + result = build_prompt(request_data, included_tables) + self.assertIn("## Database Vendor / SQL Flavor is sqlite", result["user"]) + self.assertIn("## Existing SQL ##\n\nSELECT * FROM table;\n\n", result["user"]) + self.assertIn("## Query Error ##\n\nSyntax error\n\n", result["user"]) + self.assertIn("## User's Request to Assistant ##\n\nHelp me with SQL\n\n", result["user"]) + self.assertEqual(result["system"], "system prompt") + + @patch("explorer.assistant.utils.sample_rows_from_tables", return_value="sample data") + @patch("explorer.assistant.utils.fits_in_window", return_value=True) + @patch("explorer.models.ExplorerValue.objects.get_item") + def test_build_prompt_with_extra_tables_fitting_window(self, mock_get_item, mock_fits_in_window, mock_sample_rows): + mock_get_item.return_value.value = "system prompt" + + request_data = { + "connection": "default", + "sql": "SELECT * FROM table;", + "assistant_request": "Help me with SQL" + } + included_tables = ["table1", "table2"] + + result = build_prompt(request_data, included_tables) + self.assertIn("## Database Vendor / SQL Flavor is sqlite", result["user"]) + self.assertIn("## Existing SQL ##\n\nSELECT * FROM table;\n\n", result["user"]) + self.assertIn("## Table Structure with Sampled Data ##\n\nsample data\n\n", result["user"]) + self.assertIn("## User's Request to Assistant ##\n\nHelp me with SQL\n\n", result["user"]) + self.assertEqual(result["system"], "system prompt") + + @patch("explorer.assistant.utils.sample_rows_from_tables", return_value="sample data") + @patch("explorer.assistant.utils.fits_in_window", return_value=False) + @patch("explorer.assistant.utils.tables_from_schema_info", return_value="table structure") + @patch("explorer.models.ExplorerValue.objects.get_item") + def test_build_prompt_with_extra_tables_not_fitting_window(self, mock_get_item, mock_tables_from_schema_info, + mock_fits_in_window, mock_sample_rows): + mock_get_item.return_value.value = "system prompt" + + request_data = { + "connection": "default", + "sql": "SELECT * FROM table;", + "assistant_request": "Help me with SQL" + } + included_tables = ["table1", "table2"] + + result = build_prompt(request_data, included_tables) + self.assertIn("## Database Vendor / SQL Flavor is sqlite", result["user"]) + self.assertIn("## Existing SQL ##\n\nSELECT * FROM table;\n\n", result["user"]) + self.assertIn("## Table Structure ##\n\ntable structure\n\n", result["user"]) + self.assertIn("## User's Request to Assistant ##\n\nHelp me with SQL\n\n", result["user"]) + self.assertEqual(result["system"], "system prompt") + + +class TestPromptContext(TestCase): + + def test_retrieves_sample_rows(self): + SimpleQueryFactory(title="First Query") + SimpleQueryFactory(title="Second Query") + SimpleQueryFactory(title="Third Query") + SimpleQueryFactory(title="Fourth Query") + ret = sample_rows_from_table(CONN, "explorer_query") + self.assertEqual(len(ret), ROW_SAMPLE_SIZE+1) # includes header row + + @patch("explorer.assistant.utils.get_valid_connection") + def test_truncates_long_strings(self, mock_get_valid_connection): + # Mock database connection and cursor + mock_conn = MagicMock() + mock_cursor = MagicMock() + mock_get_valid_connection.return_value = mock_conn + mock_conn.cursor.return_value = mock_cursor + + long_string = "a" * 600 + mock_cursor.description = [("col1",), ("col2",)] + mock_cursor.fetchall.return_value = [(long_string, "short string")] + + ret = sample_rows_from_table(CONN, "some_table") + header, row = ret + + self.assertEqual(header, ["col1", "col2"]) + self.assertEqual(row[0], "a" * 500 + "...") + self.assertEqual(row[1], "short string") + + @patch("explorer.assistant.utils.get_valid_connection") + def test_truncates_long_binary_data(self, mock_get_valid_connection): + # Mock database connection and cursor + mock_conn = MagicMock() + mock_cursor = MagicMock() + mock_get_valid_connection.return_value = mock_conn + mock_conn.cursor.return_value = mock_cursor + + long_binary = b"a" * 600 + mock_cursor.description = [("col1",), ("col2",)] + mock_cursor.fetchall.return_value = [(long_binary, b"short binary")] + + ret = sample_rows_from_table(CONN, "some_table") + header, row = ret + + self.assertEqual(header, ["col1", "col2"]) + self.assertEqual(row[0], b"a" * 500 + b"...") + self.assertEqual(row[1], b"short binary") + + @patch("explorer.assistant.utils.get_valid_connection") + def test_handles_various_data_types(self, mock_get_valid_connection): + # Mock database connection and cursor + mock_conn = MagicMock() + mock_cursor = MagicMock() + mock_get_valid_connection.return_value = mock_conn + mock_conn.cursor.return_value = mock_cursor + + mock_cursor.description = [("col1",), ("col2",), ("col3",)] + mock_cursor.fetchall.return_value = [(123, 45.67, "normal string")] + + ret = sample_rows_from_table(CONN, "some_table") + header, row = ret + + self.assertEqual(header, ["col1", "col2", "col3"]) + self.assertEqual(row[0], 123) + self.assertEqual(row[1], 45.67) + self.assertEqual(row[2], "normal string") + + @patch("explorer.assistant.utils.get_valid_connection") + def test_handles_operational_error(self, mock_get_valid_connection): + # Mock database connection and cursor + mock_conn = MagicMock() + mock_cursor = MagicMock() + mock_get_valid_connection.return_value = mock_conn + mock_conn.cursor.return_value = mock_cursor + + mock_cursor.execute.side_effect = OperationalError("Test OperationalError") + + ret = sample_rows_from_table(CONN, "some_table") + + self.assertEqual(ret, [["Test OperationalError"]]) + + def test_format_rows_from_table(self): + from explorer.assistant.utils import format_rows_from_table + d = [ + ["col1", "col2"], + ["val1", "val2"], + ] + ret = format_rows_from_table(d) + self.assertEqual(ret, "col1 | col2\n" + "-" * 50 + "\nval1 | val2\n") + + def test_parsing_tables_from_query(self): + from explorer.assistant.utils import get_table_names_from_query + sql = "SELECT * FROM explorer_query" + ret = get_table_names_from_query(sql) + self.assertEqual(ret, ["explorer_query"]) + + def test_parsing_tables_from_query_bad_sql(self): + from explorer.assistant.utils import get_table_names_from_query + sql = "foo" + ret = get_table_names_from_query(sql) + self.assertEqual(ret, []) + + def test_schema_info_from_table_names(self): + from explorer.assistant.utils import tables_from_schema_info + ret = tables_from_schema_info(CONN, ["explorer_query"]) + expected = [("explorer_query", [ + ("id", "AutoField"), + ("title", "CharField"), + ("sql", "TextField"), + ("description", "TextField"), + ("created_at", "DateTimeField"), + ("last_run_date", "DateTimeField"), + ("created_by_user_id", "IntegerField"), + ("snapshot", "BooleanField"), + ("connection", "CharField")])] + self.assertEqual(ret, expected) + + +class TestAssistantUtils(TestCase): + + def test_sample_rows_from_tables(self): + from explorer.assistant.utils import sample_rows_from_tables + SimpleQueryFactory(title="First Query") + SimpleQueryFactory(title="Second Query") + QueryLogFactory() + ret = sample_rows_from_tables(CONN, ["explorer_query", "explorer_querylog"]) + self.assertTrue("First Query" in ret) + self.assertTrue("Second Query" in ret) + self.assertTrue("explorer_querylog" in ret) + + def test_sample_rows_from_tables_no_tables(self): + from explorer.assistant.utils import sample_rows_from_tables + SimpleQueryFactory(title="First Query") + SimpleQueryFactory(title="Second Query") + ret = sample_rows_from_tables(CONN, []) + self.assertEqual(ret, "") diff --git a/explorer/tests/test_connections.py b/explorer/tests/test_connections.py new file mode 100644 index 00000000..06b702a2 --- /dev/null +++ b/explorer/tests/test_connections.py @@ -0,0 +1,71 @@ +from importlib import reload +from django.test import TestCase +from django.db import transaction +from unittest.mock import patch +from explorer.models import DatabaseConnection +from explorer import connections + + +class ExplorerConnectionsTests(TestCase): + + def test_a_new_get_connection_from_user_created(self): + # Create a proper DatabaseConnection object + DatabaseConnection.objects.create( + alias="user_db", + name="User Database", + engine="django.db.backends.sqlite3", + host="", + port="", + user="", + password="", + ) + + conn = connections.new_get_connection("user_db") + self.assertEqual(conn.settings_dict["NAME"], "User Database") + DatabaseConnection.objects.all().delete() + + def test_monkey_patch_transaction_get_connection(self): + original_get_connection = transaction.get_connection + + reload(connections) # Reload the module to apply the patch + self.assertEqual(transaction.get_connection, connections.new_get_connection) + + # Restore original state + transaction.get_connection = original_get_connection + + +class ConnectionsBlendingTests(TestCase): + + def setUp(self): + DatabaseConnection.objects.all().delete() + + @patch("explorer.connections.app_settings") + def test_connections_blended_when_enabled(self, mocked_app_settings): + mocked_app_settings.EXPLORER_CONNECTIONS = {"default": "default", "alt": "alt"} + DatabaseConnection.objects.create( + alias="user_db1", + name="User Database", + engine="django.db.backends.sqlite3", + host="", + port="", + user="", + password="", + ) + + DatabaseConnection.objects.create( + alias="user_db2", + name="User Database", + engine="django.db.backends.sqlite3", + host="", + port="", + user="", + password="", + ) + + conn_dict = connections.connections() + self.assertIn("default", conn_dict) + self.assertIn("alt", conn_dict) + self.assertIn("user_db1", conn_dict) + self.assertIn("user_db2", conn_dict) + self.assertEqual(len(conn_dict), 4) + diff --git a/explorer/tests/test_models.py b/explorer/tests/test_models.py index b2ffebb7..3569ba5a 100644 --- a/explorer/tests/test_models.py +++ b/explorer/tests/test_models.py @@ -7,7 +7,7 @@ from explorer import app_settings from explorer.app_settings import EXPLORER_DEFAULT_CONNECTION as CONN -from explorer.models import ColumnHeader, ColumnSummary, Query, QueryLog, QueryResult +from explorer.models import ColumnHeader, ColumnSummary, Query, QueryLog, QueryResult, DatabaseConnection from explorer.tests.factories import SimpleQueryFactory @@ -220,3 +220,19 @@ def test_handles_null_as_zero(self): def test_empty_data(self): res = ColumnSummary("foo", []) self.assertEqual(res.stats, {"Min": 0, "Max": 0, "Avg": 0, "Sum": 0, "NUL": 0}) + + +class TestDatabaseConnection(TestCase): + + def test_cant_create_a_connection_with_conflicting_name(self): + thrown = False + try: + DatabaseConnection.objects.create(name="default") + except ValidationError: + thrown = True + self.assertTrue(thrown) + + def test_create_db_connection_from_django_connection(self): + c = DatabaseConnection.from_django_connection(app_settings.EXPLORER_DEFAULT_CONNECTION) + self.assertEqual(c.name, "tst1") + self.assertEqual(c.alias, "default") diff --git a/explorer/tests/test_tasks.py b/explorer/tests/test_tasks.py index 9fb50328..e0b56b97 100644 --- a/explorer/tests/test_tasks.py +++ b/explorer/tests/test_tasks.py @@ -1,7 +1,9 @@ import unittest from datetime import datetime, timedelta from io import StringIO -from unittest.mock import patch +from unittest.mock import patch, MagicMock +import os +import tempfile from django.core import mail from django.test import TestCase @@ -9,7 +11,7 @@ from explorer import app_settings from explorer.app_settings import EXPLORER_DEFAULT_CONNECTION as CONN from explorer.models import QueryLog -from explorer.tasks import build_schema_cache_async, execute_query, snapshot_queries, truncate_querylogs +from explorer.tasks import build_schema_cache_async, execute_query, snapshot_queries, truncate_querylogs, remove_unused_sqlite_dbs from explorer.tests.factories import SimpleQueryFactory @@ -89,3 +91,105 @@ def test_build_schema_cache_async(self, mocked_build): schema = build_schema_cache_async(CONN) assert mocked_build.called self.assertEqual(schema, [("table", [("column", "Integer")]),]) + + +class RemoveUnusedSQLiteDBsTestCase(TestCase): + + @patch("explorer.tasks.QueryLog") + @patch("explorer.tasks.DatabaseConnection") + def test_remove_unused_sqlite_dbs(self, MockDatabaseConnection, MockQueryLog): + # Mock the settings + days_inactivity = 30 + with patch("explorer.tasks.app_settings.EXPLORER_PRUNE_LOCAL_UPLOAD_COPY_DAYS_INACTIVITY", days_inactivity): + with tempfile.TemporaryDirectory() as temp_dir: + + # Create a temporary SQLite file + temp_db_path = os.path.join(temp_dir, "test_db1.sqlite") + + with open(temp_db_path, "w") as temp_db: + temp_db.write("") + + # Mock DatabaseConnection objects + db1 = MagicMock() + db1.local_name = temp_db_path + db1.alias = "test_alias1" + db1.host = "some_host" + + MockDatabaseConnection.objects.filter.return_value = [db1] + + recent_time = datetime.now() - timedelta(days=days_inactivity + 1) + query_log1 = MagicMock() + query_log1.run_at = recent_time + + mock_queryset = MagicMock() + mock_queryset.first.return_value = query_log1 + + def filter_side_effect(*args, **kwargs): + if kwargs.get("connection") == "test_alias1": + return mock_queryset + return MagicMock(first=MagicMock(return_value=None)) + + MockQueryLog.objects.filter.side_effect = filter_side_effect + + # Call the function + remove_unused_sqlite_dbs() + + # Assertions + self.assertFalse(os.path.exists(temp_db_path)) + + @patch("explorer.tasks.QueryLog") + @patch("explorer.tasks.DatabaseConnection") + def test_do_not_remove_recently_used_db(self, MockDatabaseConnection, MockQueryLog): + days_inactivity = 30 + with patch("explorer.tasks.app_settings.EXPLORER_PRUNE_LOCAL_UPLOAD_COPY_DAYS_INACTIVITY", days_inactivity): + with tempfile.TemporaryDirectory() as temp_dir: + + temp_db_path = os.path.join(temp_dir, "test_db1.sqlite") + + with open(temp_db_path, "w") as temp_db: + temp_db.write("") + + # Mock DatabaseConnection objects + db1 = MagicMock() + db1.local_name = temp_db_path + db1.alias = "test_alias1" + db1.host = "some_host" + + MockDatabaseConnection.objects.filter.return_value = [db1] + + recent_time = datetime.now() - timedelta(days=days_inactivity - 1) + query_log1 = MagicMock() + query_log1.run_at = recent_time + + mock_queryset = MagicMock() + mock_queryset.first.return_value = query_log1 + + def filter_side_effect(*args, **kwargs): + if kwargs.get("connection") == "test_alias1": + return mock_queryset + return MagicMock(first=MagicMock(return_value=None)) + + MockQueryLog.objects.filter.side_effect = filter_side_effect + + # Call the function + remove_unused_sqlite_dbs() + + # Assertions + self.assertTrue(os.path.exists(temp_db_path)) + + @patch("explorer.tasks.DatabaseConnection") + def test_handle_missing_file_gracefully(self, MockDatabaseConnection): + + days_inactivity = 30 + with patch("explorer.tasks.app_settings.EXPLORER_PRUNE_LOCAL_UPLOAD_COPY_DAYS_INACTIVITY", days_inactivity): + + db1 = MagicMock() + db1.local_name = "non_existent_db.sqlite" + db1.alias = "test_alias1" + db1.host = "some_host" + MockDatabaseConnection.objects.filter.return_value = [db1] + + remove_unused_sqlite_dbs() + + # Since the file does not exist, we just ensure no exception is raised + self.assertFalse(os.path.exists("non_existent_db.sqlite")) diff --git a/explorer/tests/test_telemetry.py b/explorer/tests/test_telemetry.py index 51888213..ea39fa05 100644 --- a/explorer/tests/test_telemetry.py +++ b/explorer/tests/test_telemetry.py @@ -7,6 +7,9 @@ class TestTelemetry(TestCase): + def setUp(self): + cache.delete("last_stat_sent_time") + def test_instance_identifier(self): v = instance_identifier() self.assertEqual(len(v), 36) @@ -26,6 +29,8 @@ def test_stats_not_sent_too_frequently(self, mocked_app_settings, mocked_thread) mocked_app_settings.EXPLORER_ENABLE_ANONYMOUS_STATS = True mocked_app_settings.UNSAFE_RENDERING = True mocked_app_settings.EXPLORER_CHARTS_ENABLED = True + mocked_app_settings.has_assistant = MagicMock(return_value=True) + mocked_app_settings.db_connections_enabled = MagicMock(return_value=True) mocked_app_settings.ENABLE_TASKS = True s1 = Stat(StatNames.QUERY_RUN, {"foo": "bar"}) s2 = Stat(StatNames.QUERY_RUN, {"mux": "qux"}) @@ -44,6 +49,14 @@ def test_stats_not_sent_too_frequently(self, mocked_app_settings, mocked_thread) s3.track() self.assertEqual(mocked_thread.call_count, 3) + @patch("explorer.telemetry.threading.Thread") + @patch("explorer.app_settings") + def test_stats_not_sent_if_disabled(self, mocked_app_settings, mocked_thread): + mocked_app_settings.EXPLORER_ENABLE_ANONYMOUS_STATS = False + s1 = Stat(StatNames.QUERY_RUN, {"foo": "bar"}) + s1.track() + self.assertEqual(mocked_thread.call_count, 0) + @patch("explorer.telemetry.MigrationRecorder.Migration.objects.filter") def test_get_install_quarter_with_no_migrations(self, mock_filter): mock_filter.return_value.order_by.return_value.first.return_value = None diff --git a/explorer/tests/test_user_connection_utils.py b/explorer/tests/test_user_connection_utils.py new file mode 100644 index 00000000..7888357d --- /dev/null +++ b/explorer/tests/test_user_connection_utils.py @@ -0,0 +1,192 @@ +from django.test import TestCase +from django.core.files.uploadedfile import SimpleUploadedFile +import pandas as pd +import os +import sqlite3 +from explorer.models import DatabaseConnection +from unittest.mock import patch, MagicMock +from explorer.ee.db_connections.utils import ( + get_sqlite_for_connection, + create_django_style_connection, + pandas_to_sqlite, + is_csv, + csv_to_typed_df +) + + +def _get_csv(csv_name): + current_script_dir = os.path.dirname(os.path.abspath(__file__)) + file_path = os.path.join(current_script_dir, "csvs", csv_name) + + # Open the file in binary mode and read its contents + with open(file_path, "rb") as file: + csv_bytes = file.read() + + return csv_bytes + + +class TestCsvToTypedDf(TestCase): + + def test_mixed_types(self): + df = csv_to_typed_df(_get_csv("mixed.csv")) + self.assertTrue(pd.api.types.is_object_dtype(df["Value1"])) + self.assertTrue(pd.api.types.is_object_dtype(df["Value2"])) + self.assertTrue(pd.api.types.is_object_dtype(df["Value3"])) + + def test_all_types(self): + df = csv_to_typed_df(_get_csv("all_types.csv")) + self.assertTrue(pd.api.types.is_datetime64_ns_dtype(df["Dates"])) + print(df["Integers"].dtype) + self.assertTrue(pd.api.types.is_integer_dtype(df["Integers"])) + self.assertTrue(pd.api.types.is_float_dtype(df["Floats"])) + self.assertTrue(pd.api.types.is_object_dtype(df["Strings"])) + + def test_integer_parsing(self): + df = csv_to_typed_df(_get_csv("integers.csv")) + self.assertTrue(pd.api.types.is_integer_dtype(df["Integers"])) + self.assertTrue(pd.api.types.is_integer_dtype(df["More_integers"])) + + def test_float_parsing(self): + df = csv_to_typed_df(_get_csv("floats.csv")) + self.assertTrue(pd.api.types.is_float_dtype(df["Floats"])) + + def test_date_parsing(self): + + # Will not handle these formats: + # Unix Timestamp: 1706232300 (Seconds since Unix Epoch - 1970-01-01 00:00:00 UTC) + # ISO 8601 Week Number: 2024-W04-3 (Year-WWeekNumber-Weekday) + # Day of Year: 2024-024 (Year-DayOfYear) + + df = csv_to_typed_df(_get_csv("dates.csv")) + self.assertTrue(pd.api.types.is_datetime64_ns_dtype(df["Dates"])) + + +class TestSQLiteConnection(TestCase): + + @patch("explorer.utils.get_s3_bucket") + @patch("os.path.exists") + @patch("os.getcwd") + def test_get_sqlite_for_connection_downloads_file_if_not_exists(self, mock_getcwd, mock_path_exists, + mock_get_s3_bucket): + mock_getcwd.return_value = "/tmp" + mock_path_exists.return_value = False + mock_s3 = MagicMock() + mock_get_s3_bucket.return_value = mock_s3 + + mock_explorer_connection = MagicMock() + mock_explorer_connection.name = "test_db" + mock_explorer_connection.host = "s3_bucket/test_db" + mock_explorer_connection.local_name = "/tmp/user_dbs/test_db" + + result = get_sqlite_for_connection(mock_explorer_connection) + + mock_s3.download_file.assert_called_once_with("s3_bucket/test_db", "/tmp/user_dbs/test_db") + self.assertIsNone(result.host) + self.assertEqual(result.name, "/tmp/user_dbs/test_db") + + @patch("explorer.utils.get_s3_bucket") + @patch("os.path.exists") + @patch("os.getcwd") + def test_get_sqlite_for_connection_skips_download_if_exists(self, mock_getcwd, mock_path_exists, + mock_get_s3_bucket): + mock_getcwd.return_value = "/tmp" + mock_path_exists.return_value = True + mock_s3 = MagicMock() + mock_get_s3_bucket.return_value = mock_s3 + + mock_explorer_connection = MagicMock() + mock_explorer_connection.name = "test_db" + mock_explorer_connection.host = "s3_bucket/test_db" + mock_explorer_connection.local_name = "/tmp/user_dbs/test_db" + + result = get_sqlite_for_connection(mock_explorer_connection) + + mock_s3.download_file.assert_not_called() + self.assertIsNone(result.host) + self.assertEqual(result.name, "/tmp/user_dbs/test_db") + + +class TestDjangoStyleConnection(TestCase): + + @patch("explorer.ee.db_connections.utils.get_sqlite_for_connection") + @patch("explorer.ee.db_connections.utils.load_backend") + def test_create_django_style_connection_sqlite(self, mock_load_backend, mock_get_sqlite_for_connection): + mock_explorer_connection = MagicMock() + mock_explorer_connection.engine = DatabaseConnection.SQLITE + mock_explorer_connection.host = "s3_bucket/test_db" + mock_get_sqlite_for_connection.return_value = mock_explorer_connection + + mock_backend = MagicMock() + mock_load_backend.return_value = mock_backend + + create_django_style_connection(mock_explorer_connection) + + mock_get_sqlite_for_connection.assert_called_once_with(mock_explorer_connection) + mock_backend.DatabaseWrapper.assert_called_once() + + @patch("explorer.ee.db_connections.utils.load_backend") + def test_create_django_style_connection_non_sqlite(self, mock_load_backend): + mock_explorer_connection = MagicMock() + mock_explorer_connection.is_upload = False + mock_explorer_connection.engine = "django.db.backends.postgresql" + + mock_backend = MagicMock() + mock_load_backend.return_value = mock_backend + + create_django_style_connection(mock_explorer_connection) + + mock_load_backend.assert_called_once_with("django.db.backends.postgresql") + mock_backend.DatabaseWrapper.assert_called_once() + + +class TestPandasToSQLite(TestCase): + + def test_pandas_to_sqlite(self): + # Create a sample DataFrame + data = { + "column1": [1, 2, 3], + "column2": ["A", "B", "C"] + } + df = pd.DataFrame(data) + + # Convert the DataFrame to SQLite and get the BytesIO buffer + db_buffer = pandas_to_sqlite(df) + + # Write the buffer to a temporary file to simulate reading it back + temp_db_path = "temp_test_database.db" + with open(temp_db_path, "wb") as f: + f.write(db_buffer.getbuffer()) + + # Connect to the SQLite database and verify its content + con = sqlite3.connect(temp_db_path) + try: + cursor = con.cursor() + cursor.execute("SELECT * FROM data") + rows = cursor.fetchall() + + # Verify the content of the SQLite database + self.assertEqual(len(rows), 3) + self.assertEqual(rows[0], (1, "A")) + self.assertEqual(rows[1], (2, "B")) + self.assertEqual(rows[2], (3, "C")) + finally: + con.close() + os.remove(temp_db_path) + + +class TestIsCsvFunction(TestCase): + + def test_is_csv_with_csv_file(self): + # Create a SimpleUploadedFile with content_type set to "text/csv" + csv_file = SimpleUploadedFile("test.csv", b"column1,column2\n1,A\n2,B", content_type="text/csv") + self.assertTrue(is_csv(csv_file)) + + def test_is_csv_with_non_csv_file(self): + # Create a SimpleUploadedFile with content_type set to "text/plain" + txt_file = SimpleUploadedFile("test.txt", b"Just some text", content_type="text/plain") + self.assertFalse(is_csv(txt_file)) + + def test_is_csv_with_empty_content_type(self): + # Create a SimpleUploadedFile with an empty content_type + empty_file = SimpleUploadedFile("test.csv", b"column1,column2\n1,A\n2,B", content_type="") + self.assertFalse(is_csv(empty_file)) diff --git a/explorer/tests/test_utils.py b/explorer/tests/test_utils.py index 0d067e45..1fa2ebfa 100644 --- a/explorer/tests/test_utils.py +++ b/explorer/tests/test_utils.py @@ -269,5 +269,5 @@ def test_only_registered_connections_are_in_connections(self): from explorer.app_settings import EXPLORER_DEFAULT_CONNECTION from explorer.connections import connections - self.assertTrue(EXPLORER_DEFAULT_CONNECTION in connections) - self.assertNotEqual(len(connections), len([c for c in djcs])) + self.assertTrue(EXPLORER_DEFAULT_CONNECTION in connections()) + self.assertNotEqual(len(connections()), len([c for c in djcs])) diff --git a/explorer/tests/test_views.py b/explorer/tests/test_views.py index 1b184413..9fc30df8 100644 --- a/explorer/tests/test_views.py +++ b/explorer/tests/test_views.py @@ -2,10 +2,12 @@ import json import time import unittest -from unittest.mock import Mock, patch +from unittest.mock import Mock, patch, MagicMock from django.contrib.auth.models import User +from django.core.files.uploadedfile import SimpleUploadedFile from django.core.cache import cache +from django.conf import settings from django.db import connections from django.forms.models import model_to_dict from django.shortcuts import redirect @@ -15,9 +17,10 @@ from explorer import app_settings from explorer.app_settings import EXPLORER_DEFAULT_CONNECTION as CONN from explorer.app_settings import EXPLORER_TOKEN -from explorer.models import MSG_FAILED_BLACKLIST, Query, QueryFavorite, QueryLog +from explorer.models import MSG_FAILED_BLACKLIST, Query, QueryFavorite, QueryLog, DatabaseConnection from explorer.tests.factories import QueryLogFactory, SimpleQueryFactory from explorer.utils import user_can_see_query +from explorer.ee.db_connections.views import DatabaseConnectionsListView def reload_app_settings(): @@ -68,7 +71,7 @@ def test_run_count(self): for _ in range(0, 4): q.log() resp = self.client.get(reverse("explorer_index")) - self.assertContains(resp, "") + self.assertContains(resp, "4") class TestQueryCreateView(TestCase): @@ -343,13 +346,13 @@ def test_multiple_connections_integration(self): from explorer.connections import connections c1_alias = EXPLORER_CONNECTIONS["SQLite"] - conn = connections[c1_alias] + conn = connections()[c1_alias] c = conn.cursor() c.execute("CREATE TABLE IF NOT EXISTS animals (name text NOT NULL);") c.execute("INSERT INTO animals ( name ) VALUES ('peacock')") c2_alias = EXPLORER_CONNECTIONS["Another"] - conn = connections[c2_alias] + conn = connections()[c2_alias] c = conn.cursor() c.execute("CREATE TABLE IF NOT EXISTS animals (name text NOT NULL);") c.execute("INSERT INTO animals ( name ) VALUES ('superchicken')") @@ -582,13 +585,13 @@ def test_sql_download_respects_connection(self): from explorer.connections import connections c1_alias = EXPLORER_CONNECTIONS["SQLite"] - conn = connections[c1_alias] + conn = connections()[c1_alias] c = conn.cursor() c.execute("CREATE TABLE IF NOT EXISTS animals (name text NOT NULL);") c.execute("INSERT INTO animals ( name ) VALUES ('peacock')") c2_alias = EXPLORER_CONNECTIONS["Another"] - conn = connections[c2_alias] + conn = connections()[c2_alias] c = conn.cursor() c.execute("CREATE TABLE IF NOT EXISTS animals (name text NOT NULL);") c.execute("INSERT INTO animals ( name ) VALUES ('superchicken')") @@ -895,3 +898,178 @@ def test_toggle(self): ) resp = json.loads(resp.content.decode("utf-8")) self.assertFalse(resp["is_favorite"]) + + +class UploadDbViewTest(TestCase): + + def setUp(self): + DatabaseConnection.objects.all().delete() + self.user = User.objects.create_superuser( + "admin", "admin@admin.com", "pwd" + ) + self.client.login(username="admin", password="pwd") + + def test_post_csv_file(self): + file_content = "col1,col2\nval1,val2\nval3,val4" + uploaded_file = SimpleUploadedFile("test.csv", file_content.encode(), content_type="text/csv") + + with patch("explorer.ee.db_connections.views.is_csv", return_value=True), \ + patch("explorer.ee.db_connections.views.csv_to_typed_df") as mock_csv_to_typed_df, \ + patch("explorer.ee.db_connections.views.pandas_to_sqlite") as mock_pandas_to_sqlite, \ + patch("explorer.ee.db_connections.views.upload_sqlite") as mock_upload_sqlite: + mock_csv_to_typed_df.return_value = MagicMock() + mock_pandas_to_sqlite.return_value = b"sqlite_content" + + response = self.client.post(reverse("explorer_upload"), {"file": uploaded_file}) + + self.assertEqual(response.status_code, 200) + self.assertJSONEqual(response.content, {"success": True}) + self.assertTrue(DatabaseConnection.objects.filter(alias="test.db").exists()) + mock_upload_sqlite.assert_called_once() + mock_csv_to_typed_df.assert_called_once() + + def test_post_no_file(self): + response = self.client.post(reverse("explorer_upload")) + + self.assertEqual(response.status_code, 400) + self.assertJSONEqual(response.content, {"error": "No file provided"}) + + def test_delete_existing_connection(self): + dbc = DatabaseConnection.objects.create( + alias="test.db", + engine=DatabaseConnection.SQLITE, + name="test.db", + host="s3_path/test.db" + ) + + with patch("explorer.ee.db_connections.views.delete_from_s3") as mock_delete_from_s3: + response = self.client.delete(reverse("explorer_connection_delete", kwargs={"pk": dbc.id})) + + self.assertEqual(response.status_code, 302) + self.assertFalse(DatabaseConnection.objects.filter(alias="test.db").exists()) + mock_delete_from_s3.assert_called_once_with("s3_path/test.db") + + def test_delete_non_existing_connection(self): + response = self.client.delete("/upload/?alias=nonexistent.db") + + self.assertEqual(response.status_code, 404) + + +class DatabaseConnectionsListViewTestCase(TestCase): + + def setUp(self): + # Setting up the app settings for testing + self.original_explorer_connections = settings.EXPLORER_CONNECTIONS + settings.EXPLORER_CONNECTIONS = { + "default": "default_connection_alias", + } + + def tearDown(self): + # Restore the original app settings + settings.EXPLORER_CONNECTIONS = self.original_explorer_connections + + def test_queryset_combines_connections(self): + # Create a DatabaseConnection instance + db_conn = DatabaseConnection.objects.create(name="Test Connection", alias="test_alias") + + # Instantiate the view and call get_queryset + view = DatabaseConnectionsListView() + queryset = view.get_queryset() + + # Convert queryset to a list for easy comparison + queryset_list = list(queryset) + + # Check if both the created connection and the connections from app settings are in the queryset + self.assertIn(db_conn, queryset_list) + self.assertTrue(any(conn.alias == "default" for conn in queryset_list)) + + +class DatabaseConnectionValidateViewTestCase(TestCase): + + def setUp(self): + + self.user = User.objects.create_superuser( + "admin", "admin@admin.com", "pwd" + ) + self.client.login(username="admin", password="pwd") + self.url = reverse("explorer_connection_validate") + self.valid_data = { + "alias": "test_alias", + "engine": "django.db.backends.sqlite3", + "name": ":memory:", + "user": "", + "password": "", + "host": "", + "port": "", + } + self.invalid_data = { + "alias": "", + "engine": "", + "name": "", + "user": "", + "password": "", + "host": "", + "port": "", + } + + def test_validate_connection_success(self): + response = self.client.post(self.url, data=self.valid_data) + self.assertEqual(response.status_code, 200) + self.assertJSONEqual(response.content, {"success": True}) + + def test_validate_connection_invalid_form(self): + response = self.client.post(self.url, data=self.invalid_data) + self.assertEqual(response.status_code, 200) + self.assertJSONEqual(response.content, {"success": False, "error": "Invalid form data"}) + + def test_update_existing_connection(self): + DatabaseConnection.objects.create(alias="test_alias", engine="django.db.backends.sqlite3", name=":memory:") + response = self.client.post(self.url, data=self.valid_data) + self.assertEqual(response.status_code, 200) + self.assertJSONEqual(response.content, {"success": True}) + + @patch("explorer.ee.db_connections.views.create_django_style_connection") + def test_database_connection_error(self, mock_create_connection): + mock_create_connection.side_effect = Exception("Connection error") + response = self.client.post(self.url, data=self.valid_data) + self.assertEqual(response.status_code, 200) + self.assertJSONEqual(response.content, {"success": False, "error": "Connection error"}) + + +class SimpleViewTests(TestCase): + + def setUp(self): + self.user = User.objects.create_superuser( + "admin", "admin@admin.com", "pwd" + ) + self.client.login(username="admin", password="pwd") + self.connection = DatabaseConnection.objects.create( + alias="test_alias", + engine="django.db.backends.sqlite3", + name=":memory:", + user="", + password="", + host="", + port="" + ) + + def test_database_connection_detail_view(self): + response = self.client.get(reverse("explorer_connection_detail", args=[self.connection.pk])) + self.assertEqual(response.status_code, 200) + + def test_database_connection_create_view(self): + response = self.client.get(reverse("explorer_connection_create")) + self.assertEqual(response.status_code, 200) + + def test_database_connection_update_view(self): + response = self.client.get(reverse("explorer_connection_update", args=[self.connection.pk])) + self.assertEqual(response.status_code, 200) + + def test_database_connections_list_view(self): + response = self.client.get(reverse("explorer_connections")) + self.assertEqual(response.status_code, 200) + + def test_database_connection_delete_view(self): + response = self.client.get(reverse("explorer_connection_delete", args=[self.connection.pk])) + self.assertEqual(response.status_code, 200) + diff --git a/explorer/urls.py b/explorer/urls.py index 18b6d794..36fec7bb 100644 --- a/explorer/urls.py +++ b/explorer/urls.py @@ -1,12 +1,12 @@ from django.urls import path +from explorer.ee.urls import ee_urls from explorer.views import ( CreateQueryView, DeleteQueryView, DownloadFromSqlView, DownloadQueryView, EmailCsvQueryView, ListQueryLogView, ListQueryView, PlayQueryView, QueryFavoritesView, QueryFavoriteView, QueryView, SchemaJsonView, SchemaView, - StreamQueryView, format_sql, + StreamQueryView, format_sql ) - -from .assistant.views import assistant_help +from explorer.assistant.views import AssistantHelpView urlpatterns = [ path( @@ -43,5 +43,7 @@ path("favorites/", QueryFavoritesView.as_view(), name="query_favorites"), path("favorite/", QueryFavoriteView.as_view(), name="query_favorite"), path("", ListQueryView.as_view(), name="explorer_index"), - path("assistant/", assistant_help, name="assistant"), + path("assistant/", AssistantHelpView.as_view(), name="assistant"), ] + +urlpatterns += ee_urls diff --git a/explorer/utils.py b/explorer/utils.py index 88dcf0f7..7cb1607c 100644 --- a/explorer/utils.py +++ b/explorer/utils.py @@ -182,30 +182,47 @@ def get_valid_connection(alias=None): from explorer.connections import connections if not alias: - return connections[app_settings.EXPLORER_DEFAULT_CONNECTION] + return connections()[app_settings.EXPLORER_DEFAULT_CONNECTION] - if alias not in connections: + if alias not in connections(): raise InvalidExplorerConnectionException( f"Attempted to access connection {alias}, " f"but that is not a registered Explorer connection." ) - return connections[alias] + return connections()[alias] + + +def delete_from_s3(s3_path): + s3_bucket = get_s3_bucket() + s3_bucket.delete_objects( + Delete={ + "Objects": [ + {"Key": s3_path} + ] + } + ) def get_s3_bucket(): import boto3 from botocore.client import Config + + config = Config( + signature_version=app_settings.S3_SIGNATURE_VERSION, + region_name=app_settings.S3_REGION + ) + kwargs = { "aws_access_key_id": app_settings.S3_ACCESS_KEY, "aws_secret_access_key": app_settings.S3_SECRET_KEY, - "region_name": app_settings.S3_REGION, - "config": Config( - signature_version=app_settings.S3_SIGNATURE_VERSION - ) + "config": config } + if app_settings.S3_ENDPOINT_URL: kwargs["endpoint_url"] = app_settings.S3_ENDPOINT_URL + s3 = boto3.resource("s3", **kwargs) + return s3.Bucket(name=app_settings.S3_BUCKET) diff --git a/explorer/views/__init__.py b/explorer/views/__init__.py index 1a3a8c1a..5aa0d7c0 100644 --- a/explorer/views/__init__.py +++ b/explorer/views/__init__.py @@ -10,7 +10,6 @@ from .schema import SchemaJsonView, SchemaView from .stream import StreamQueryView - __all__ = [ "CreateQueryView", "DeleteQueryView", diff --git a/explorer/views/list.py b/explorer/views/list.py index 456a251a..79f87c8d 100644 --- a/explorer/views/list.py +++ b/explorer/views/list.py @@ -106,13 +106,15 @@ def _build_queries_and_headers(self): }) rendered_headers.append(header) + lrl = q.last_run_log() model_dict.update({ "is_in_category": headers[header] > 1, "collapse_target": collapse_target, "created_at": q.created_at, "is_header": False, "run_count": q.querylog_set.count(), - "ran_successfully": q.querylog_set.first() and q.querylog_set.first().success, + "ran_successfully": lrl.success, + "last_run_at": lrl.run_at, "created_by_user": str(q.created_by_user) if q.created_by_user else None, "is_favorite": q.id in query_favorites_for_user diff --git a/explorer/views/mixins.py b/explorer/views/mixins.py index f973acbe..4cf70225 100644 --- a/explorer/views/mixins.py +++ b/explorer/views/mixins.py @@ -14,7 +14,12 @@ def gen_ctx(self): "can_change": app_settings.EXPLORER_PERMISSION_CHANGE( self.request ), - "assistant_enabled": app_settings.EXPLORER_AI_API_KEY is not None, + "can_manage_connections": app_settings.EXPLORER_PERMISSION_CONNECTIONS( + self.request + ), + "assistant_enabled": app_settings.has_assistant(), + "db_connections_enabled": app_settings.db_connections_enabled(), + "s3_enabled": app_settings.S3_BUCKET is not None, "csrf_cookie_name": settings.CSRF_COOKIE_NAME, "csrf_cookie_httponly": settings.CSRF_COOKIE_HTTPONLY, "view_name": self.request.resolver_match.view_name, diff --git a/explorer/views/schema.py b/explorer/views/schema.py index d751f2d2..d1807308 100644 --- a/explorer/views/schema.py +++ b/explorer/views/schema.py @@ -19,7 +19,7 @@ def dispatch(self, *args, **kwargs): def get(self, request, *args, **kwargs): connection = kwargs.get("connection", "") - if connection not in connections: + if connection not in connections(): raise Http404 schema = schema_info(connection) if schema: @@ -38,6 +38,6 @@ class SchemaJsonView(PermissionRequiredMixin, View): def get(self, request, *args, **kwargs): connection = kwargs.get("connection", "") - if connection not in connections: + if connection not in connections(): raise Http404 return JsonResponse(schema_json_info(connection)) diff --git a/public_key.pem b/public_key.pem new file mode 100644 index 00000000..ce8677cf --- /dev/null +++ b/public_key.pem @@ -0,0 +1,9 @@ +-----BEGIN PUBLIC KEY----- +MIIBIjANBgkqhkiG9w0BAQEFAAOCAQ8AMIIBCgKCAQEAr1Lfhtla6gpaBlKb2r9s +5udyW6zHt0my9924sFtJ8OWzWpUFzGTDHySclWfKJkxdAS//zVhID0+i3FvkMCAe +kldnU9qDvSWx40zebWZ6JlnYA/cEl2XbiXOrLyAueeYZsr0Y1ZomMbFuMee8Kmjl +gj/YV+XGzlA3BBT3ajeCsBEYuSUIuGKsn7VUG6fjKX/TsOKNcGOCCXXyE9gcsrSe +J4h63eJTN+IZBC78E9f+0y8VqWZ/k1lpmPFLMSwr/Z5oeS1mvadFIikRq6UWNBC7 +YZxh7XSsfDSC1oaTSCgWiN+DDvK+mhBC4tpm1vJJ6lFWjS/pUrCoas0BMU5r3QJp +oQIDAQAB +-----END PUBLIC KEY----- diff --git a/requirements/base.txt b/requirements/base.txt index 870c6633..5569a2e5 100644 --- a/requirements/base.txt +++ b/requirements/base.txt @@ -1,3 +1,5 @@ Django>=3.2 sqlparse>=0.4.0 -requests>=2.2 \ No newline at end of file +requests>=2.2 +django-cryptography @ git+https://github.com/saurav-codes/django-cryptography.git@ac210338dd2c84a410452e0b8e18ddee43f1920f +cryptography>=42.0 diff --git a/requirements/extra/assistant.txt b/requirements/extra/assistant.txt index fe136bf7..c5ae0910 100644 --- a/requirements/extra/assistant.txt +++ b/requirements/extra/assistant.txt @@ -1,3 +1,3 @@ openai>=1.6.1 sql_metadata>=2.10 -tiktoken>=0.6 \ No newline at end of file +tiktoken>=0.7 diff --git a/requirements/extra/user_connections.txt b/requirements/extra/user_connections.txt new file mode 100644 index 00000000..31c9bd8d --- /dev/null +++ b/requirements/extra/user_connections.txt @@ -0,0 +1,3 @@ +python-dateutil>=2.9 +pandas>=2.2 +boto3>=1.20.0 diff --git a/test_project/celery_config.py b/test_project/celery_config.py index 7ba8afb2..ab3d9beb 100644 --- a/test_project/celery_config.py +++ b/test_project/celery_config.py @@ -1,18 +1,35 @@ import os from celery import Celery +from celery.schedules import crontab -# Set the default Django settings module for the 'celery' program. +# Set the default Django settings module for the "celery" program. os.environ.setdefault("DJANGO_SETTINGS_MODULE", "test_project.settings") app = Celery("test_project") -# Using a string here means the worker doesn't have to serialize +# Using a string here means the worker doesn"t have to serialize # the configuration object to child processes. -# - namespace='CELERY' means all celery-related configuration keys +# - namespace="CELERY" means all celery-related configuration keys # should have a `CELERY_` prefix. app.config_from_object("django.conf:settings", namespace="CELERY") # Load task modules from all registered Django apps. app.autodiscover_tasks() + +app.conf.beat_schedule = { + "explorer.tasks.snapshot_queries": { + "task": "explorer.tasks.snapshot_queries", + "schedule": crontab(hour="1", minute="0") + }, + "explorer.tasks.truncate_querylogs": { + "task": "explorer.tasks.truncate_querylogs", + "schedule": crontab(hour="1", minute="10"), + "kwargs": {"days": 30} + }, + "explorer.tasks.remove_unused_sqlite_dbs": { + "task": "explorer.tasks.remove_unused_sqlite_dbs", + "schedule": crontab(hour="1", minute="20") + } +} diff --git a/test_project/settings.py b/test_project/settings.py index 2bdbcf6c..427ae17c 100644 --- a/test_project/settings.py +++ b/test_project/settings.py @@ -25,24 +25,18 @@ "NAME": "tmp2" } }, - # "northwinds": { - # "ENGINE": "django.db.backends.sqlite3", - # "NAME": "northwind.db", - # "TEST": { - # "NAME": "tmp3" - # } - # }, } EXPLORER_CONNECTIONS = { "SQLite": "default", "Another": "alt", - #"Northwinds": "northwinds" } EXPLORER_DEFAULT_CONNECTION = "default" ROOT_URLCONF = "test_project.urls" +PROJECT_PATH = os.path.realpath(os.path.dirname(__file__)) + TEMPLATES = [ { "BACKEND": "django.template.backends.django.DjangoTemplates", @@ -111,8 +105,11 @@ # Tasks disabled by default, but if you have celery installed # make sure the broker URL is set correctly EXPLORER_TASKS_ENABLED = False -CELERY_BROKER_URL = 'redis://localhost:6379/0' +CELERY_BROKER_URL = os.environ.get('CELERY_BROKER_URL') -EXPLORER_S3_BUCKET = "thisismybucket.therearemanylikeit.butthisoneismine" +EXPLORER_S3_BUCKET = os.environ.get("EXPLORER_S3_BUCKET") +EXPLORER_S3_ACCESS_KEY = os.environ.get("EXPLORER_S3_ACCESS_KEY") +EXPLORER_S3_SECRET_KEY = os.environ.get("EXPLORER_S3_SECRET_KEY") EXPLORER_AI_API_KEY = os.environ.get("AI_API_KEY") EXPLORER_ASSISTANT_BASE_URL = os.environ.get("AI_BASE_URL") +EXPLORER_DB_CONNECTIONS_ENABLED = True diff --git a/tox.ini b/tox.ini index ab90e8c9..f6ff1554 100644 --- a/tox.ini +++ b/tox.ini @@ -18,7 +18,7 @@ deps = dev: -r requirements/dev.txt commands = {envpython} --version - {env:COMMAND:coverage} run manage.py test --settings=tests.settings --noinput + {env:COMMAND:coverage} run manage.py test --settings=explorer.tests.settings --noinput ignore_outcome = djmain: True ignore_errors =
    {% trans "Query" %}{% trans "Created" %}{% trans "Query" %}{% trans "Created" %}{% trans "Email" %}{% trans "Email" %}{% trans "CSV" %} {% trans "Delete" %}{% trans "Favorite" %}{% trans "Ran Successfully" %}{% trans "Run Count" %}{% trans "Last Run" %}{% trans "Run Count" %}{% trans "Connection" %}
    + {{ object.title }} {{ object.created_at|date:"SHORT_DATE_FORMAT" }} + {{ object.created_at|date:"m/d/y" }} {% if object.created_by_user %} {% blocktrans trimmed with cuser=object.created_by_user %} by {{cuser}} @@ -101,7 +108,7 @@

    {% trans "All Queries" %}

    {% if can_change %}     		- + @@ -111,13 +118,15 @@

    {% trans "All Queries" %}

    		 {% query_favorite_button object.id object.is_favorite 'query_favorite_toggle' %}{% if object.ran_successfully %} - + {% if object.ran_successfully %} + {% elif object.ran_successfully is not None %} - + {% endif %} + {{ object.last_run_at|date:"m/d/y" }} {{ object.run_count }}{{ object.run_count }}{{ object.connection }}
    4