Psycopg：获取查询大小

Question

我想在实际使用 psycopg 执行查询之前获取查询的大小，但我一直在点击

ProgrammingError: the last operation didn't produce a result

我目前拥有的是以下 python 脚本，该脚本应该运行 SQL 查询并返回一个表及其大小作为结果。

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()

PGAdmin 上的查询确实返回了结果。

1. 我错过了什么？
2. 是否有其他方法可以实现这一目标？

解决方案

感谢弗兰克·海肯斯。

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()

更好的是，使用 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()

Answer 1

只需执行这4条语句：

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

您还可以在 Python 代码中与

conn.transaction():

一起使用。请参阅https://www.psycopg.org/psycopg3/docs/basic/transactions.html