| import plpy |
| import pipes |
| from control import MinWarning |
| |
| def __get_madlib_temp_tables(target_schema): |
| sql_get_tables_to_drop = """ |
| SELECT quote_ident(tablename) AS tablename |
| FROM pg_tables |
| WHERE tablename LIKE E'%madlib\_temp%' |
| AND quote_ident(schemaname) = '{target_schema}' |
| """.format(**locals()) |
| return plpy.execute(sql_get_tables_to_drop) |
| |
| def __get_madlib_temp_views(target_schema): |
| sql_get_tables_to_drop = """ |
| SELECT quote_ident(viewname) AS viewname |
| FROM pg_views |
| WHERE viewname LIKE E'%madlib\_temp%' |
| AND quote_ident(schemaname) = '{target_schema}' |
| """.format(**locals()) |
| return plpy.execute(sql_get_tables_to_drop) |
| |
| # ------------------------------------------------------------------------------ |
| def cleanup_madlib_temp_tables(schema_madlib, target_schema, **kwargs): |
| """ Drop all tables matching '%madlib_temp%' in the given schema |
| |
| Args: |
| @param target_schema |
| |
| Returns: |
| None |
| """ |
| to_drop_list = __get_madlib_temp_tables(target_schema) |
| if len(to_drop_list) == 0: |
| plpy.info("No madlib temp tables found in schema {target_schema}.".format(**locals())) |
| return None |
| sql_drop_all = 'DROP TABLE ' |
| sql_drop_all += ",".join(["{target_schema}.{tablename}".format( |
| tablename=row['tablename'], **locals()) for row in to_drop_list]) |
| sql_drop_all += ";" |
| plpy.notice("Dropping {0} tables ...".format(len(to_drop_list))) |
| try: |
| with MinWarning(): |
| plpy.execute(sql_drop_all) |
| except plpy.spiexceptions.OutOfMemory: |
| escaped_sql = pipes.quote( |
| 'SELECT {schema_madlib}.cleanup_madlib_temp_tables_script(\'{target_schema}\')' |
| .format(**locals())) |
| plpy.error( |
| "You have too many tables and cannot be dropped all at once.\n" |
| "Use the following bash commands instead:\n" |
| " $ psql -A -t -c {escaped_sql} > /tmp/drop_all_madlib_temp.sql\n" |
| " $ psql -f /tmp/drop_all_madlib_temp.sql" |
| .format(**locals())) |
| plpy.notice("Finished dropping {0} tables.".format(len(to_drop_list))) |
| |
| # ------------------------------------------------------------------------------ |
| def cleanup_madlib_temp_tables_script(schema_madlib, target_schema, **kwargs): |
| """ Create a sql script to drop all tables matching '%madlib_temp%' in the given schema |
| |
| Args: |
| @param target_schema |
| |
| Returns: |
| None |
| """ |
| to_drop_list = __get_madlib_temp_tables(target_schema) |
| if len(to_drop_list) == 0: |
| plpy.info("No madlib temp tables found in schema {target_schema}.") |
| return None |
| sql_content = '' |
| for row in to_drop_list: |
| tablename = row['tablename'] |
| sql_drop = "DROP TABLE {target_schema}.{tablename};".format(**locals()) |
| sql_content += sql_drop + "\n" |
| return sql_content |
| |
| # ------------------------------------------------------------------------------ |
| def cleanup_madlib_temp_views(schema_madlib, target_schema, **kwargs): |
| to_drop_list = __get_madlib_temp_views(target_schema) |
| if len(to_drop_list) == 0: |
| plpy.info("No madlib temp views found in schema {target_schema}.".format(**locals())) |
| return None |
| sql_drop_all = 'DROP VIEW IF EXISTS ' |
| sql_drop_all += ",".join(["{target_schema}.{viewname}".format( |
| viewname=row['viewname'], **locals()) for row in to_drop_list]) |
| sql_drop_all += " CASCADE;" |
| plpy.notice("Dropping {0} views ...".format(len(to_drop_list))) |
| plpy.execute(sql_drop_all) |