Python psycopg2.extras.RealDictCursor() Examples

def execute_query(conn, query):
    try:
        cursor = conn.cursor(cursor_factory=RealDictCursor)
        cursor.execute(query)
        rows = list(cursor.fetchall())

        return rows
    except Exception:
        # TODO this kind of thing is only necessary if we re-use connections

        # If any exception occurs during query execution, close the
        # connection to ensure it is not in an invalid state. The
        # connection pool knows to create new connections to replace
        # those that are closed
        try:
            conn.close()
        except Exception:
            pass
        raise 

def database_connection(database_url):
    parsed = urlparse(database_url)
    user = parsed.username
    password = parsed.password
    host = parsed.hostname
    port = parsed.port
    database = parsed.path.strip('/')

    try:
        connection = psycopg2.connect(
            host=host,
            port=port,
            user=user,
            password=password,
            database=database,
            cursor_factory=RealDictCursor)
    except psycopg2.OperationalError:
        raise UnableToConnectToDatabase
    connection.set_session(autocommit=True)

    return connection 

def get_personal_timetable_rows(upi):
    set_id = settings.ROOMBOOKINGS_SETID

    # Get from Django's ORM to raw psycopg2 so that a new cursor
    # factory can be used to fetch dicts.

    wrapped_connection = connections['gencache']
    if wrapped_connection.connection is None:
        cursor = wrapped_connection.cursor()

    raw_connection = wrapped_connection.connection

    bucket = 'a' if Lock.objects.all()[0].a else 'b'

    with raw_connection.cursor(cursor_factory=RealDictCursor) as cursor:
        cursor.callproc(
            'get_student_timetable_' + bucket,
            [
                upi,
                set_id
            ]
        )
        rows = cursor.fetchall()
        return rows 

def execute_cursor(self, sql, arguments):
        with self._get_connection() as conn:
            cursor = conn.cursor(cursor_factory=RealDictCursor)
            try:
                cursor.execute(sql, arguments)
                conn.commit()
                return cursor
            except:
                if not conn.closed:
                    conn.rollback()
                # cursor.close()

                # retry execution
                conn = self.connectionPool.getconn()
                try:
                    cursor = conn.cursor(cursor_factory=RealDictCursor)
                    cursor.execute(sql, arguments)
                    conn.commit()
                except Exception as e:
                    if not conn.closed:
                        conn.rollback()
                    print(e) 

def insert(self, sql, values):
        with self._get_connection() as conn:
            cursor = conn.cursor()
            try:
                execute_values(cursor, sql, values)
                conn.commit()
            except:
                if not conn.closed:
                    conn.rollback()
                # cursor.close()

                # retry execution
                conn = self.connectionPool.getconn()
                try:
                    cursor = conn.cursor(cursor_factory=RealDictCursor)
                    execute_values(cursor, sql, values)
                    conn.commit()
                except Exception as e:
                    if not conn.closed:
                        conn.rollback()
                    print(e) 

def fetch_block(self, block_num):
        fetch = """
        SELECT block_num, block_id FROM blocks WHERE block_num = {}
        """.format(block_num)

        with self._conn.cursor(cursor_factory=RealDictCursor) as cursor:
            cursor.execute(fetch)
            block = cursor.fetchone()

        return block 

def connect(self):
        # Check if connection already exists
        if self.connection != None:
            # Is Connection Closed?
            if self.connection.closed == 0:
                # If Not, Return the current cursor
                return self.cursor
        # Create a New Connection
        try:
            # Connect to Database
            self.connection = psycopg2.connect(host= DB_INFO['HOST_NAME'],
                user=DB_INFO['USERNAME'],
                database=DB_INFO['DATABASE_NAME'],
                password=DB_INFO['PASSWORD'])
            # Enable AutoCommit
            self.connection.autocommit = True
            # Set Cursor to DictCursor
            self.cursor = self.connection.cursor(cursor_factory = psycopg2.extras.DictCursor)
            self.RealDictCursor = self.connection.cursor(cursor_factory = rdc)

            logger.debug(__name__+" Connected to Database")
            
            # Try Initializing the Database
            if not self.initilaize():
                logger.error(__name__ + " Database Cannot be initialiazed automatically. Try it Manually.")
            
            return self.cursor
        except Exception as e:
            logger.critical(__name__+" Database Connection Error! Msg: " + str(e))
            if self.connection != None:
                self.connection.close()
                self.connection = None

        return False 

