multi_tenant_auto_increment_poc.py

# import pip
# pip.main(['install', 'psycopg2', 'pandas'])

import sys
import subprocess
import psycopg2
import threading

DSN = "dbname=postgres user=piti host=localhost"

def setup_db():
    conn = psycopg2.connect(DSN)
    cur = conn.cursor()

    # Counter table
    cur.execute("""
    DROP TABLE IF EXISTS todo_counter;
    CREATE TABLE IF NOT EXISTS todo_counter (
        user_id bigint PRIMARY KEY,
        next_val bigint NOT NULL DEFAULT 1
    );
    """)

    # Todo table
    cur.execute("""
    DROP TABLE IF EXISTS todo;
    CREATE TABLE IF NOT EXISTS todo (
        id bigserial PRIMARY KEY,
        user_id bigint NOT NULL,
        user_todo_id bigint NOT NULL,
        created_at timestamptz NOT NULL DEFAULT now(),
        title text
    );
    """)

    # Trigger function using INSERT ... ON CONFLICT
    cur.execute("""
    CREATE OR REPLACE FUNCTION assign_user_todo_id()
    RETURNS trigger AS $$
    DECLARE
        v_next bigint;
    BEGIN
        INSERT INTO todo_counter(user_id, next_val)
        VALUES (NEW.user_id, 2)  -- first assigned ID = 1
        ON CONFLICT(user_id)
        DO UPDATE SET next_val = todo_counter.next_val + 1
        RETURNING next_val - 1 INTO v_next;

        NEW.user_todo_id := v_next;
        RETURN NEW;
    END;
    $$ LANGUAGE plpgsql;
    """)

    # Trigger
    cur.execute("""
    DROP TRIGGER IF EXISTS trg_assign_user_todo_id ON todo;
    CREATE TRIGGER trg_assign_user_todo_id
    BEFORE INSERT ON todo
    FOR EACH ROW
    EXECUTE FUNCTION assign_user_todo_id();
    """)

    conn.commit()
    conn.close()

# Worker thread: insert multiple rows for a user
def insert_rows(user_id, count, results):
    conn = psycopg2.connect(DSN)
    cur = conn.cursor()
    for _ in range(count):
        cur.execute("""
            INSERT INTO todo(user_id, title)
            VALUES (%s, 'task')
            RETURNING id, user_todo_id, created_at;
        """, (user_id,))
        results.append(cur.fetchone())
    conn.commit()
    conn.close()

# Run test with concurrency
def run_test():
    setup_db()
    NUM_THREADS = 10
    INSERTS_PER_THREAD = 50
    USER_ID = [999,888]

    threads = []
    results = []

    for t in range(NUM_THREADS):
        t = threading.Thread(target=insert_rows, args=(USER_ID[t%2], INSERTS_PER_THREAD, results))
        t.start()
        threads.append(t)

    for t in threads:
        t.join()

    print(f"Inserted {len(results)} rows.")

    # Check duplicates for per-user counter
    user_todo_ids = [(r[0], r[1]) for r in results]
    dup = len(user_todo_ids) - len(set(user_todo_ids))
    if dup == 0:
        print("✅ No duplicates — per-user counter is safe!")
    else:
        print("❌ Duplicates detected:", dup)

    print("Max per-user todo_id:", max(user_todo_ids))
    print("Max global id:", max(r[0] for r in results))

run_test()

import pandas as pd
import psycopg2

def print_table_df(dsn, table_name, limit=None):
    """
    Fetch a table from Postgres and print it as a Pandas DataFrame.
    
    :param dsn: Postgres DSN
    :param table_name: Table name
    :param limit: Optional row limit
    """
    conn = psycopg2.connect(dsn)
    
    sql = f"SELECT * FROM {table_name}"
    if limit:
        sql += f" LIMIT {limit}"
    
    df = pd.read_sql(sql, conn)
    conn.close()
    
    print(df)  # prints nicely like a DataFrame
    return df  # return df for further analysis
print_table_df(DSN, 'todo')