Python psycopg2.extensions.cursor() Examples

def as_string(self, context):
        # is it a connection or cursor?
        if isinstance(context, ext.connection):
            conn = context
        elif isinstance(context, ext.cursor):
            conn = context.connection
        else:
            raise TypeError("context must be a connection or a cursor")

        a = ext.adapt(self._wrapped)
        if hasattr(a, 'prepare'):
            a.prepare(conn)

        rv = a.getquoted()
        if sys.version_info[0] >= 3 and isinstance(rv, bytes):
            rv = rv.decode(ext.encodings[conn.encoding])

        return rv 

def _pipeDbCursorToTsvFile(self, dbCursor, tsvFile, include_columns=True):
        """
        Pipe any arbitrary DB cursor to a TSV file.
        """
        # Extract DB columns.
        columns = dbCursor.description
        numColumns = len(columns)

        if include_columns:
            # Write TSV header.
            for i in range(numColumns - 1):
                # 0th index is column name.
                tsvFile.write("%s\t" % columns[i][0])
            tsvFile.write("%s\n" % columns[numColumns - 1][0])

        # By default, cursor iterates through both header and data rows.
        self._numRows = 0
        row = dbCursor.fetchone()
        while row is not None:
            for i in range(numColumns - 1):
                tsvFile.write("%s\t" % row[i])
            tsvFile.write("%s\n" % row[numColumns - 1])
            row = dbCursor.fetchone()
            self._numRows += 1 

def _processPatientEpisodeDbCursor(self):
        """
        Convert self.patientEpisodeInput from DB cursor to TSV file.
        """
        # Instantiate patientEpisodeTempFile.
        patientEpisodeTempFile = open(self._patientEpisodeTempFileName, "w")
        self._pipeDbCursorToTsvFile(self.patientEpisodeInput, patientEpisodeTempFile)
        patientEpisodeTempFile.close()
        self.patientsProcessed = True

        if LocalEnv.DATABASE_CONNECTOR_NAME == 'psycopg2':
            return self.patientEpisodeInput.rowcount
        elif LocalEnv.DATABASE_CONNECTOR_NAME == 'sqlite3':
        # In sqlite3, rowcount is somehow "always" -1; See for details:
        # https://docs.python.org/3.0/library/sqlite3.html#sqlite3.Cursor.rowcount
            return self._numRows #self.patientEpisodeInput.rowcount 

def test_with_error(self):
        try:
            with self.conn as conn:
                with conn.cursor() as curs:
                    curs.execute("insert into test_with values (5)")
                    1/0
        except ZeroDivisionError:
            pass

        self.assertEqual(self.conn.status, ext.STATUS_READY)
        self.assert_(not self.conn.closed)
        self.assert_(curs.closed)

        curs = self.conn.cursor()
        curs.execute("select * from test_with")
        self.assertEqual(curs.fetchall(), []) 

def test_subclass_rollback(self):
        rollbacks = []
        class MyConn(ext.connection):
            def rollback(self):
                rollbacks.append(None)
                super(MyConn, self).rollback()

        try:
            with self.connect(connection_factory=MyConn) as conn:
                curs = conn.cursor()
                curs.execute("insert into test_with values (11)")
                1/0
        except ZeroDivisionError:
            pass
        else:
            self.assert_("exception not raised")

        self.assertEqual(conn.status, ext.STATUS_READY)
        self.assert_(rollbacks)

        curs = conn.cursor()
        curs.execute("select * from test_with")
        self.assertEqual(curs.fetchall(), []) 

