Python sqlalchemy.engine() Examples

def test_dml(engine, session, table_dml):
    # test insert
    engine.execute(table_dml.insert(ONE_ROW_CONTENTS_DML))
    result = table_dml.select().execute().fetchall()
    assert len(result) == 1

    # test update
    session.query(table_dml)\
        .filter(table_dml.c.string == 'test')\
        .update({'string': 'updated_row'}, synchronize_session=False)
    updated_result = table_dml.select().execute().fetchone()
    assert updated_result['test_pybigquery.sample_dml_string'] == 'updated_row'

    # test delete
    session.query(table_dml).filter(table_dml.c.string == 'updated_row').delete(synchronize_session=False)
    result = table_dml.select().execute().fetchall()
    assert len(result) == 0 

def test_reflect_table_include_columns(self, engine, conn):
        one_row_complex = Table("one_row_complex", MetaData(bind=engine))
        version = float(
            re.search(r"^([\d]+\.[\d]+)\..+", sqlalchemy.__version__).group(1)
        )
        if version <= 1.2:
            engine.dialect.reflecttable(
                conn, one_row_complex, include_columns=["col_int"], exclude_columns=[]
            )
        else:
            # https://docs.sqlalchemy.org/en/13/changelog/changelog_13.html#
            # change-64ac776996da1a5c3e3460b4c0f0b257
            engine.dialect.reflecttable(
                conn,
                one_row_complex,
                include_columns=["col_int"],
                exclude_columns=[],
                resolve_fks=True,
            )
        self.assertEqual(len(one_row_complex.c), 1)
        self.assertIsNotNone(one_row_complex.c.col_int)
        self.assertRaises(AttributeError, lambda: one_row_complex.c.col_tinyint) 

def test_get_column_type(self, engine, conn):
        dialect = engine.dialect
        self.assertEqual(dialect._get_column_type("boolean"), "boolean")
        self.assertEqual(dialect._get_column_type("tinyint"), "tinyint")
        self.assertEqual(dialect._get_column_type("smallint"), "smallint")
        self.assertEqual(dialect._get_column_type("integer"), "integer")
        self.assertEqual(dialect._get_column_type("bigint"), "bigint")
        self.assertEqual(dialect._get_column_type("real"), "real")
        self.assertEqual(dialect._get_column_type("double"), "double")
        self.assertEqual(dialect._get_column_type("varchar"), "varchar")
        self.assertEqual(dialect._get_column_type("timestamp"), "timestamp")
        self.assertEqual(dialect._get_column_type("date"), "date")
        self.assertEqual(dialect._get_column_type("varbinary"), "varbinary")
        self.assertEqual(dialect._get_column_type("array(integer)"), "array")
        self.assertEqual(dialect._get_column_type("map(integer, integer)"), "map")
        self.assertEqual(dialect._get_column_type("row(a integer, b integer)"), "row")
        self.assertEqual(dialect._get_column_type("decimal(10,1)"), "decimal") 

def test_contain_percents_character_query_with_parameter(self, engine, conn):
        query = sqlalchemy.sql.text(
            """
            SELECT date_parse('20191030', '%Y%m%d'), :word
            """
        )
        result = engine.execute(query, word="cat")
        self.assertEqual(result.fetchall(), [(datetime(2019, 10, 30), "cat")])

        query = sqlalchemy.sql.text(
            """
            SELECT col_string FROM one_row_complex
            WHERE col_string LIKE 'a%' OR col_string LIKE :param
            """
        )
        result = engine.execute(query, param="b%")
        self.assertEqual(result.fetchall(), [("a string",)]) 

def receive_before_cursor_execute(**kw):
        try:
            parent_span = tracer.active_span

            # If we're not tracing, just return
            if parent_span is None:
                return

            scope = tracer.start_active_span("sqlalchemy", child_of=parent_span)
            context = kw['context']
            context._stan_scope = scope

            conn = kw['conn']
            url = str(conn.engine.url)
            scope.span.set_tag('sqlalchemy.sql', kw['statement'])
            scope.span.set_tag('sqlalchemy.eng', conn.engine.name)
            scope.span.set_tag('sqlalchemy.url', url_regexp.sub('//', url))
        except Exception as e:
            logger.debug(e)
        finally:
            return 

def _instrument(self, **kwargs):
        """Instruments SQLAlchemy engine creation methods and the engine
        if passed as an argument.

        Args:
            **kwargs: Optional arguments
                ``engine``: a SQLAlchemy engine instance
                ``tracer_provider``: a TracerProvider, defaults to global
                ``service``: the name of the service to trace.

        Returns:
            An instrumented engine if passed in as an argument, None otherwise.
        """
        _w("sqlalchemy", "create_engine", _wrap_create_engine)
        _w("sqlalchemy.engine", "create_engine", _wrap_create_engine)
        if kwargs.get("engine") is not None:
            return EngineTracer(
                _get_tracer(
                    kwargs.get("engine"), kwargs.get("tracer_provider")
                ),
                kwargs.get("service"),
                kwargs.get("engine"),
            )
        return None 

