Python pandas.read_sql_query() Examples

def test_read_sql_iris(self):
        iris_frame = sql.read_sql_query(
            "SELECT * FROM iris", self.conn)
        self._check_iris_loaded_frame(iris_frame) 

def test_date_and_index(self):
        # Test case where same column appears in parse_date and index_col

        df = sql.read_sql_query("SELECT * FROM types_test_data", self.conn,
                                index_col='DateCol',
                                parse_dates=['DateCol', 'IntDateCol'])

        assert issubclass(df.index.dtype.type, np.datetime64)
        assert issubclass(df.IntDateCol.dtype.type, np.datetime64) 

def count_sources_for_bad_persistence():
    """
    Count the number of extracted objects for each id and look for fields
    with few objects, which are usually problems with the persistence mask
    """

    import pandas as pd
    from grizli.aws import db as grizli_db
    from grizli import utils
    engine = grizli_db.get_db_engine(echo=False)

    # Number of matches per field
    counts = pd.read_sql_query("select root, COUNT(root) as n from redshift_fit, photometry_apcorr where phot_root = p_root AND id = p_id AND bic_diff > 5 AND mag_auto < 24 group by root;", engine)

    counts = utils.GTable.from_pandas(counts)
    so = np.argsort(counts['n'])

    sh = """
    BUCKET=grizli-v
    root=j113812m1134

    aws s3 rm --recursive s3://grizli-v1/Pipeline/${root}/ --include "*"
    grism_run_single.sh ${root} --run_fine_alignment=True --extra_filters=g800l --bucket=grizli-v1 --preprocess_args.skip_single_optical_visits=True --mask_spikes=True --persistence_args.err_threshold=1
    """ 

def test_nan_fullcolumn(self):
        # full NaN column (numeric float column)
        df = DataFrame({'A': [0, 1, 2], 'B': [np.nan, np.nan, np.nan]})
        df.to_sql('test_nan', self.conn, index=False)

        # with read_table
        result = sql.read_sql_table('test_nan', self.conn)
        tm.assert_frame_equal(result, df)

        # with read_sql -> not type info from table -> stays None
        df['B'] = df['B'].astype('object')
        df['B'] = None
        result = sql.read_sql_query('SELECT * FROM test_nan', self.conn)
        tm.assert_frame_equal(result, df) 

def add_to_charge():

    engine = grizli_db.get_db_engine()

    p = pd.read_sql_query('select distinct p_root from photometry_apcorr', engine)
    f = pd.read_sql_query('select distinct field_root from charge_fields', engine)

    new_fields = []
    for root in p['p_root'].values:
        if root not in f['field_root'].values:
            print(root)
            new_fields.append(root)

    df = pd.DataFrame()
    df['field_root'] = new_fields
    df['comment'] = 'CANDELS'
    ix = df['field_root'] == 'j214224m4420'
    df['comment'][ix] = 'Rafelski UltraDeep'

    df.to_sql('charge_fields', engine, index=False, if_exists='append', method='multi') 

def test_datetime_time(self):
        # test support for datetime.time
        df = DataFrame([time(9, 0, 0), time(9, 1, 30)], columns=["a"])
        df.to_sql('test_time', self.conn, index=False)
        res = read_sql_table('test_time', self.conn)
        tm.assert_frame_equal(res, df)

        # GH8341
        # first, use the fallback to have the sqlite adapter put in place
        sqlite_conn = TestSQLiteFallback.connect()
        sql.to_sql(df, "test_time2", sqlite_conn, index=False)
        res = sql.read_sql_query("SELECT * FROM test_time2", sqlite_conn)
        ref = df.applymap(lambda _: _.strftime("%H:%M:%S.%f"))
        tm.assert_frame_equal(ref, res)  # check if adapter is in place
        # then test if sqlalchemy is unaffected by the sqlite adapter
        sql.to_sql(df, "test_time3", self.conn, index=False)
        if self.flavor == 'sqlite':
            res = sql.read_sql_query("SELECT * FROM test_time3", self.conn)
            ref = df.applymap(lambda _: _.strftime("%H:%M:%S.%f"))
            tm.assert_frame_equal(ref, res)
        res = sql.read_sql_table("test_time3", self.conn)
        tm.assert_frame_equal(df, res) 

