Python psycopg2.extensions.ISOLATION_LEVEL_AUTOCOMMIT Examples

def notify(self, name, sec=0, payload=None):
        """Send a notification to the database, eventually after some time."""
        if payload is None:
            payload = ''
        else:
            payload = ", %r" % payload

        script = ("""\
import time
time.sleep(%(sec)s)
import %(module)s as psycopg2
import %(module)s.extensions as ext
conn = psycopg2.connect(%(dsn)r)
conn.set_isolation_level(ext.ISOLATION_LEVEL_AUTOCOMMIT)
print conn.get_backend_pid()
curs = conn.cursor()
curs.execute("NOTIFY " %(name)r %(payload)r)
curs.close()
conn.close()
""" % {
        'module': psycopg2.__name__,
        'dsn': dsn, 'sec': sec, 'name': name, 'payload': payload})

        return Popen([sys.executable, '-c', script_to_py3(script)], stdout=PIPE) 

def handle(self, *args, **options):
        self.stdout.write(self.style.SUCCESS('Starting to drop DB..'))

        dbname = settings.DATABASES['default']['NAME']
        user = settings.DATABASES['default']['USER']
        password = settings.DATABASES['default']['PASSWORD']
        host = settings.DATABASES['default']['HOST']

        self.stdout.write(self.style.SUCCESS('Connecting to host..'))
        con = connect(dbname='postgres', user=user, host=host, password=password)
        con.set_isolation_level(ISOLATION_LEVEL_AUTOCOMMIT)

        self.stdout.write(self.style.SUCCESS("Dropping database '{}'".format(dbname)))
        cur = con.cursor()
        cur.execute('DROP DATABASE ' + dbname)
        cur.close()

        con.close()

        self.stdout.write(self.style.SUCCESS('All done!')) 

def handle(self, *args, **options):
        self.stdout.write(self.style.SUCCESS('Starting DB creation..'))

        dbname = settings.DATABASES['default']['NAME']
        user = settings.DATABASES['default']['USER']
        password = settings.DATABASES['default']['PASSWORD']
        host = settings.DATABASES['default']['HOST']

        self.stdout.write(self.style.SUCCESS('Connecting to host..'))
        con = connect(dbname='postgres', user=user, host=host, password=password)
        con.set_isolation_level(ISOLATION_LEVEL_AUTOCOMMIT)

        self.stdout.write(self.style.SUCCESS('Creating database'))
        cur = con.cursor()
        cur.execute('CREATE DATABASE ' + dbname)
        cur.close()

        con.close()

        self.stdout.write(self.style.SUCCESS('All done!')) 

def handle(self, *args, **options):
        self.stdout.write(self.style.SUCCESS('Starting schema deletion...'))

        dbname = settings.DATABASES['default']['NAME']
        user = settings.DATABASES['default']['USER']
        password = settings.DATABASES['default']['PASSWORD']
        host = settings.DATABASES['default']['HOST']

        con = connect(dbname=dbname, user=user, host=host, password=password)

        self.stdout.write(self.style.SUCCESS('Removing schema {schema} from database {dbname}'
                                             .format(schema=settings.SCHEMA, dbname=dbname)))
        con.set_isolation_level(ISOLATION_LEVEL_AUTOCOMMIT)

        cur = con.cursor()
        cur.execute('DROP SCHEMA {schema} CASCADE;'.format(schema=settings.SCHEMA))
        cur.close()

        con.close()

        self.stdout.write(self.style.SUCCESS('All done.')) 

def handle(self, *args, **options):
        self.stdout.write(self.style.SUCCESS('Starting Schema creation..'))

        dbname = settings.DATABASES['default']['NAME']
        user = settings.DATABASES['default']['USER']
        password = settings.DATABASES['default']['PASSWORD']
        host = settings.DATABASES['default']['HOST']

        con = connect(dbname=dbname, user=user, host=host, password=password)

        self.stdout.write(self.style.SUCCESS('Adding schema {schema} to database {dbname}'
                                             .format(schema=settings.SCHEMA, dbname=dbname)))
        con.set_isolation_level(ISOLATION_LEVEL_AUTOCOMMIT)

        cur = con.cursor()
        cur.execute('CREATE SCHEMA {schema};'.format(schema=settings.SCHEMA))
        cur.close()

        con.close()

        self.stdout.write(self.style.SUCCESS('All Done!')) 

