Runs psql and returns a CSVReader object from the query This CSVReader includes header names as the first record in all situations. The output is fully buffered into Python.
(sql_command, error_handler=None)
| 32 | |
| 33 | |
| 34 | def psql_csv_run(sql_command, error_handler=None): |
| 35 | """ |
| 36 | Runs psql and returns a CSVReader object from the query |
| 37 | |
| 38 | This CSVReader includes header names as the first record in all |
| 39 | situations. The output is fully buffered into Python. |
| 40 | |
| 41 | """ |
| 42 | csv_query = ('COPY ({query}) TO STDOUT WITH CSV HEADER;' |
| 43 | .format(query=sql_command)) |
| 44 | |
| 45 | new_env = os.environ.copy() |
| 46 | new_env.setdefault('PGOPTIONS', '') |
| 47 | new_env["PGOPTIONS"] += ' --statement-timeout=0' |
| 48 | psql_proc = popen_nonblock([PSQL_BIN, '-d', 'postgres', '--no-password', |
| 49 | '--no-psqlrc', '-c', csv_query], |
| 50 | stdout=PIPE, |
| 51 | env=new_env) |
| 52 | stdout = psql_proc.communicate()[0].decode('utf-8') |
| 53 | |
| 54 | if psql_proc.returncode != 0: |
| 55 | if error_handler is not None: |
| 56 | error_handler(psql_proc) |
| 57 | else: |
| 58 | assert error_handler is None |
| 59 | raise UserException( |
| 60 | 'could not csv-execute a query successfully via psql', |
| 61 | 'Query was "{query}".'.format(sql_command), |
| 62 | 'You may have to set some libpq environment ' |
| 63 | 'variables if you are sure the server is running.') |
| 64 | |
| 65 | # Previous code must raise any desired exceptions for non-zero |
| 66 | # exit codes |
| 67 | assert psql_proc.returncode == 0 |
| 68 | |
| 69 | # Fake enough iterator interface to get a CSV Reader object |
| 70 | # that works. |
| 71 | return csv.reader(iter(stdout.strip().split('\n'))) |
| 72 | |
| 73 | |
| 74 | class PgBackupStatements(object): |
no test coverage detected