def getData(conn, query, params=None):
    "Use this for non-geometry SELECTs, produces plain json based on DB field names"
    with conn.cursor(cursor_factory=RealDictCursor) as cur:
        if (params):
            cur.execute(query, params)
        else:
            cur.execute(query)
	return json.dumps(cur.fetchall(), indent=2)


#removes CR LF characters from string, for safer logging 

def update_url_prioritization():
    """
    """
    log.info("Started update_url_prioritization")
    conn = connect_db(conf)
    cur = conn.cursor(cursor_factory=RealDictCursor)

    log.info("Regenerating URL prioritization file")
    sql = """SELECT priority, domain, url, cc, category_code FROM citizenlab"""
    cur.execute(sql)
    entries = list(cur.fetchall())
    conn.rollback()
    conn.close()

    # Create dict: cc -> category_code -> [entry, ... ]
    entries_by_country = {}
    for e in entries:
        country = e["cc"].upper()
        if country not in entries_by_country:
            entries_by_country[country] = {}
        ccode = e["category_code"]
        entries_by_country[country].setdefault(ccode, []).append(e)

    # merge ZZ into each country
    zz = entries_by_country.pop("ZZ")
    for ccode, country_dict in entries_by_country.items():
        for category_code, test_items in zz.items():
            country_dict.setdefault(category_code, []).extend(test_items)

    log.info("Update done")
    return entries_by_country 

def fetch_all_record_resources(self):
        fetch_records = """
        SELECT record_id FROM records
        WHERE ({0}) >= start_block_num
        AND ({0}) < end_block_num;
        """.format(LATEST_BLOCK_NUM)

        async with self._conn.cursor(cursor_factory=RealDictCursor) as cursor:
            try:
                await cursor.execute(fetch_records)
                records = await cursor.fetchall()

                for record in records:
                    fetch_record_locations = """
                    SELECT latitude, longitude, timestamp
                    FROM record_locations
                    WHERE record_id='{0}'
                    AND ({1}) >= start_block_num
                    AND ({1}) < end_block_num;
                    """.format(record['record_id'], LATEST_BLOCK_NUM)

                    fetch_record_owners = """
                    SELECT agent_id, timestamp
                    FROM record_owners
                    WHERE record_id='{0}'
                    AND ({1}) >= start_block_num
                    AND ({1}) < end_block_num;
                    """.format(record['record_id'], LATEST_BLOCK_NUM)

                    await cursor.execute(fetch_record_locations)
                    record['locations'] = await cursor.fetchall()

                    await cursor.execute(fetch_record_owners)
                    record['owners'] = await cursor.fetchall()

                return records
            except TypeError:
                return [] 

def fetch_auth_resource(self, public_key):
        fetch = """
        SELECT * FROM auth WHERE public_key='{}'
        """.format(public_key)

        async with self._conn.cursor(cursor_factory=RealDictCursor) as cursor:
            await cursor.execute(fetch)
            return await cursor.fetchone() 

def fetch_all_agent_resources(self):
        fetch = """
        SELECT public_key, name, timestamp FROM agents
        WHERE ({0}) >= start_block_num
        AND ({0}) < end_block_num;
        """.format(LATEST_BLOCK_NUM)

        async with self._conn.cursor(cursor_factory=RealDictCursor) as cursor:
            await cursor.execute(fetch)
            return await cursor.fetchall() 

def fetch_agent_resource(self, public_key):
        fetch = """
        SELECT public_key, name, timestamp FROM agents
        WHERE public_key='{0}'
        AND ({1}) >= start_block_num
        AND ({1}) < end_block_num;
        """.format(public_key, LATEST_BLOCK_NUM)

        async with self._conn.cursor(cursor_factory=RealDictCursor) as cursor:
            await cursor.execute(fetch)
            return await cursor.fetchone() 

def execute(conn, query, query_vars=()):
    """Execute given query and return fetched results"""
    with conn.cursor(cursor_factory=RealDictCursor) as cursor:
        cursor.execute(query, query_vars)
        return cursor.fetchall() 

def fetch_last_known_blocks(self, count):
        """Fetches the specified number of most recent blocks
        """
        fetch = """
        SELECT block_num, block_id FROM blocks
        ORDER BY block_num DESC LIMIT {}
        """.format(count)

        with self._conn.cursor(cursor_factory=RealDictCursor) as cursor:
            cursor.execute(fetch)
            blocks = cursor.fetchall()

        return blocks 

