common.py

'''
common.py
Copyright (c) 2016-2024, Postgres Professional
'''

importpsycopg2
importpsycopg2.extensions
importselect
importtime

BACKEND_IS_IDLE_INFO='INFO: state of backend is idle\n'
BACKEND_IS_ACTIVE_INFO='INFO: state of backend is active\n'

defwait(conn):
"""wait for some event on connection to postgres"""
while1:
state=conn.poll()
ifstate==psycopg2.extensions.POLL_OK:
break
elifstate==psycopg2.extensions.POLL_WRITE:
select.select([], [conn.fileno()], [])
elifstate==psycopg2.extensions.POLL_READ:
select.select([conn.fileno()], [], [])
else:
raisepsycopg2.OperationalError("poll() returned %s"%state)

defn_async_connect(config, n=1):
"""establish n asynchronious connections to the postgres with specified config"""

aconfig=config.copy()
aconfig['async'] =True

result= []
for_inrange(n):
conn=psycopg2.connect(**aconfig)
wait(conn)
result.append(conn)
returnresult

defn_close(conns):
"""close connections to postgres"""

forconninconns:
conn.close()

defpg_query_state_locks(config, pid, conn, verbose=False, costs=False, timing=False, \
buffers=False, triggers=False, format='text'):
"""
 Get query state from backend with specified pid and optional parameters.
 Save any warning, info, notice and log data in global variable 'notices'
 """

curs=conn.cursor()
curs.callproc('pg_query_state', (pid, verbose, costs, timing, buffers, triggers, format))
wait(conn)
result=curs.fetchall()
notices=conn.notices[:]

returnresult, notices

defpg_query_state(config, pid, verbose=False, costs=False, timing=False, \
buffers=False, triggers=False, format='text'):
"""
 Get query state from backend with specified pid and optional parameters.
 Save any warning, info, notice and log data in global variable 'notices'
 """

conn=psycopg2.connect(**config)
curs=conn.cursor()
curs.callproc('pg_query_state', (pid, verbose, costs, timing, buffers, triggers, format))
result=curs.fetchall()
notices=conn.notices[:]
conn.close()

returnresult, notices

defonetime_query_state_locks(config, acon_query, acon_pg, query, args={}, num_workers=0):
"""
 Get intermediate state of 'query' on connection 'acon_query' after number of 'steps'
 of node executions from start of query
 """

curs_query=acon_query.cursor()
curs_pg=acon_pg.cursor()
curs_query.execute("select pg_advisory_lock(1);")
curs_pg.execute("select pg_advisory_lock(2);")
wait(acon_query)
wait(acon_pg)
curs_pg.execute("select pg_advisory_lock(1);")
set_guc(acon_query, 'enable_mergejoin', 'off')
set_guc(acon_query, 'max_parallel_workers_per_gather', num_workers)
curs_query.execute(query)
# extract current state of query progress
MAX_PG_QS_RETRIES=10
DELAY_BETWEEN_RETRIES=0.1
pg_qs_args= {
'config': config,
'pid': acon_query.get_backend_pid(),
'conn': acon_pg
 }
fork, vinargs.items():
pg_qs_args[k] =v
n_retries=0

wait(acon_pg)

whileTrue:
result, notices=pg_query_state_locks(**pg_qs_args)
n_retries+=1
iflen(result) >0:
break
ifn_retries>=MAX_PG_QS_RETRIES:
# pg_query_state callings don't return any result, more likely run
# query has completed
break
time.sleep(DELAY_BETWEEN_RETRIES)

curs_pg.execute("select pg_advisory_unlock(2);")
wait(acon_pg)
wait(acon_query)

set_guc(acon_query, 'enable_mergejoin', 'on')
curs_query.execute("select pg_advisory_unlock(2);")
curs_pg.execute("select pg_advisory_unlock(1);")
returnresult, notices

defonetime_query_state(config, async_conn, query, args={}, num_workers=0):
"""
 Get intermediate state of 'query' on connection 'async_conn' after number of 'steps'
 of node executions from start of query
 """

acurs=async_conn.cursor()

set_guc(async_conn, 'enable_mergejoin', 'off')
set_guc(async_conn, 'max_parallel_workers_per_gather', num_workers)
acurs.execute(query)

# extract current state of query progress
MAX_PG_QS_RETRIES=10
DELAY_BETWEEN_RETRIES=0.1
pg_qs_args= {
'config': config,
'pid': async_conn.get_backend_pid()
 }
fork, vinargs.items():
pg_qs_args[k] =v
n_retries=0
whileTrue:
result, notices=pg_query_state(**pg_qs_args)
n_retries+=1
iflen(result) >0:
break
ifn_retries>=MAX_PG_QS_RETRIES:
# pg_query_state callings don't return any result, more likely run
# query has completed
break
time.sleep(DELAY_BETWEEN_RETRIES)
wait(async_conn)

set_guc(async_conn, 'enable_mergejoin', 'on')
returnresult, notices

defset_guc(async_conn, param, value):
acurs=async_conn.cursor()
acurs.execute('set %s to %s'% (param, value))
wait(async_conn)