def test_connectable_issue_example(self):
        # This tests the example raised in issue
        # https://github.com/pandas-dev/pandas/issues/10104

        def foo(connection):
            query = 'SELECT test_foo_data FROM test_foo_data'
            return sql.read_sql_query(query, con=connection)

        def bar(connection, data):
            data.to_sql(name='test_foo_data',
                        con=connection, if_exists='append')

        def main(connectable):
            with connectable.connect() as conn:
                with conn.begin():
                    foo_data = conn.run_callable(foo)
                    conn.run_callable(bar, foo_data)

        DataFrame({'test_foo_data': [0, 1, 2]}).to_sql(
            'test_foo_data', self.conn)
        main(self.conn) 

def test_temporary_table(self):
        test_data = u'Hello, World!'
        expected = DataFrame({'spam': [test_data]})
        Base = declarative.declarative_base()

        class Temporary(Base):
            __tablename__ = 'temp_test'
            __table_args__ = {'prefixes': ['TEMPORARY']}
            id = sqlalchemy.Column(sqlalchemy.Integer, primary_key=True)
            spam = sqlalchemy.Column(sqlalchemy.Unicode(30), nullable=False)

        Session = sa_session.sessionmaker(bind=self.conn)
        session = Session()
        with session.transaction:
            conn = session.connection()
            Temporary.__table__.create(conn)
            session.add(Temporary(spam=test_data))
            session.flush()
            df = sql.read_sql_query(
                sql=sqlalchemy.select([Temporary.spam]),
                con=conn,
            )

        tm.assert_frame_equal(df, expected) 

def test_temporary_table(self):
        test_data = u'Hello, World!'
        expected = DataFrame({'spam': [test_data]})
        Base = declarative.declarative_base()

        class Temporary(Base):
            __tablename__ = 'temp_test'
            __table_args__ = {'prefixes': ['TEMPORARY']}
            id = sqlalchemy.Column(sqlalchemy.Integer, primary_key=True)
            spam = sqlalchemy.Column(sqlalchemy.Unicode(30), nullable=False)

        Session = sa_session.sessionmaker(bind=self.conn)
        session = Session()
        with session.transaction:
            conn = session.connection()
            Temporary.__table__.create(conn)
            session.add(Temporary(spam=test_data))
            session.flush()
            df = sql.read_sql_query(
                sql=sqlalchemy.select([Temporary.spam]),
                con=conn,
            )

        tm.assert_frame_equal(df, expected) 

def test_connectable_issue_example(self):
        # This tests the example raised in issue
        # https://github.com/pandas-dev/pandas/issues/10104

        def foo(connection):
            query = 'SELECT test_foo_data FROM test_foo_data'
            return sql.read_sql_query(query, con=connection)

        def bar(connection, data):
            data.to_sql(name='test_foo_data',
                        con=connection, if_exists='append')

        def main(connectable):
            with connectable.connect() as conn:
                with conn.begin():
                    foo_data = conn.run_callable(foo)
                    conn.run_callable(bar, foo_data)

        DataFrame({'test_foo_data': [0, 1, 2]}).to_sql(
            'test_foo_data', self.conn)
        main(self.conn) 