def query_cursor(self, q, lazy_fetch=False, commit=True):
        """Execute a query and yield a cursor.

        All execution performed by the Postgres object uses this method.

        Args:
            q (str): SQL query
            lazy_fetch (bool): whether to use a server-side cursor (lazily fetches results).
        """
        self.cursors_opened += 1

        if self.verbose:
            logging.debug(q)

        if self.debug:
            empty_cursor = Bunch()
            empty_cursor.fetchmany = lambda size: []
            empty_cursor.fetchall = lambda: []
            yield empty_cursor
            return

        cursor_name = 'server_side_{}'.format(self.cursors_opened) if lazy_fetch else None
        with self.connection.cursor(cursor_name, cursor_factory=RealDictCursor) as cursor:
            cursor.execute(q)
            yield cursor

        if commit:
            self.commit() 

def query_cursor(self, q, lazy_fetch=False, commit=True):
        """Execute a query and yield a cursor.

        All execution performed by the Postgres object uses this method.

        Args:
            q (str): SQL query
            lazy_fetch (bool): whether to use a server-side cursor (lazily fetches results).
        """
        self.cursors_opened += 1

        if self.verbose:
            logging.debug(q)

        if self.debug:
            empty_cursor = Bunch()
            empty_cursor.fetchmany = lambda size: []
            empty_cursor.fetchall = lambda: []
            yield empty_cursor
            return

        cursor_name = 'server_side_{}'.format(self.cursors_opened) if lazy_fetch else None
        with self.connection.cursor(cursor_name, cursor_factory=RealDictCursor) as cursor:
            cursor.execute(q)
            yield cursor

        if commit:
            self.commit() 

def dictCursor():
    """"Pull a dictionary cursor from the connection."""
    return connect().cursor(cursor_factory=pg2ext.RealDictCursor) 

def dictCursor():
    """"Pull a dictionary cursor from the connection."""
    return connect().cursor(cursor_factory=pg2ext.RealDictCursor) 

def execute(conn, query):
    """Execute given query and return fetched results"""
    with conn.cursor(cursor_factory=RealDictCursor) as cursor:
        cursor.execute(query)

        return cursor.fetchall() 

def parse_source(self, configuration):
        """This will Query the SQL Server to find data

        Args:
            configuration (dict): Configuration of Source. See Class Documentation above for more info

        Returns:
            bool: If True data will be scheduled for ingestion after deduplication. If False the engine will bail out

        """
        ioc = GreaseContainer()
        if configuration.get('hour'):
            if datetime.datetime.utcnow().hour != int(configuration.get('hour')):
                # it is not the correct hour
                return True
        if configuration.get('minute'):
            if datetime.datetime.utcnow().minute != int(configuration.get('minute')):
                # it is not the correct hour
                return True
        if configuration.get('type') != 'postgresql':
            ioc.getLogger().error("Unsupported SQL Server Type; Currently Only supporting PostgreSQL", notify=False)
            return False
        else:
            # Attempt to get the DSN for the connection
            if os.environ.get(configuration.get('dsn')) and configuration.get('query'):
                # ensure the DSN is setup and the query is present
                try:
                    DSN = os.environ.get(configuration.get('dsn'))
                    with psycopg2.connect(DSN) as conn:
                        with conn.cursor(cursor_factory=RealDictCursor) as cursor:
                            cursor.execute(configuration.get('query'))
                            data = cursor.fetchall()
                            for row in data:
                                self._data.append(row)
                            del ioc
                    return True
                except Exception as e:
                    # Naked except to prevent issues around connections
                    ioc.getLogger().error("Error processing configuration; Error [{0}]".format(e.message), notify=False)
                    del ioc
                    return False
            else:
                # could not get the DSN
                ioc.getLogger().error("Failed to locate the DSN variable", notify=False)
                del ioc
                return False 

def execute(self, sql, arguments, numReturn=None):
        with self._get_connection() as conn:
            cursor = conn.cursor(cursor_factory=RealDictCursor)

            # execute statement
            try:
                cursor.execute(sql, arguments)
                conn.commit()
            except Exception as e:
                if not conn.closed:
                    conn.rollback()
                # self.connectionPool.putconn(conn, close=False)    #TODO: this still causes connection to close
                conn = self.connectionPool.getconn()

                # retry execution
                try:
                    cursor = conn.cursor(cursor_factory=RealDictCursor)
                    cursor.execute(sql, arguments)
                    conn.commit()
                except:
                    if not conn.closed:
                        conn.rollback()
                    print(e)

            # get results
            try:
                returnValues = []
                if numReturn is None:
                    # cursor.close()
                    return
                
                elif numReturn == 'all':
                    returnValues = cursor.fetchall()
                    # cursor.close()
                    return returnValues

                else:
                    for _ in range(numReturn):
                        rv = cursor.fetchone()
                        if rv is None:
                            return returnValues
                        returnValues.append(rv)
        
                    # cursor.close()
                    return returnValues
            except Exception as e:
                print(e) 