def setup_method(self, datapath):
        super(_EngineToConnMixin, self).setup_method(datapath)
        engine = self.conn
        conn = engine.connect()
        self.__tx = conn.begin()
        self.pandasSQL = sql.SQLDatabase(conn)
        self.__engine = engine
        self.conn = conn

        yield

        self.__tx.rollback()
        self.conn.close()
        self.conn = self.__engine
        self.pandasSQL = sql.SQLDatabase(self.__engine)
        # XXX:
        # super(_EngineToConnMixin, self).teardown_method(method) 

def setup_method(self, load_iris_data):
        super(_EngineToConnMixin, self).load_test_data_and_sql()
        engine = self.conn
        conn = engine.connect()
        self.__tx = conn.begin()
        self.pandasSQL = sql.SQLDatabase(conn)
        self.__engine = engine
        self.conn = conn

        yield

        self.__tx.rollback()
        self.conn.close()
        self.conn = self.__engine
        self.pandasSQL = sql.SQLDatabase(self.__engine)
        # XXX:
        # super(_EngineToConnMixin, self).teardown_method(method) 

def build_filter(cls, engine, table, tree):
        try:
            operator, nodes = list(tree.items())[0]
        except Exception:
            raise indexer.QueryError()

        try:
            op = cls.multiple_operators[operator]
        except KeyError:
            try:
                op = cls.binary_operators[operator]
            except KeyError:
                try:
                    op = cls.unary_operators[operator]
                except KeyError:
                    raise indexer.QueryInvalidOperator(operator)
                return cls._handle_unary_op(engine, table, op, nodes)
            return cls._handle_binary_op(engine, table, op, nodes)
        return cls._handle_multiple_op(engine, table, op, nodes) 

def test_char_length(self, engine, conn):
        one_row_complex = Table("one_row_complex", MetaData(bind=engine), autoload=True)
        result = (
            sqlalchemy.select(
                [sqlalchemy.func.char_length(one_row_complex.c.col_string)]
            )
            .execute()
            .scalar()
        )
        self.assertEqual(result, len("a string")) 

def assertColumnNotExists(self, engine, table_name, column):
        self.assertFalse(oslodbutils.column_exists(engine, table_name, column),
                        'Column %s.%s should not exist' % (table_name, column)) 

def test_reserved_words(self, engine, conn):
        """Presto uses double quotes, not backticks"""
        fake_table = Table(
            "select", MetaData(bind=engine), Column("current_timestamp", STRINGTYPE)
        )
        query = str(fake_table.select(fake_table.c.current_timestamp == "a"))
        self.assertIn('"select"', query)
        self.assertIn('"current_timestamp"', query)
        self.assertNotIn("`select`", query)
        self.assertNotIn("`current_timestamp`", query) 

def test_contain_percents_character_query(self, engine, conn):
        query = sqlalchemy.sql.text(
            """
            SELECT date_parse('20191030', '%Y%m%d')
            """
        )
        result = engine.execute(query)
        self.assertEqual(result.fetchall(), [(datetime(2019, 10, 30),)]) 

def test_query_with_parameter(self, engine, conn):
        query = sqlalchemy.sql.text(
            """
            SELECT :word
            """
        )
        result = engine.execute(query, word="cat")
        self.assertEqual(result.fetchall(), [("cat",)]) 

def assertIndexExists(self, engine, table_name, index):
        self.assertTrue(oslodbutils.index_exists(engine, table_name, index),
                        'Index %s on table %s does not exist' %
                        (index, table_name)) 

def create_db_engine(
    dialect: str,
    driver: str,
    username: str,
    password: str,
    host: str,
    dbname: str,
):
    """Create SQLAlchemy DB Engine to connect Pandas <-> DB.

    Function that creates the engine object to connect to the database. The
    object is required by the pandas functions to_sql and from_sql

    :param dialect: Dialect for the engine (oracle, mysql, postgresql, etc)
    :param driver: DBAPI driver (psycopg2, ...)
    :param username: Username to connect with the database
    :param password: Password to connect with the database
    :param host: Host to connect with the database
    :param dbname: database name
    :return: the engine
    """
    database_url = '{dial}{drv}://{usr}:{pwd}@{h}/{dbname}'.format(
        dial=dialect,
        drv=driver,
        usr=username,
        pwd=password,
        h=host,
        dbname=dbname,
    )

    if settings.DEBUG:
        LOGGER.debug('Creating engine: %s', database_url)

    return sqlalchemy.create_engine(
        database_url,
        client_encoding=str('utf8'),
        encoding=str('utf8'),
        echo=False,
        paramstyle='format') 

def destroy_db_engine(db_engine=None):
    """Destroys the DB SQAlchemy engine.

    :param db_engine: Engine to destroy
    :return: Nothing
    """
    if db_engine:
        db_engine.dispose()
    else:
        if getattr(OnTaskSharedState, 'engine', None):
            OnTaskSharedState.engine.dispose() 