def test_datetime_time(self):
        # test support for datetime.time
        df = DataFrame([time(9, 0, 0), time(9, 1, 30)], columns=["a"])
        df.to_sql('test_time', self.conn, index=False)
        res = read_sql_table('test_time', self.conn)
        tm.assert_frame_equal(res, df)

        # GH8341
        # first, use the fallback to have the sqlite adapter put in place
        sqlite_conn = TestSQLiteFallback.connect()
        sql.to_sql(df, "test_time2", sqlite_conn, index=False)
        res = sql.read_sql_query("SELECT * FROM test_time2", sqlite_conn)
        ref = df.applymap(lambda _: _.strftime("%H:%M:%S.%f"))
        tm.assert_frame_equal(ref, res)  # check if adapter is in place
        # then test if sqlalchemy is unaffected by the sqlite adapter
        sql.to_sql(df, "test_time3", self.conn, index=False)
        if self.flavor == 'sqlite':
            res = sql.read_sql_query("SELECT * FROM test_time3", self.conn)
            ref = df.applymap(lambda _: _.strftime("%H:%M:%S.%f"))
            tm.assert_frame_equal(ref, res)
        res = sql.read_sql_table("test_time3", self.conn)
        tm.assert_frame_equal(df, res) 

def test_datetime_NaT(self):
        df = DataFrame({'A': date_range('2013-01-01 09:00:00', periods=3),
                        'B': np.arange(3.0)})
        df.loc[1, 'A'] = np.nan
        df.to_sql('test_datetime', self.conn, index=False)

        # with read_table -> type information from schema used
        result = sql.read_sql_table('test_datetime', self.conn)
        tm.assert_frame_equal(result, df)

        # with read_sql -> no type information -> sqlite has no native
        result = sql.read_sql_query('SELECT * FROM test_datetime', self.conn)
        if self.flavor == 'sqlite':
            assert isinstance(result.loc[0, 'A'], string_types)
            result['A'] = to_datetime(result['A'], errors='coerce')
            tm.assert_frame_equal(result, df)
        else:
            tm.assert_frame_equal(result, df) 

def test_datetime(self):
        df = DataFrame({'A': date_range('2013-01-01 09:00:00', periods=3),
                        'B': np.arange(3.0)})
        df.to_sql('test_datetime', self.conn)

        # with read_table -> type information from schema used
        result = sql.read_sql_table('test_datetime', self.conn)
        result = result.drop('index', axis=1)
        tm.assert_frame_equal(result, df)

        # with read_sql -> no type information -> sqlite has no native
        result = sql.read_sql_query('SELECT * FROM test_datetime', self.conn)
        result = result.drop('index', axis=1)
        if self.flavor == 'sqlite':
            assert isinstance(result.loc[0, 'A'], string_types)
            result['A'] = to_datetime(result['A'])
            tm.assert_frame_equal(result, df)
        else:
            tm.assert_frame_equal(result, df) 

def wait_on_db_update(root, t0=60, dt=30, n_iter=60, engine=None):
    """
    Wait for db to stop updating on root
    """
    import pandas as pd
    from astropy.table import Table
    from grizli.aws import db as grizli_db
    import numpy as np
    import time

    if engine is None:
        engine = grizli_db.get_db_engine(echo=False)

    n_i, n6_i, checksum_i = -1, -1, -1

    for i in range(n_iter):
        res = pd.read_sql_query("SELECT root, id, status FROM redshift_fit WHERE root = '{0}'".format(root), engine)
        checksum = (2**res['status']).sum()
        n = len(res)
        n6 = (res['status'] == 6).sum()
        n5 = (res['status'] == 5).sum()
        if (n == n_i) & (checksum == checksum_i) & (n6 == n6_i):
            break

        now = time.ctime()
        print('{0}, {1}: n={2:<5d} n5={5:<5d} n6={3:<5d} checksum={4}'.format(root, now, n, n6, checksum, n5))
        n_i, n6_i, checksum_i = n, n6, checksum
        if i == 0:
            time.sleep(t0)
        else:
            time.sleep(dt)

    return res

## 