def _update_url_prioritization(conf):
    log.info("UPU: Started _update_url_prioritization")
    outdir = conf.output_directory / "url_prioritization"
    outdir.mkdir(parents=True, exist_ok=True)
    conn = connect_db(conf.standby)
    cur = conn.cursor(cursor_factory=RealDictCursor)

    log.info("UPU: Regenerating URL prioritization file")
    sql = """SELECT priority, domain, url, cc, category_code FROM citizenlab"""
    cur.execute(sql)
    entries = list(cur.fetchall())
    conn.rollback()
    conn.close()

    entries_by_country = {}  # cc -> [entry, ... ]
    for e in entries:
        cc = e["cc"]
        if cc not in entries_by_country:
            entries_by_country[cc] = []
        entries_by_country[cc].append(e)

    del entries

    k = 10

    for cc, entry in entries_by_country.items():
        candidates = entries_by_country[cc] + entries_by_country["ZZ"]
        selected = algo_chao(candidates, k)

        test_items = []
        for entry in selected:
            cc = entry["cc"]
            test_items.append(
                {
                    "category_code": entry["category_code"],
                    "url": entry["url"],
                    "country_code": "XX" if cc == "ZZ" else cc,
                }
            )

        out = {
            "metadata": {
                "count": len(test_items),
                "current_page": -1,
                "limit": -1,
                "next_url": "",
                "pages": 1,
            },
            "results": test_items,
        }
        f = outdir / ("url_list.%s.json" % cc)
        tmpf = f.with_suffix(".tmp")
        j = json.dumps(out, sort_keys=True)
        tmpf.write_text(j)
        tmpf.rename(f)

    log.info("UPU: Done") 

def _standby_status_query(self, instance, db_conn):
        """Status query that is executed on the standby node"""
        f_result = None
        result = {"fetch_time": get_iso_timestamp(), "connection": False}
        if not db_conn:
            db_conn = self._connect_to_db(instance, self.config["remote_conns"].get(instance))
            if not db_conn:
                return result
        try:
            phase = "querying status from"
            self.log.debug("%s %r", phase, instance)
            c = db_conn.cursor(cursor_factory=RealDictCursor)
            if db_conn.server_version >= 100000:
                fields = [
                    "now() AS db_time",
                    "pg_is_in_recovery()",
                    "pg_last_xact_replay_timestamp()",
                    "pg_last_wal_receive_lsn() AS pg_last_xlog_receive_location",
                    "pg_last_wal_replay_lsn() AS pg_last_xlog_replay_location",
                ]
            else:
                fields = [
                    "now() AS db_time",
                    "pg_is_in_recovery()",
                    "pg_last_xact_replay_timestamp()",
                    "pg_last_xlog_receive_location()",
                    "pg_last_xlog_replay_location()",
                ]
            query = "SELECT {}".format(", ".join(fields))
            c.execute(query)
            wait_select(c.connection)
            f_result = c.fetchone()
            if not f_result['pg_is_in_recovery']:
                # This is only run on masters to create txid traffic every db_poll_interval
                phase = "updating transaction on"
                self.log.debug("%s %r", phase, instance)
                # With pg_current_wal_lsn we simulate replay_location on the master
                # With txid_current we force a new transaction to occur every poll interval to ensure there's
                # a heartbeat for the replication lag.
                if db_conn.server_version >= 100000:
                    c.execute("SELECT txid_current(), pg_current_wal_lsn() AS pg_last_xlog_replay_location")
                else:
                    c.execute("SELECT txid_current(), pg_current_xlog_location() AS pg_last_xlog_replay_location")
                wait_select(c.connection)
                master_result = c.fetchone()
                f_result["pg_last_xlog_replay_location"] = master_result["pg_last_xlog_replay_location"]
        except (PglookoutTimeout, psycopg2.DatabaseError, psycopg2.InterfaceError, psycopg2.OperationalError) as ex:
            self.log.warning("%s (%s) %s %s", ex.__class__.__name__, str(ex).strip(), phase, instance)
            db_conn.close()
            self.db_conns[instance] = None
            # Return "no connection" result in case of any error. If we get an error for master server after the initial
            # query we'd end up returning completely invalid value for master's current position
            return result

        result.update(self._parse_status_query_result(f_result))
        return result