def _pipeDbCursorToTsvFile(self, dbCursor, tsvFile):
        """
        Pipe any arbitrary DB cursor to a TSV file.
        """
        # Extract DB columns.
        columns = dbCursor.description
        numColumns = len(columns)

        # Write TSV header.
        for i in range(numColumns - 1):
            # 0th index is column name.
            tsvFile.write("%s\t" % columns[i][0])
        tsvFile.write("%s\n" % columns[numColumns - 1][0])

        # By default, cursor iterates through both header and data rows.
        self._numRows = 0
        row = dbCursor.fetchone()
        while row is not None:
            for i in range(numColumns - 1):
                tsvFile.write("%s\t" % row[i])
            tsvFile.write("%s\n" % row[numColumns - 1])
            row = dbCursor.fetchone()
            self._numRows += 1 

def setPatientEpisodeInput(self, patientEpisodeInput, \
        patientIdColumn = "patient_id", timestampColumn = "item_date"):
        """
        Define the input patient episode list for the feature matrix.
        patientEpisodeInput: TSV file descriptor or DB cursor.
        patientIdColumn: Name of TSV column or DB column.
        timestampColumn: Name of TSV column or DB column.
        """
        # Verify patientEpisodeInput is TSV file or DB cursor.
        if not isinstance(patientEpisodeInput, cursor) and \
            not isinstance(patientEpisodeInput, file):
            raise TypeError("patientEpisodeInput must be DB cursor or TSV file.")

        self.patientEpisodeInput = patientEpisodeInput
        self.patientEpisodeIdColumn = patientIdColumn
        self.patientEpisodeTimeColumn = timestampColumn 

def _processPatientListTsvFile(self):
        """
        Convert self.patientListInput from DB cursor to TSV file.
        """
        # Instantiate patientListTempFile.
        patientListTempFile = open(self._patientListTempFileName, "w")
        patientListTempFile.write("%s\n" % self.patientIdColumn)

        # Iterate through all rows in TSV file.
        # Extract patientId from dictionary.
        for row in TabDictReader(self.patientListInput):
            patientId = int(row[self.patientIdColumn])
            patientListTempFile.write("%s\n" % patientId)

        patientListTempFile.close() 

def _processPatientListDbCursor(self):
        """
        Convert self.patientListInput from DB cursor to TSV file.
        """
        patientListTempFile = open(self._patientListTempFileName, "w")
        self._pipeDbCursorToTsvFile(self.patientListInput, patientListTempFile)
        patientListTempFile.close() 

def processPatientListInput(self):
        """
        Convert patient list input to a TSV file.
        """
        if self.patientListInput is None:
            raise ValueError("FeatureMatrixFactory.patientListInput is None.")

        if isinstance(self.patientListInput, cursor):
            return self._processPatientListDbCursor()
        elif isinstance(self.patientListInput, file):
            return self._processPatientListTsvFile() 

def processClinicalItemInput(self, clinicalItemName):
        """
        Process clinicalItemName for DB cursor.
        """
        patientEpisodes = self._readPatientEpisodesFile()
        pass 

def _processPatientEpisodeDbCursor(self):
        """
        Convert self.patientEpisodeInput from DB cursor to TSV file.
        """
        # Instantiate patientEpisodeTempFile.
        patientEpisodeTempFile = open(self._patientEpisodeTempFileName, "w")
        self._pipeDbCursorToTsvFile(self.patientEpisodeInput, patientEpisodeTempFile)
        patientEpisodeTempFile.close()
        self.patientsProcessed = True

        return self.patientEpisodeInput.rowcount 

def _queryMichiganItemsByCategory(self, label, tableName): #
        """
        Query for all patient items that match with the given clinical item
        category ID.
        """
        # Identify which columns to pull from patient_item table.

        # return sth like
        # clinicalItemEvents=[[-3384542270496665494, u'2009-07-07 13:00:00'], [1262980084096039344, u'2003-01-22 12:29:00'], ...]
        self._patientItemIdColumn = "pat_id"
        self._patientItemTimeColumn = label

        patientIds = set()
        patientEpisodes = self.getPatientEpisodeIterator()
        for episode in patientEpisodes:
            patientIds.add(episode[self.patientEpisodeIdColumn])

        query_str = "SELECT %s, %s " % (self._patientItemIdColumn, self._patientItemTimeColumn)

        query_str += " FROM %s " % tableName

        query_str += "WHERE pat_id IN "
        pat_list_str = "("
        for pat_id in patientIds:
            pat_list_str += str(pat_id) + ","
        pat_list_str = pat_list_str[:-1] + ") "
        query_str += pat_list_str

        query_str += "ORDER BY pat_id, %s " % label

        results = DBUtil.connection().cursor().execute(query_str).fetchall()

        clinicalItemEvents = [list(row) for row in results]
        return clinicalItemEvents 