def random_database():
    """
    Creates a random database in the testing Postgres instance and returns the
    name of the database.
    """
    # Setup connection with default credentials for testing.
    with connect(dbname='holo', user='holocleanuser', password='abcd1234', host='localhost') as conn:
        conn.set_isolation_level(ISOLATION_LEVEL_AUTOCOMMIT)
        with conn.cursor() as cur:
            while True:
                # Generate a random DB name that is not already in Postgres.
                db_name = 'test_holo_{}'.format(random.randint(0, 1e6))
                cur.execute("""
                    SELECT EXISTS(
                        SELECT datname FROM pg_catalog.pg_database
                        WHERE datname = '{db_name}'
                    );
                """.format(db_name=db_name))
                if cur.fetchall()[0][0]:
                    continue

                cur.execute("CREATE DATABASE {db_name}".format(db_name=db_name))
                return db_name 

def __cleanup_schema(self):
        with psycopg2.connect(os.environ['GREASE_TEST_DSN_ORIGINAL']) as conn:
            conn.set_isolation_level(ISOLATION_LEVEL_AUTOCOMMIT)
            with conn.cursor() as cursor:
                cursor.execute("""
                  SELECT  
                    pg_terminate_backend(pid) 
                    FROM pg_stat_activity 
                    WHERE datname='test_data'
                """)
                try:
                    cursor.execute("""
                        DROP DATABASE test_data;
                    """)
                except psycopg2.ProgrammingError as e:
                    print("Schema Does Not Exist: {0}".format(e.pgerror))
        os.environ['GREASE_TEST_DSN'] = os.environ['GREASE_TEST_DSN_ORIGINAL'] 

def runAutoCommit(self, func, *args, **kwargs):
        """Call the given function in auto-commit mode. Needed to execute
        statements that cannot be run in a transaction, like CREATE
        DATABASE.

        WARNING: This will commit any open transaction!
        """
        old_level = self.dbh.isolation_level
        try:
            if self.inTransaction():
                self.dbh.commit()
            self.dbh.set_isolation_level(psy_ext.ISOLATION_LEVEL_AUTOCOMMIT)
            return func(*args, **kwargs)
        finally:
            self.dbh.set_isolation_level(old_level)

    # resetting the auto increment values of primary keys 

def create_db_if_needed(cls, db):
        if cls.SCHEMA == 'postgresql':
            con = psycopg2.connect(user='postgres', host='localhost', password='postgres')
            con.set_isolation_level(ISOLATION_LEVEL_AUTOCOMMIT)
            cur = con.cursor()
            try:
                cur.execute(sql.SQL("CREATE DATABASE {}").format(
                    sql.Identifier(db))
                )
            except Exception:
                pass
            finally:
                cur.close()
                con.close()
        elif cls.SCHEMA == 'mysql':
            con = mysql.connector.connect(user='root', host='localhost')
            cur = con.cursor()
            cur.execute('CREATE DATABASE IF NOT EXISTS {}'.format(db))
            con.commit()
            cur.close()
            con.close() 

def create_init_db():
    """
    create the initial role to own the db
    create the new db as the db super user assign owner to new role
    :return: new db with new owner
    """
    con = connect(dbname='postgres', user=db_superuser, host=dbhost, port=dbport, password=db_superuser_pwd)
    con.set_isolation_level(ISOLATION_LEVEL_AUTOCOMMIT)
    cur = con.cursor()
    cur.execute('CREATE ROLE \"{0}\" LOGIN ENCRYPTED PASSWORD \'{1}\''.format(dbowner_name, dbowner_pwd))
    print("creating role")
    cur.execute('''CREATE DATABASE \"{0}\" WITH OWNER = \'{1}\'
                    ENCODING = \'UTF8\' TEMPLATE=template0 CONNECTION LIMIT = -1;'''.format(db_name, dbowner_name))
    print("creating database")
    cur.close()
    con.close() 

def connectToPostgres():
    """
    If globusmonitor database does not exist yet:
        $ initdb /home/globusmonitor/postgres/data
        $ pg_ctl -D /home/globusmonitor/postgres/data -l /home/globusmonitor/postgres/log
        $   createdb globusmonitor
    """
    try:
        conn = psycopg2.connect(dbname='globusmonitor')
    except:
        # Attempt to create database if not found
        conn = psycopg2.connect(dbname='postgres')
        conn.set_isolation_level(ISOLATION_LEVEL_AUTOCOMMIT)
        curs = conn.cursor()
        curs.execute('CREATE DATABASE globusmonitor;')
        curs.close()
        conn.commit()
        conn.close()

        conn = psycopg2.connect(dbname='globusmonitor')
        initializeDatabase(conn)

    return conn 

