citus

gistfile1.txt

import tqdm
import psycopg2
import random
from multiprocessing import Pool



def go():
    conn = psycopg2.connect("dbname=postgres user=postgres host=127.0.0.1")
    conn.set_session(autocommit=True)
    conn.set_isolation_level(0)

    cursor = conn.cursor()

    cursor.execute('DROP TABLE IF EXISTS message')
    cursor.execute('DROP TABLE IF EXISTS chat')

    cursor.execute('''
create table chat
(
	id bigserial not null
		constraint chat_pkey
			primary key,
	title text not null
)
    ''')

    cursor.execute('''
    
CREATE TABLE message (
  id bigserial,
  chat_id bigint REFERENCES chat (id),
  content text NOT NULL,

  PRIMARY KEY (chat_id, id)

--   FOREIGN KEY (company_id, ad_id)    -- added
--   REFERENCES ads (company_id, id)
);



    ''')

    cursor.execute('''
    SELECT create_distributed_table('chat',   'id');
    ''')
    cursor.execute('''    SELECT create_distributed_table('message',   'chat_id');
''')


    cursor.close()
    conn.close()


def go2():
    conn = psycopg2.connect("dbname=postgres user=postgres host=127.0.0.1")
    conn.set_session(autocommit=True)
    conn.set_isolation_level(0)

    cursor = conn.cursor()
    cursor.execute('DROP DATABASE IF EXISTS db_boner')
    cursor.execute('CREATE DATABASE db_boner')
    conn.close()

    conn = psycopg2.connect("dbname=db_boner user=postgres host=127.0.0.1")
    conn.set_session(autocommit=True)
    conn.set_isolation_level(0)
    cursor = conn.cursor()
    cursor.execute('DROP TABLE IF EXISTS message')
    cursor.execute('DROP TABLE IF EXISTS chat')

    cursor.execute('''
            create table chat
            (
                id bigserial PRIMARY KEY,
                title text not null
            )
    ''')

    cursor.execute('''
        CREATE TABLE message (
          id bigserial PRIMARY KEY,
          chat_id bigint REFERENCES chat (id),
          content text NOT NULL
        
        
        --   FOREIGN KEY (company_id, ad_id)    -- added
        --   REFERENCES ads (company_id, id)
        );

    ''')


    cursor.close()
    conn.close()

def run_on_citus(i):
    print ('citus', i)
    conn = psycopg2.connect("dbname=postgres user=postgres host=127.0.0.1")
    conn.set_session(autocommit=True)
    conn.set_isolation_level(0)

    cursor = conn.cursor()

    for i in range(300):
        cursor.execute(
            'INSERT INTO chat (title) VALUES (\'title' + str(i) + '\')');

    for i in range(1000):
        cursor.execute(
            'INSERT INTO message (content, chat_id) VALUES (\'title' + str(
                i) + '\', ' + str(random.randint(1, 299)) + ')');

    cursor.close()
    conn.close()

def run_on_original(i):
    print(i)
    conn = psycopg2.connect("dbname=db_boner user=postgres host=127.0.0.1")
    conn.set_session(autocommit=True)
    conn.set_isolation_level(0)
    cursor = conn.cursor()

    for i in range(300):
        cursor.execute(
            'INSERT INTO chat (title) VALUES (\'title' + str(i) + '\')');

    for i in range(1000):
        cursor.execute(
            'INSERT INTO message (content, chat_id) VALUES (\'title' + str(
                i) + '\', ' + str(random.randint(1, 299)) + ')');

    cursor.close()
    conn.close()

import time

# TEST ORIGINAL
go2()
start_time = time.time()
with Pool(4) as p:
    p.map(run_on_original, list(range(4)))

print("--- %s seconds ---" % (time.time() - start_time))


# TEST CITUS
go()

start_time = time.time()
with Pool(4) as p:
    p.map(run_on_citus, list(range(4)))

print("--- %s seconds ---" % (time.time() - start_time))