def as_string(self, context):
        """
        Return the string value of the object.

        :param context: the context to evaluate the string into.
        :type context: `connection` or `cursor`

        The method is automatically invoked by `~cursor.execute()` and
        `~cursor.executemany()` if a `!Composable` is passed instead of the
        query string.
        """
        raise NotImplementedError 

def processPatientEpisodeInput(self):
        """
        Convert patient episode input to a TSV file.
        """
        if self.patientEpisodeInput is None:
            raise ValueError("FeatureMatrixFactory.patientEpisodeInput is None.")

        if isinstance(self.patientEpisodeInput, cursor):
            return self._processPatientEpisodeDbCursor()
        elif isinstance(self.patientEpisodeInput, IOBase):
            return self._processPatientEpisodeTsvFile() 

def _processPatientListTsvFile(self):
        """
        Convert self.patientListInput from DB cursor to TSV file.
        """
        # Instantiate patientListTempFile.
        patientListTempFile = open(self._patientListTempFileName, "w")
        patientListTempFile.write("%s\n" % self.patientIdColumn)

        # Iterate through all rows in TSV file.
        # Extract patientId from dictionary.
        for row in TabDictReader(self.patientListInput):
            patientId = int(row[self.patientIdColumn])
            patientListTempFile.write("%s\n" % patientId)

        patientListTempFile.close() 

def processClinicalItemInput(self, clinicalItemName):
        """
        Process clinicalItemName for DB cursor.
        """
        patientEpisodes = self._readPatientEpisodesFile()
        pass 

def _processPatientListDbCursor(self):
        """
        Convert self.patientListInput from DB cursor to TSV file.
        """
        patientListTempFile = open(self._patientListTempFileName, "w")
        self._pipeDbCursorToTsvFile(self.patientListInput, patientListTempFile)
        patientListTempFile.close() 

def processPatientListInput(self):
        """
        Convert patient list input to a TSV file.
        """
        if self.patientListInput is None:
            raise ValueError("FeatureMatrixFactory.patientListInput is None.")

        if isinstance(self.patientListInput, cursor):
            return self._processPatientListDbCursor()
        elif isinstance(self.patientListInput, IOBase):
            return self._processPatientListTsvFile() 

def setPatientListInput(self, patientListInput, patientIdColumn = "patient_id"):
        """
        Define the input patient list for the feature matrix.
        patientListInput: TSV file descriptor or DB cursor
        patientIdColumn: Name of TSV column or DB column.
        """
        # Verify patientListInput is TSV file or DB cursor.
        if not isinstance(patientListInput, cursor) and \
                not isinstance(patientListInput, IOBase):
            raise TypeError("patientListInput must be DB cursor or TSV file.")

        self.patientListInput = patientListInput
        self.patientIdColumn = patientIdColumn
        pass 

def test_exception_swallow(self):
        # bug #262: __exit__ calls cur.close() that hides the exception
        # with another error.
        try:
            with self.conn as conn:
                with conn.cursor('named') as cur:
                    cur.execute("select 1/0")
                    cur.fetchone()
        except psycopg2.DataError, e:
            self.assertEqual(e.pgcode, '22012') 