def autocommit(self, conn):
        """Set a connection in autocommit mode."""
        conn.set_isolation_level(extensions.ISOLATION_LEVEL_AUTOCOMMIT) 

def __ensure_schema(self):
        # Helper function
        try:
            if not os.environ.get('GREASE_TEST_DSN'):
                os.environ['GREASE_TEST_DSN'] = "host=localhost user=postgres"
            with psycopg2.connect(os.environ['GREASE_TEST_DSN']) as conn:
                conn.set_isolation_level(ISOLATION_LEVEL_AUTOCOMMIT)
                with conn.cursor() as cursor:
                    # just to make sure nothing exists
                    try:
                        cursor.execute("""
                        SELECT  
                            pg_terminate_backend(pid) 
                            FROM pg_stat_activity 
                            WHERE datname='test_data'
                        """)
                        cursor.execute("""
                            DROP DATABASE test_data;
                        """)
                    except:
                        print("Exception occurred during ensure schema... most of the time this is fine")
                    try:
                        cursor.execute("""
                            CREATE DATABASE test_data;
                        """)
                    except psycopg2.ProgrammingError as e:
                        print("Schema Exists: {0}".format(e.pgerror))
            if not os.environ.get('GREASE_TEST_DSN_ORIGINAL'):
                os.environ['GREASE_TEST_DSN_ORIGINAL'] = os.environ.get('GREASE_TEST_DSN')
            os.environ['GREASE_TEST_DSN'] = os.environ['GREASE_TEST_DSN'] + " dbname=test_data"
            
        except psycopg2.Error as e:
            print("psycopg2 Exception occurred, {}".format(e.pgerror)) 

def _isolation_lookup(self):
        extensions = self._psycopg2_extensions()
        return {
            "AUTOCOMMIT": extensions.ISOLATION_LEVEL_AUTOCOMMIT,
            "READ COMMITTED": extensions.ISOLATION_LEVEL_READ_COMMITTED,
            "READ UNCOMMITTED": extensions.ISOLATION_LEVEL_READ_UNCOMMITTED,
            "REPEATABLE READ": extensions.ISOLATION_LEVEL_REPEATABLE_READ,
            "SERIALIZABLE": extensions.ISOLATION_LEVEL_SERIALIZABLE,
        } 

def create_postgres_db(connection_dict, config):
    if check_db_or_user_exists(connection_dict["db_name"], connection_dict["db_username"], config):
        raise ValueError("db or user already exists")
    with _create_pg_connection(config) as con:
        con.set_isolation_level(ISOLATION_LEVEL_AUTOCOMMIT)
        with con.cursor() as cur:
            try:
                cur.execute(SQL("CREATE USER {} WITH PASSWORD %s;").format(
                    Identifier(connection_dict['db_username']),
                ), (
                    connection_dict['db_pwd'],
                ))
                cur.execute(SQL("GRANT {} TO {};").format(
                    Identifier(connection_dict['db_username']),
                    Identifier(config['username']),
                ))
            except psycopg2.ProgrammingError as e:
                raise ValueError(e.args[0])
            # cleanup role in case database creation fails
            # sadly 'CREATE DATABASE' cannot run inside a transaction block
            try:
                cur.execute(SQL("CREATE DATABASE {} OWNER {};").format(
                    Identifier(connection_dict['db_name']),
                    Identifier(connection_dict['db_username']),
                ))
            except psycopg2.ProgrammingError as e:
                cur.execute(SQL("DROP USER {};").format(
                    Identifier(connection_dict['db_username']),
                ))
                raise ValueError(e.args[0]) 

def delete_database(db_name, config):
    with _create_pg_connection(config) as con:
        con.set_isolation_level(ISOLATION_LEVEL_AUTOCOMMIT)
        with con.cursor() as cur:
            try:
                cur.execute(SQL('DROP DATABASE {};').format(
                    Identifier(db_name),
                ))
            except psycopg2.ProgrammingError as e:
                raise ValueError(e.args[0]) 

def delete_user(username, config):
    with _create_pg_connection(config) as con:
        con.set_isolation_level(ISOLATION_LEVEL_AUTOCOMMIT)
        with con.cursor() as cur:
            try:
                cur.execute(SQL('DROP USER {};').format(
                   Identifier(get_normalized_username(username)),
                ))
            except psycopg2.ProgrammingError as e:
                raise ValueError(e.args[0]) 