def add_missing_rows(root='j004404m2034', engine=None):
    """
    Add rows that were completed but that aren't in the table
    """
    import glob
    from astropy.table import vstack, Table

    from grizli.aws import db as grizli_db

    if engine is None:
        engine = grizli_db.get_db_engine(echo=False)

    os.system('aws s3 sync s3://grizli-v1/Pipeline/{0}/Extractions/ ./ --exclude "*" --include "*row.fits"'.format(root))

    row_files = glob.glob('{0}*row.fits'.format(root))
    row_files.sort()

    res = pd.read_sql_query("SELECT root, id, status FROM redshift_fit WHERE root = '{0}' AND status=6".format(root), engine)

    res_ids = res['id'].to_list()
    tabs = []

    print('\n\n NROWS={0}, NRES={1}\n\n'.format(len(row_files), len(res)))

    for row_file in row_files:
        id_i = int(row_file.split('.row.fits')[0][-5:])
        if id_i not in res_ids:
            grizli_db.add_redshift_fit_row(row_file, engine=engine, verbose=True) 

def test_datetime_date(self):
        # test support for datetime.date
        df = DataFrame([date(2014, 1, 1), date(2014, 1, 2)], columns=["a"])
        df.to_sql('test_date', self.conn, index=False)
        res = read_sql_query('SELECT * FROM test_date', self.conn)
        if self.flavor == 'sqlite':
            # comes back as strings
            tm.assert_frame_equal(res, df.astype(str))
        elif self.flavor == 'mysql':
            tm.assert_frame_equal(res, df) 

def test_read_procedure(self):
        # see GH7324. Although it is more an api test, it is added to the
        # mysql tests as sqlite does not have stored procedures
        df = DataFrame({'a': [1, 2, 3], 'b': [0.1, 0.2, 0.3]})
        df.to_sql('test_procedure', self.conn, index=False)

        proc = """DROP PROCEDURE IF EXISTS get_testdb;

        CREATE PROCEDURE get_testdb ()

        BEGIN
            SELECT * FROM test_procedure;
        END"""

        connection = self.conn.connect()
        trans = connection.begin()
        try:
            r1 = connection.execute(proc)  # noqa
            trans.commit()
        except:
            trans.rollback()
            raise

        res1 = sql.read_sql_query("CALL get_testdb();", self.conn)
        tm.assert_frame_equal(df, res1)

        # test delegation to read_sql_query
        res2 = sql.read_sql("CALL get_testdb();", self.conn)
        tm.assert_frame_equal(df, res2) 

def test_roundtrip_chunksize(self):
        sql.to_sql(self.test_frame1, 'test_frame_roundtrip', con=self.conn,
                   index=False, chunksize=2)
        result = sql.read_sql_query(
            'SELECT * FROM test_frame_roundtrip',
            con=self.conn)
        tm.assert_frame_equal(result, self.test_frame1) 

def column_comments():

    from collections import OrderedDict
    import yaml

    tablename = 'redshift_fit'

    cols = pd.read_sql_query('select * from {0} where false'.format(tablename), engine)

    d = {}  # OrderedDict{}
    for c in cols.columns:
        d[c] = '---'

    if not os.path.exists('{0}_comments.yml'.format(tablename)):
        print('Init {0}_comments.yml'.format(tablename))
        fp = open('{0}_comments.yml'.format(tablename), 'w')
        yaml.dump(d, stream=fp, default_flow_style=False)
        fp.close()

    # Edit file
    comments = yaml.load(open('{0}_comments.yml'.format(tablename)))
    SQL = ""
    upd = "COMMENT ON COLUMN {0}.{1} IS '{2}';\n"
    for col in comments:
        if comments[col] != '---':
            SQL += upd.format(tablename, col, comments[col])
        else:
            print('Skip ', col) 

def test_nan_string(self):
        # NaNs in string column
        df = DataFrame({'A': [0, 1, 2], 'B': ['a', 'b', np.nan]})
        df.to_sql('test_nan', self.conn, index=False)

        # NaNs are coming back as None
        df.loc[2, 'B'] = None

        # with read_table
        result = sql.read_sql_table('test_nan', self.conn)
        tm.assert_frame_equal(result, df)

        # with read_sql
        result = sql.read_sql_query('SELECT * FROM test_nan', self.conn)
        tm.assert_frame_equal(result, df) 

