Psycopg: get query size

I would like to get size of a query before actually performing it with psycopg but I keep hitting

ProgrammingError: the last operation didn't produce a result

What i have currently is the following python script that should run a SQL query and return a table with its size as a result.

import dotenv
import os
import psycopg
from psycopg.rows import dict_row

_ = dotenv.load_dotenv("../env/norther.env")
__ = dotenv.load_dotenv("../.env.development")
if _ and __:
    print("Environment Variables Loaded Successfully")
else:
    raise Exception("Environment Variables Not Loaded")

_q = """
-- DROP FUNCTION IF EXISTS get_query_size;

CREATE OR REPLACE FUNCTION get_query_size() 
RETURNS TABLE(total_size BIGINT, table_size BIGINT, index_size BIGINT) AS $$
BEGIN
  -- Create the temporary table
  CREATE TEMP TABLE QueryDataTable AS
  -- Generic SQL query to get data from a table with a JSONB column
  SELECT * from my_table;
  
  -- Gather size information
  RETURN QUERY
  SELECT 
    pg_total_relation_size('QueryDataTable') AS total_size,
    pg_relation_size('QueryDataTable') AS table_size,
    pg_total_relation_size('QueryDataTable') - pg_relation_size('QueryDataTable') AS index_size;
  
  -- Drop the temporary table
  PERFORM pg_sleep(0); -- Ensures no empty result error
  DROP TABLE QueryDataTable;
END $$ LANGUAGE plpgsql;

-- To call the function and get the size
SELECT * FROM get_query_size()
"""

# Connect to your PostgreSQL database
conn = psycopg.connect(
        user=os.getenv("_DATABASE_USERNAME"),
        password=os.getenv("_DATABASE_PASSWORD"),
        host=os.getenv("_DATABASE_SERVER"),
        dbname=os.getenv("_DATABASE_NAME"),
        port=os.getenv("_PORT"),
        row_factory=dict_row,
        cursor_factory=psycopg.ClientCursor,
    )

# Create a cursor
cur = conn.cursor()

# Execute the function call
cur.execute(_q)


# Fetch the result
result = cur.fetchall()

# Print the result
print(result)

# Close the cursor and connection
cur.close()
conn.close()

The query on PGAdmin does return a result, though.

1. What am I missing?
2. Is there an alternative approach to achieve this?

Solution

Thanks to Frank Heikens.

connection = psycopg.connect(
    user=os.getenv("_DATABASE_USERNAME"),
    password=os.getenv("_DATABASE_PASSWORD"),
    host=os.getenv("_DATABASE_SERVER"),
    dbname=os.getenv("_DATABASE_NAME"),
    port=os.getenv("_PORT"),
    row_factory=dict_row,
)

try:
    with connection as conn:
        cur = conn.cursor()
        cur.execute("BEGIN;")

        # Create the temporary table
        cur.execute("""
            CREATE TEMP TABLE QueryDataTable ON COMMIT DROP AS
            SELECT * from my_table;
        """)

        # Gather size information
        cur.execute("""
            SELECT
                pg_total_relation_size('QueryDataTable') AS total_size,
                pg_relation_size('QueryDataTable') AS table_size,
                pg_indexes_size('QueryDataTable') AS index_size;
        """)

        # Fetch the result
        result = cur.fetchall()

        # Print the result
        print(result)

        # Commit the transaction
        cur.execute("COMMIT;")

except Exception as e:
    # If an error occurs, rollback the transaction
    conn.rollback()
    print(f"An error occurred: {e}")
finally:
    # Close the connection
    conn.close()

Even better, using Psycopg Transactions:

try:
    with connection as conn:
        cur = conn.cursor()
        with conn.transaction():
            # Create the temporary table
            cur.execute("""
                CREATE TEMP TABLE QueryDataTable ON COMMIT DROP AS
                    SELECT * FROM my_table;
            """)

            # Gather size information
            cur.execute("""
                SELECT
                    pg_total_relation_size('QueryDataTable') AS total_size,
                    pg_relation_size('QueryDataTable') AS table_size,
                    pg_indexes_size('QueryDataTable') AS index_size;
            """)

            # Fetch the result
            result = cur.fetchall()

            # Print the result
            print(result)
except Exception as e:
    # If an error occurs, print the error and rollback is automatic due to context manager
    print(f"An error occurred: {e}")

finally:
    # Close the connection
    conn.close()

Just execute these 4 statements:

BEGIN ; -- first
  CREATE TEMP TABLE QueryDataTable ON COMMIT DROP AS
  -- Generic SQL query to get data from a table with a JSONB column
  SELECT * from data.fact_account_card; -- second

  -- Gather size information
  SELECT
    pg_total_relation_size('QueryDataTable') AS total_size,
    pg_relation_size('QueryDataTable') AS table_size,
    pg_indexes_size('QueryDataTable') AS index_size; -- third, fetch the result

COMMIT; -- fourth, this will also drop the temp table

You can also use with conn.transaction(): in your python code. See https://www.psycopg.org/psycopg3/docs/basic/transactions.html