def test_database_creation(self):

        create_database(database=self.db_name)

        conn = psycopg2.connect(database=self.userdir['main_db'],
                                user=self.userdir['main_user'],
                                password=self.userdir['main_password'],
                                host=self.userdir['main_host'],
                                port=self.userdir['main_port'])
        conn.set_isolation_level(ISOLATION_LEVEL_AUTOCOMMIT)

        with conn:
            cur = conn.cursor()
            cur.execute("""SELECT datname FROM pg_catalog.pg_database
                        WHERE lower(datname)=lower('%s')""" % self.db_name)
            database_exist = cur.fetchone()

            self.assertEqual(len(database_exist), 1)

            # cur.execute("""SELECT pg_terminate_backend(pg_stat_activity.pid)
            #             FROM pg_stat_activity
            #             WHERE datname = current_database()
            #             AND pid <> pg_backend_pid()""")
            cur.execute("""DROP DATABASE IF EXISTS %s""" % self.db_name)
            cur.close()
        conn.close() 

def init_postgresql(_):
    con = psycopg2.connect(user='postgres', host='localhost', password='postgres')
    con.set_isolation_level(ISOLATION_LEVEL_AUTOCOMMIT)
    cur = con.cursor()
    cur.execute(sql.SQL("DROP DATABASE IF EXISTS {}").format(
        sql.Identifier(DATABASE_NAME))
    )
    cur.execute(sql.SQL("CREATE DATABASE {}").format(
        sql.Identifier(DATABASE_NAME))
    )
    cur.close()
    con.close() 

def db_remove(cls):
        close_all_sessions()
        if cls.SCHEMA == 'sqlite':
            for db in [DATABASEFILE_UNITTESTS, DATABASEFILE_UNITTESTS_2]:
                if os.path.isfile(db):
                    os.remove(db)
        elif cls.SCHEMA == 'postgresql':
            for db in [DATABASE_NAME, DATABASE_NAME_2]:
                cls.create_db_if_needed(db)
                con = psycopg2.connect(user='postgres', host='localhost', password='postgres', database=db)
                con.set_isolation_level(ISOLATION_LEVEL_AUTOCOMMIT)
                cur = con.cursor()
                try:
                    # drop all tables
                    cur.execute(sql.SQL("""
                        DO $$ DECLARE
                            r RECORD;
                        BEGIN
                            FOR r IN (SELECT tablename FROM pg_tables WHERE schemaname = current_schema()) LOOP
                                EXECUTE 'DROP TABLE IF EXISTS ' || quote_ident(r.tablename) || ' CASCADE';
                            END LOOP;
                        END $$;"""
                    ))
                finally:
                    cur.close()
                    con.close()
        elif cls.SCHEMA == 'mysql':
            for db in [DATABASE_NAME, DATABASE_NAME_2]:
                cls.create_db_if_needed(db)
                con = mysql.connector.connect(user='root', host='localhost', database=db, autocommit=True)
                cur = con.cursor(buffered=True)
                try:
                    cur.execute("DROP DATABASE {};".format(db))
                    cur.execute("CREATE DATABASE {};".format(db))
                finally:
                    cur.close()
                    con.close() 

def _isolation_lookup(self):
        extensions = self._psycopg2_extensions()
        return {
            'AUTOCOMMIT': extensions.ISOLATION_LEVEL_AUTOCOMMIT,
            'READ COMMITTED': extensions.ISOLATION_LEVEL_READ_COMMITTED,
            'READ UNCOMMITTED': extensions.ISOLATION_LEVEL_READ_UNCOMMITTED,
            'REPEATABLE READ': extensions.ISOLATION_LEVEL_REPEATABLE_READ,
            'SERIALIZABLE': extensions.ISOLATION_LEVEL_SERIALIZABLE
        } 

def _isolation_lookup(self):
        from psycopg2 import extensions
        return {
            'AUTOCOMMIT': extensions.ISOLATION_LEVEL_AUTOCOMMIT,
            'READ COMMITTED': extensions.ISOLATION_LEVEL_READ_COMMITTED,
            'READ UNCOMMITTED': extensions.ISOLATION_LEVEL_READ_UNCOMMITTED,
            'REPEATABLE READ': extensions.ISOLATION_LEVEL_REPEATABLE_READ,
            'SERIALIZABLE': extensions.ISOLATION_LEVEL_SERIALIZABLE
        } 