def test_subclass(self):
        closes = []
        class MyCurs(ext.cursor):
            def close(self):
                closes.append(None)
                super(MyCurs, self).close()

        with self.conn.cursor(cursor_factory=MyCurs) as curs:
            self.assert_(isinstance(curs, MyCurs))

        self.assert_(curs.closed)
        self.assert_(closes) 

def test_with_ok(self):
        with self.conn as conn:
            with conn.cursor() as curs:
                curs.execute("insert into test_with values (4)")
                self.assert_(not curs.closed)
            self.assertEqual(self.conn.status, ext.STATUS_BEGIN)
            self.assert_(curs.closed)

        self.assertEqual(self.conn.status, ext.STATUS_READY)
        self.assert_(not self.conn.closed)

        curs = self.conn.cursor()
        curs.execute("select * from test_with")
        self.assertEqual(curs.fetchall(), [(4,)]) 

def test_with_error_python(self):
        def f():
            with self.conn as conn:
                curs = conn.cursor()
                curs.execute("insert into test_with values (3)")
                1/0

        self.assertRaises(ZeroDivisionError, f)
        self.assertEqual(self.conn.status, ext.STATUS_READY)
        self.assert_(not self.conn.closed)

        curs = self.conn.cursor()
        curs.execute("select * from test_with")
        self.assertEqual(curs.fetchall(), []) 

def test_with_error_db(self):
        def f():
            with self.conn as conn:
                curs = conn.cursor()
                curs.execute("insert into test_with values ('a')")

        self.assertRaises(psycopg2.DataError, f)
        self.assertEqual(self.conn.status, ext.STATUS_READY)
        self.assert_(not self.conn.closed)

        curs = self.conn.cursor()
        curs.execute("select * from test_with")
        self.assertEqual(curs.fetchall(), []) 

def test_with_connect_idiom(self):
        with self.connect() as conn:
            self.assertEqual(conn.status, ext.STATUS_READY)
            curs = conn.cursor()
            curs.execute("insert into test_with values (2)")
            self.assertEqual(conn.status, ext.STATUS_BEGIN)

        self.assertEqual(self.conn.status, ext.STATUS_READY)
        self.assert_(not self.conn.closed)

        curs = self.conn.cursor()
        curs.execute("select * from test_with")
        self.assertEqual(curs.fetchall(), [(2,)]) 

def test_with_ok(self):
        with self.conn as conn:
            self.assert_(self.conn is conn)
            self.assertEqual(conn.status, ext.STATUS_READY)
            curs = conn.cursor()
            curs.execute("insert into test_with values (1)")
            self.assertEqual(conn.status, ext.STATUS_BEGIN)

        self.assertEqual(self.conn.status, ext.STATUS_READY)
        self.assert_(not self.conn.closed)

        curs = self.conn.cursor()
        curs.execute("select * from test_with")
        self.assertEqual(curs.fetchall(), [(1,)]) 

def setUp(self):
        ConnectingTestCase.setUp(self)
        curs = self.conn.cursor()
        try:
            curs.execute("delete from test_with")
            self.conn.commit()
        except psycopg2.ProgrammingError:
            # assume table doesn't exist
            self.conn.rollback()
            curs.execute("create table test_with (id integer primary key)")
            self.conn.commit() 

def cursor(self, *args: Any, **kwargs: Any) -> TimeTrackingCursor:
        kwargs.setdefault('cursor_factory', TimeTrackingCursor)
        return connection.cursor(self, *args, **kwargs) 

def _solve_conn_curs(conn_or_curs):
    """Return the connection and a DBAPI cursor from a connection or cursor."""
    if conn_or_curs is None:
        raise psycopg2.ProgrammingError("no connection or cursor provided")

    if hasattr(conn_or_curs, 'execute'):
        conn = conn_or_curs.connection
        curs = conn.cursor(cursor_factory=_cursor)
    else:
        conn = conn_or_curs
        curs = conn.cursor(cursor_factory=_cursor)

    return conn, curs