def test_nan_numeric(self):
        # NaNs in numeric float column
        df = DataFrame({'A': [0, 1, 2], 'B': [0.2, np.nan, 5.6]})
        df.to_sql('test_nan', self.conn, index=False)

        # with read_table
        result = sql.read_sql_table('test_nan', self.conn)
        tm.assert_frame_equal(result, df)

        # with read_sql
        result = sql.read_sql_query('SELECT * FROM test_nan', self.conn)
        tm.assert_frame_equal(result, df) 

def test_date_and_index(self):
        # Test case where same column appears in parse_date and index_col

        df = sql.read_sql_query("SELECT * FROM types_test_data", self.conn,
                                index_col='DateCol',
                                parse_dates=['DateCol', 'IntDateCol'])

        assert issubclass(df.index.dtype.type, np.datetime64)
        assert issubclass(df.IntDateCol.dtype.type, np.datetime64) 

def test_timedelta(self):

        # see #6921
        df = to_timedelta(
            Series(['00:00:01', '00:00:03'], name='foo')).to_frame()
        with tm.assert_produces_warning(UserWarning):
            df.to_sql('test_timedelta', self.conn)
        result = sql.read_sql_query('SELECT * FROM test_timedelta', self.conn)
        tm.assert_series_equal(result['foo'], df['foo'].astype('int64')) 

def test_to_sql_index_label(self):
        temp_frame = DataFrame({'col1': range(4)})

        # no index name, defaults to 'index'
        sql.to_sql(temp_frame, 'test_index_label', self.conn)
        frame = sql.read_sql_query('SELECT * FROM test_index_label', self.conn)
        assert frame.columns[0] == 'index'

        # specifying index_label
        sql.to_sql(temp_frame, 'test_index_label', self.conn,
                   if_exists='replace', index_label='other_label')
        frame = sql.read_sql_query('SELECT * FROM test_index_label', self.conn)
        assert frame.columns[0] == "other_label"

        # using the index name
        temp_frame.index.name = 'index_name'
        sql.to_sql(temp_frame, 'test_index_label', self.conn,
                   if_exists='replace')
        frame = sql.read_sql_query('SELECT * FROM test_index_label', self.conn)
        assert frame.columns[0] == "index_name"

        # has index name, but specifying index_label
        sql.to_sql(temp_frame, 'test_index_label', self.conn,
                   if_exists='replace', index_label='other_label')
        frame = sql.read_sql_query('SELECT * FROM test_index_label', self.conn)
        assert frame.columns[0] == "other_label"

        # index name is integer
        temp_frame.index.name = 0
        sql.to_sql(temp_frame, 'test_index_label', self.conn,
                   if_exists='replace')
        frame = sql.read_sql_query('SELECT * FROM test_index_label', self.conn)
        assert frame.columns[0] == "0"

        temp_frame.index.name = None
        sql.to_sql(temp_frame, 'test_index_label', self.conn,
                   if_exists='replace', index_label=0)
        frame = sql.read_sql_query('SELECT * FROM test_index_label', self.conn)
        assert frame.columns[0] == "0" 

def test_read_sql_delegate(self):
        iris_frame1 = sql.read_sql_query("SELECT * FROM iris", self.conn)
        iris_frame2 = sql.read_sql("SELECT * FROM iris", self.conn)
        tm.assert_frame_equal(iris_frame1, iris_frame2)

        pytest.raises(sql.DatabaseError, sql.read_sql, 'iris', self.conn) 