def create_db(dbname):
    args = CONNECT_ARGS.copy()
    del args['database']
    conn = psycopg2.connect(**args)
    conn.set_isolation_level(ISOLATION_LEVEL_AUTOCOMMIT)
    with conn.cursor() as curs:
        curs.execute('CREATE DATABASE ' + dbname)
    conn.close() 

def _isolation_lookup(self):
        extensions = self._psycopg2_extensions()
        return {
            'AUTOCOMMIT': extensions.ISOLATION_LEVEL_AUTOCOMMIT,
            'READ COMMITTED': extensions.ISOLATION_LEVEL_READ_COMMITTED,
            'READ UNCOMMITTED': extensions.ISOLATION_LEVEL_READ_UNCOMMITTED,
            'REPEATABLE READ': extensions.ISOLATION_LEVEL_REPEATABLE_READ,
            'SERIALIZABLE': extensions.ISOLATION_LEVEL_SERIALIZABLE
        } 

def _create_database_postgres(self):
        conn = connect(dbname='postgres', user=self.username,
                       password=self.password, host=self.postgres_host, port=self.postgres_port)
        conn.set_isolation_level(ISOLATION_LEVEL_AUTOCOMMIT)
        cursor = conn.cursor()
        cursor.execute('DROP DATABASE IF EXISTS ' + self.dbname)
        cursor.execute('CREATE DATABASE ' + self.dbname)
        conn.commit()
        cursor.close()
        conn.close() 

def __listen(socketio):
    conn = connect_db()
    conn.set_isolation_level(ISOLATION_LEVEL_AUTOCOMMIT)
    cur = conn.cursor()
    cur.execute("LISTEN job_update")

    while True:
        trampoline(conn, read=True)
        conn.poll()
        while conn.notifies:
            n = conn.notifies.pop()
            socketio.start_background_task(__handle_event,
                                           json.loads(n.payload),
                                           socketio) 

def prepare_db():
    """
    Create a database with name in .env
    """
    try:
        con = psycopg2.connect(dbname='postgres', user=USER, password=PASSWORD)
    except psycopg2.Error as e:
        raise e
    logging.info('Connected to database postgres')
    con.set_isolation_level(ISOLATION_LEVEL_AUTOCOMMIT)
    cur = con.cursor()
    try:
        cur.execute('CREATE DATABASE ' + DB_NAME)
    except psycopg2.Error as e:
        logging.info('DROP OLD DATABASE')
        logging.info('CREATE NEW DATABASE')
        cur.execute('DROP DATABASE ' + DB_NAME)
        cur.execute('CREATE DATABASE ' + DB_NAME)
    cur.close()
    con.close()

    con = psycopg2.connect(dbname=DB_NAME, user=USER, password=PASSWORD)
    cur = con.cursor()
    cur.execute('CREATE EXTENSION CUBE')
    cur.execute('CREATE TABLE images (id serial, name text, url text, vector cube);')
    con.commit()
    cur.close()
    con.close() 

def autocommit(self, on_off=1):
        """autocommit(on_off=1) -> switch autocommit on (1) or off (0)"""
        if on_off > 0:
            self.set_isolation_level(_ext.ISOLATION_LEVEL_AUTOCOMMIT)
        else:
            self.set_isolation_level(_ext.ISOLATION_LEVEL_READ_COMMITTED) 

def delete_database(db_name):
    with connect(dbname='holo', user='holocleanuser', password='abcd1234', host='localhost') as conn:
        conn.set_isolation_level(ISOLATION_LEVEL_AUTOCOMMIT)
        with conn.cursor() as cur:
            # Terminate un-closed connections.
            cur.execute("""
            SELECT pid, pg_terminate_backend(pid)
            FROM pg_stat_activity
            WHERE datname = '{db_name}' AND pid <> pg_backend_pid();""".format(db_name=db_name))
            # Drop the database.
            cur.execute("DROP DATABASE IF EXISTS {db_name}".format(db_name=db_name)) 

def _isolation_lookup(self):
        extensions = self._psycopg2_extensions()
        return {
            'AUTOCOMMIT': extensions.ISOLATION_LEVEL_AUTOCOMMIT,
            'READ COMMITTED': extensions.ISOLATION_LEVEL_READ_COMMITTED,
            'READ UNCOMMITTED': extensions.ISOLATION_LEVEL_READ_UNCOMMITTED,
            'REPEATABLE READ': extensions.ISOLATION_LEVEL_REPEATABLE_READ,
            'SERIALIZABLE': extensions.ISOLATION_LEVEL_SERIALIZABLE
        }