def load_table(
    table_name: str,
    columns: Optional[List[str]] = None,
    filter_exp: Optional[Dict] = None,
) -> Optional[pd.DataFrame]:
    """Load a Pandas data frame from the SQL DB.

    :param table_name: Table name
    :param columns: Optional list of columns to load (all if NOne is given)
    :param filter_exp: JSON expression to filter a subset of rows
    :return: data frame
    """
    if table_name not in connection.introspection.table_names():
        return None

    if settings.DEBUG:
        LOGGER.debug('Loading table %s', table_name)

    if columns or filter_exp:
        # A list of columns or a filter exp is given
        query, query_fields = sql.get_select_query_txt(
            table_name,
            column_names=columns,
            filter_formula=filter_exp)
        return pd.read_sql_query(
            query,
            OnTaskSharedState.engine,
            params=query_fields)

    # No special fields given, load the whole thing
    return pd.read_sql_table(table_name, OnTaskSharedState.engine) 

def assertColumnExists(self, engine, table_name, column):
        self.assertTrue(oslodbutils.column_exists(engine, table_name, column),
                        'Column %s.%s does not exist' % (table_name, column)) 

def test_get_columns(self, engine, conn):
        insp = sqlalchemy.inspect(engine)
        actual = insp.get_columns(table_name="one_row", schema=SCHEMA)[0]
        self.assertEqual(actual["name"], "number_of_rows")
        self.assertTrue(isinstance(actual["type"], INTEGER))
        self.assertTrue(actual["nullable"])
        self.assertIsNone(actual["default"])
        self.assertEqual(actual["ordinal_position"], 1)
        self.assertIsNone(actual["comment"]) 

def test_has_table(self, engine, conn):
        self.assertTrue(Table("one_row", MetaData(bind=engine)).exists())
        self.assertFalse(
            Table("this_table_does_not_exist", MetaData(bind=engine)).exists()
        ) 

def test_get_table_names(self, engine, conn):
        meta = MetaData()
        meta.reflect(bind=engine)
        print(meta.tables)
        self.assertIn("one_row", meta.tables)
        self.assertIn("one_row_complex", meta.tables)

        insp = sqlalchemy.inspect(engine)
        self.assertIn(
            "many_rows", insp.get_table_names(schema=SCHEMA),
        ) 

def test_reflect_schemas(self, engine, conn):
        insp = sqlalchemy.inspect(engine)
        schemas = insp.get_schema_names()
        self.assertIn(SCHEMA, schemas)
        self.assertIn("default", schemas) 

def _check_001(self, engine, data):
        self.assertColumnExists(engine, 'failover_segments', 'uuid')
        self.assertColumnExists(engine, 'failover_segments', 'name')
        self.assertColumnExists(engine, 'failover_segments', 'service_type')
        self.assertColumnExists(engine, 'failover_segments', 'description')
        self.assertColumnExists(engine, 'failover_segments',
                                'recovery_method')
        self.assertIndexMembers(engine, 'failover_segments',
                                'segments_service_type_idx', ['service_type']) 

def test_reflect_table_with_schema(self, engine, conn):
        one_row = Table("one_row", MetaData(bind=engine), schema=SCHEMA, autoload=True)
        self.assertEqual(len(one_row.c), 1)
        self.assertIsNotNone(one_row.c.number_of_rows) 

def test_reflect_table(self, engine, conn):
        one_row = Table("one_row", MetaData(bind=engine), autoload=True)
        self.assertEqual(len(one_row.c), 1)
        self.assertIsNotNone(one_row.c.number_of_rows) 

def test_reflect_no_such_table(self, engine, conn):
        self.assertRaises(
            NoSuchTableError,
            lambda: Table("this_does_not_exist", MetaData(bind=engine), autoload=True),
        )
        self.assertRaises(
            NoSuchTableError,
            lambda: Table(
                "this_does_not_exist",
                MetaData(bind=engine),
                schema="also_does_not_exist",
                autoload=True,
            ),
        ) 

def test_basic_query(self, engine, conn):
        rows = conn.execute("SELECT * FROM one_row").fetchall()
        self.assertEqual(len(rows), 1)
        self.assertEqual(rows[0].number_of_rows, 1)
        self.assertEqual(len(rows[0]), 1) 

def _check_002(self, engine, data):
        self.assertColumnExists(engine, 'hosts', 'uuid')
        self.assertColumnExists(engine, 'hosts', 'name')
        self.assertColumnExists(engine, 'hosts', 'reserved')
        self.assertColumnExists(engine, 'hosts', 'type')
        self.assertColumnExists(engine, 'hosts', 'control_attributes')
        self.assertColumnExists(engine, 'hosts', 'failover_segment_id')
        self.assertColumnExists(engine, 'hosts', 'on_maintenance')
        self.assertColumnExists(engine, 'hosts', 'type')
        self.assertIndexMembers(engine, 'hosts', 'hosts_type_idx', ['type']) 

def db_sync(self, engine):
        sa_migration.db_sync(engine=self.migrate_engine)