<commit_before>from sqlshare_rest.util.db import get_backend
from sqlshare_rest.models import Query
from sqlshare_rest.dao.dataset import reset_dataset_account_access
from django.utils import timezone
from django.conf import settings
from django import db
from time import sleep
from sqlshare_rest.util.queue_triggers import trigger_query_queue_processing
from sqlshare_rest.util.queue_triggers import QUERY_QUEUE_PORT_NUMBER
from sqlshare_rest.logger import getLogger
from django.db.utils import DatabaseError
import atexit
import signal
import json
import time
import sys
import os

import socket
from threading import Thread

import six
TERMINATE_TRIGGER_FILE = getattr(settings,
                                 "SQLSHARE_TERMINATE_QUERY_QUEUE_PATH",
                                 "/tmp/sqlshare_terminate_query_queue")


def process_queue(thread_count=0, run_once=True, verbose=False):
    def start_query(query, background=True):
        query.is_started = True
        query.save()
        query_id = query.pk

        if background:
            from django.db import connection
            connection.close()

            if os.fork():
                # This is the main process
                return

            os.setsid()

            if os.fork():
                # Double fork the daemon
                sys.exit(0)

            # Close stdin/out/err
            sys.stdin.flush()
            sys.stdout.flush()
            sys.stderr.flush()
            null = os.open(os.devnull, os.O_RDWR)
            os.dup2(null, sys.stdin.fileno())
            os.dup2(null, sys.stdout.fileno())
            os.dup2(null, sys.stderr.fileno())
            os.close(null)

        try:
            process_query(query_id)
        except Exception as ex:
            try:
                query = Query.objects.get(pk=query_id)
                query.has_error = True
                query.error = str(ex)
                query.is_finished = True
                query.save()
            except:
                # That try is just trying to get info out to the user, it's
                # relatively ok if that fails
                pass
            logger = getLogger(__name__)
            logger.error("Error on %s: %s" % (query_id, str(ex)))

        if background:
            sys.exit(0)

    def get_column_names_from_cursor(cursor):
        index = 0
        names = []
        for col in cursor.description:
            index += 1
            column_name = col[0]
            if column_name == "":
                column_name = "COLUMN%s" % index

            names.append(column_name)

        return names

    def process_query(query_id):
        logger = getLogger(__name__)
        query = Query.objects.get(pk=query_id)
        # queries can be cancelled before we see them.  clean it up now.

        if query.terminated:
            query.is_finished = True
            query.has_error = True
            query.error = "Query cancelled"
            query.save()
            return

        pid = os.getpid()
        query.process_queue_id = pid
        query.save()

        msg = "Processing query id %s, in process %s" % (
            query.pk,
            pid
        )
        logger.info(msg)
        if verbose:
            print(msg)
        user = query.owner
        row_count = 0
        backend = get_backend()
        try:
            start = timezone.now()
            query_plan = backend.get_query_plan(query.sql, user)

            t1 = time.time()

            sql = query.sql
            if query.is_ui_preview:
                sql = backend.get_preview_sql_for_query(sql)

            cursor = backend.run_query(sql,
                                       user,
                                       return_cursor=True)

            t2 = time.time()
            try:
                all_data = []
                for row in cursor:
                    all_data.append(list(row))
                    row_count += 1

                columns = get_column_names_from_cursor(cursor)
                formatted = json.dumps({"columns": columns, "data": all_data})
                query.preview_content = formatted
                t3 = time.time()
<commit_msg>Resolves sqlshr-224<commit_after>from sqlshare_rest.util.db import get_backend
from sqlshare_rest.models import Query
from sqlshare_rest.dao.dataset import reset_dataset_account_access
from django.utils import timezone
from django.conf import settings
from django import db
from time import sleep
from sqlshare_rest.util.queue_triggers import trigger_query_queue_processing
from sqlshare_rest.util.queue_triggers import QUERY_QUEUE_PORT_NUMBER
from sqlshare_rest.logger import getLogger
from django.db.utils import DatabaseError
import datetime
import atexit
import signal
import json
import time
import sys
import os

import socket
from threading import Thread

import six
TERMINATE_TRIGGER_FILE = getattr(settings,
                                 "SQLSHARE_TERMINATE_QUERY_QUEUE_PATH",
                                 "/tmp/sqlshare_terminate_query_queue")


def process_queue(thread_count=0, run_once=True, verbose=False):
    def start_query(query, background=True):
        query.is_started = True
        query.save()
        query_id = query.pk

        if background:
            from django.db import connection
            connection.close()

            if os.fork():
                # This is the main process
                return

            os.setsid()

            if os.fork():
                # Double fork the daemon
                sys.exit(0)

            # Close stdin/out/err
            sys.stdin.flush()
            sys.stdout.flush()
            sys.stderr.flush()
            null = os.open(os.devnull, os.O_RDWR)
            os.dup2(null, sys.stdin.fileno())
            os.dup2(null, sys.stdout.fileno())
            os.dup2(null, sys.stderr.fileno())
            os.close(null)

        try:
            process_query(query_id)
        except Exception as ex:
            try:
                query = Query.objects.get(pk=query_id)
                query.has_error = True
                query.error = str(ex)
                query.is_finished = True
                query.save()
            except:
                # That try is just trying to get info out to the user, it's
                # relatively ok if that fails
                pass
            logger = getLogger(__name__)
            logger.error("Error on %s: %s" % (query_id, str(ex)))

        if background:
            sys.exit(0)

    def get_column_names_from_cursor(cursor):
        index = 0
        names = []
        for col in cursor.description:
            index += 1
            column_name = col[0]
            if column_name == "":
                column_name = "COLUMN%s" % index

            names.append(column_name)

        return names

    def custom_encode(obj):
        if isinstance(obj, datetime.datetime):
            return str(obj)

    def process_query(query_id):
        logger = getLogger(__name__)
        query = Query.objects.get(pk=query_id)
        # queries can be cancelled before we see them.  clean it up now.

        if query.terminated:
            query.is_finished = True
            query.has_error = True
            query.error = "Query cancelled"
            query.save()
            return

        pid = os.getpid()
        query.process_queue_id = pid
        query.save()

        msg = "Processing query id %s, in process %s" % (
            query.pk,
            pid
        )
        logger.info(msg)
        if verbose:
            print(msg)
        user = query.owner
        row_count = 0
        backend = get_backend()
        try:
            start = timezone.now()
            query_plan = backend.get_query_plan(query.sql, user)

            t1 = time.time()

            sql = query.sql
            if query.is_ui_preview:
                sql = backend.get_preview_sql_for_query(sql)

            cursor = backend.run_query(sql,
                                       user,
                                       return_cursor=True)

            t2 = time.time()
            try:
                all_data = []
                for row in cursor:
                    all_data.append(list(row))
                    row_count += 1

                columns = get_column_names_from_cursor(cursor)
                formatted = json.dumps({"columns": columns, "data": all_data},
                                       default=custom_encode)
                query.preview_content = formatted
                t3 = time.time()

{'fraction_non_alphanumeric': 0.0764115725618292, 'fraction_numerical': 0.0036164255716285582, 'mean_word_length': 1.5454275534441806, 'pattern_counts': {'":': 6, '<': 3, '<?xml version=': 0, '>': 3, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 9, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '4312246', 'n_tokens_mistral': 2376, 'n_tokens_neox': 2260, 'n_words': 571}