def test_database_uri_string(self):

        # Test read_sql and .to_sql method with a database URI (GH10654)
        test_frame1 = self.test_frame1
        # db_uri = 'sqlite:///:memory:' # raises
        # sqlalchemy.exc.OperationalError: (sqlite3.OperationalError) near
        # "iris": syntax error [SQL: 'iris']
        with tm.ensure_clean() as name:
            db_uri = 'sqlite:///' + name
            table = 'iris'
            test_frame1.to_sql(table, db_uri, if_exists='replace', index=False)
            test_frame2 = sql.read_sql(table, db_uri)
            test_frame3 = sql.read_sql_table(table, db_uri)
            query = 'SELECT * FROM iris'
            test_frame4 = sql.read_sql_query(query, db_uri)
        tm.assert_frame_equal(test_frame1, test_frame2)
        tm.assert_frame_equal(test_frame1, test_frame3)
        tm.assert_frame_equal(test_frame1, test_frame4)

        # using driver that will not be installed on Travis to trigger error
        # in sqlalchemy.create_engine -> test passing of this error to user
        try:
            # the rest of this test depends on pg8000's being absent
            import pg8000  # noqa
            pytest.skip("pg8000 is installed")
        except ImportError:
            pass

        db_uri = "postgresql+pg8000://user:pass@host/dbname"
        with tm.assert_raises_regex(ImportError, "pg8000"):
            sql.read_sql("select * from table", db_uri) 

def test_not_reflect_all_tables(self):
        # create invalid table
        qry = """CREATE TABLE invalid (x INTEGER, y UNKNOWN);"""
        self.conn.execute(qry)
        qry = """CREATE TABLE other_table (x INTEGER, y INTEGER);"""
        self.conn.execute(qry)

        with warnings.catch_warnings(record=True) as w:
            # Cause all warnings to always be triggered.
            warnings.simplefilter("always")
            # Trigger a warning.
            sql.read_sql_table('other_table', self.conn)
            sql.read_sql_query('SELECT * FROM other_table', self.conn)
            # Verify some things
            assert len(w) == 0 

def test_read_sql_delegate(self):
        iris_frame1 = sql.read_sql_query(
            "SELECT * FROM iris", self.conn)
        iris_frame2 = sql.read_sql(
            "SELECT * FROM iris", self.conn)
        tm.assert_frame_equal(iris_frame1, iris_frame2)

        iris_frame1 = sql.read_sql_table('iris', self.conn)
        iris_frame2 = sql.read_sql('iris', self.conn)
        tm.assert_frame_equal(iris_frame1, iris_frame2) 

def test_to_sql_index_label_multiindex(self):
        temp_frame = DataFrame({'col1': range(4)},
                               index=MultiIndex.from_product(
                                   [('A0', 'A1'), ('B0', 'B1')]))

        # no index name, defaults to 'level_0' and 'level_1'
        sql.to_sql(temp_frame, 'test_index_label', self.conn)
        frame = sql.read_sql_query('SELECT * FROM test_index_label', self.conn)
        assert frame.columns[0] == 'level_0'
        assert frame.columns[1] == 'level_1'

        # specifying index_label
        sql.to_sql(temp_frame, 'test_index_label', self.conn,
                   if_exists='replace', index_label=['A', 'B'])
        frame = sql.read_sql_query('SELECT * FROM test_index_label', self.conn)
        assert frame.columns[:2].tolist() == ['A', 'B']

        # using the index name
        temp_frame.index.names = ['A', 'B']
        sql.to_sql(temp_frame, 'test_index_label', self.conn,
                   if_exists='replace')
        frame = sql.read_sql_query('SELECT * FROM test_index_label', self.conn)
        assert frame.columns[:2].tolist() == ['A', 'B']

        # has index name, but specifying index_label
        sql.to_sql(temp_frame, 'test_index_label', self.conn,
                   if_exists='replace', index_label=['C', 'D'])
        frame = sql.read_sql_query('SELECT * FROM test_index_label', self.conn)
        assert frame.columns[:2].tolist() == ['C', 'D']

        # wrong length of index_label
        pytest.raises(ValueError, sql.to_sql, temp_frame,
                      'test_index_label', self.conn, if_exists='replace',
                      index_label='C') 

def _get_index_columns(self, tbl_name):
        ixs = sql.read_sql_query(
            "SELECT * FROM sqlite_master WHERE type = 'index' " +
            "AND tbl_name = '%s'" % tbl_name, self.conn)
        ix_cols = []
        for ix_name in ixs.name:
            ix_info = sql.read_sql_query(
                "PRAGMA index_info(%s)" % ix_name, self.conn)
            ix_cols.append(ix_info.name.tolist())
        return ix_cols