import duckdb
from datetime import datetime
import threading
import queue as Queue

class DuckDBThreaded:
    def __init__(self,duckdb_insert_thread_count,thread_function):
        self.duckdb_insert_thread_count = duckdb_insert_thread_count
        self.threads = []
        self.thread_function = thread_function
        
    def multithread_test(self,size,if_all_true=True):
        duckdb_conn = duckdb.connect(check_same_thread=False)
        queue = Queue.Queue()
        return_value = False

        for i in range(0,self.duckdb_insert_thread_count):
            self.threads.append(threading.Thread(target=self.thread_function, args=(duckdb_conn,queue,size),name='duckdb_thread_'+str(i)))

        for i in range(0,len(self.threads)):
            self.threads[i].start()
            if not if_all_true:
                if queue.get():
                    return_value = True
            else:
                if i == 0 and queue.get():
                    return_value = True
                elif queue.get() and return_value:
                    return_value = True
            
        for i in range(0,len(self.threads)):
            self.threads[i].join()

        assert (return_value)


def fetchone_query(duckdb_conn, queue, size):
    try:
        a = duckdb_conn.execute("select sum(range) from  range ("+str(size)+")").fetchone()
        queue.put(True)
    except:
        queue.put(False)  

def run_test(num_threads, size):
    cur = datetime.now()
    duck_threads = DuckDBThreaded(num_threads,fetchone_query)
    duck_threads.multithread_test(size)
    print (str(num_threads) + "T " + str(size) + " elements")
    print (datetime.now() - cur)

run_test (10, 100000000)
run_test (20, 100000000)
run_test (100, 100000)