Query Engine

To use DataRobot’s Query Engine, ensure you have installed the datarobot[query-engine] package extra.

class datarobot._experimental.query_engine.engine.QueryEngine

Execute SQL statements against a database through DataRobot. Supports statements as strings or SQLAlchemy constructs.

Parameters:

connection_manager (IConnectionManager) – The connection manager to use to execute statements against a database.
dialect (Optional[Dialect]) – The SQL dialect to use to compile statements. Modifies how SQLAlchemy constructs are compiled to their database-specific SQL strings.
paramstyle (Optional[str]) – The parameter style used to bind parameters to the statement. By default, “qmark” is used, which uses ? placeholders for parameters. Named parameters will be substituted in the SQL string according to the parameter name.
**kwargs – Additional keyword arguments for future-proofing.

Examples

Execute a query against a DataStore using QueryEngine:

>>> from datarobot._experimental.query_engine import QueryEngine
>>> engine = QueryEngine.from_data_store(
...     data_store_id="my_data_store_id",
...     credential_id="my_credential_id",
... )
>>> result: IteratorResult = engine.execute(
...     "SELECT * FROM my_table WHERE name = :name AND status IN :statuses",
...     params={
...         "name": "John Doe",
...         "statuses": ["active", "pending"]
...     }
... )
>>> result.all()
[(1, "John Doe", "active"), (2, "Jane Doe", "pending")]

Execute an update against a MS SQL Server database through a JDBC connection using QueryEngine and SQLAlchemy constructs:

>>> from sqlalchemy import insert, table, bindparam, column
>>> from sqlalchemy.dialects import mssql
>>> USER_TABLE = table("users", column("name"), column("status"))
>>> engine = QueryEngine.from_jdbc_connection(
...     jdbc_url="jdbc:sqlserver://localhost:1433;databaseName=mydb",
...     jdbc_params={"user": "sa", "password": "myPassword"},
...     dialect=mssql.dialect(),
... )
>>> engine.execute(
...     insert(USER_TABLE).values(name="John Doe", status=bindparam("status")),
...     params={"status": "active"},
... )

execute(stmt, params=None, *, max_rows=1000, mode=None, **kwargs)

Execute a SQL statement against a database. Supports string statements and SQLAlchemy constructs. Supports named parameters only.

Uses best-efforts to determine if the statement will return rows. Use mode to override this behavior. No results are returned for non-query statements.

Notes

If a parameter is a list or tuple, it will always be expanded. Replacement of a single parameter with a list or tuple is not supported. See examples below for more details.

Parameters:

stmt (Union[str, Executable]) – The SQL statement to execute. Supports string statements and SQLAlchemy constructs.
params (Optional[Dict[str, Union[str, int, float, bool, datetime, date, None, List[Union[str, int, float, bool, datetime, date, None]], Tuple[Union[str, int, float, bool, datetime, date, None], ...]]]]) – Named parameters to bind to the statement. Supports scalar, list, and tuple values.
max_rows (int) – The maximum number of rows to return. Only used for query-type statements.
mode (Optional[QueryMode]) – The mode to execute the statement. Overrides best-efforts to determine the mode.
**kwargs (Any) – Additional keyword arguments for future-proofing.

Returns:

The result of the statement. If the statement is a query, returns an IteratorResult with the result of the query. If the statement is an update, returns an IteratorResult with an empty result.

Return type:

sqlalchemy.engine.result.IteratorResult

Examples

Execute plain SQL string:

>>> from datarobot._experimental.query_engine import QueryEngine
>>> engine = QueryEngine.from_jdbc_connection(jdbc_url="jdbc:postgresql://localhost:5432/mydb")
>>> results = engine.execute("SELECT * FROM users")
>>> results.all()
[(1, "John Doe")]

Execute SQL query with named parameters:

>>> engine.execute("SELECT * FROM users WHERE name = :name", params={"name": "John Doe"})
>>> # Compiles to: SELECT * FROM users WHERE name = ?

>>> results.all()
[(1, "John Doe")]

Execute SQL query with named parameter that will be expanded. Note the expansion of the age parameter to (?, ?, ?):

>>> engine.execute(
...     "SELECT * FROM users WHERE name = :name AND age IN :ages",
...     params={"name": "John Doe", "ages": (30, 40, 50)},
... )
>>> # Compiles to: SELECT * FROM users WHERE name = ? AND age IN (?, ?, ?)
>>> results.all()
[(1, "John Doe", 30), (1, "John Doe", 40), (1, "John Doe", 50)]

Execute SQL statement to insert record with named parameter and parameter that will be expanded:

>>> engine.execute(
...     "INSERT INTO users (name, brothers) VALUES (:name, :brother_names)",
...     params={
...         "name": "John Doe",
...         "brother_names": ["Jim Doe", "Jack Doe"]
...     },
... )
>>> # Compiles to: INSERT INTO users (name, brothers) VALUES (?, (?, ?))

Execute SQLAlchemy select with named and bound parameters:

>>> from sqlalchemy import select, bindparam, column, table
>>> USER_TABLE = table("users", column("name"), column("status"))
>>> results = engine.execute(
...     select(USER_TABLE)
...         .where(USER_TABLE.c.name == "John Doe")
...         .where(USER_TABLE.c.status == bindparam("status")),
...     params={"status": "active"},
... )
>>> results.all()
[("John Doe", "active")]

Execute SQLAlchemy insert statement with bound parameter:

>>> from sqlalchemy import insert
>>> USER_TABLE = table("users", column("name"), column("status"))
>>> results = engine.execute(
...     insert(USER_TABLE).values(name="John Doe", status=bindparam("status")),
...     params={"status": "active"},
... )

classmethod from_jdbc_connection(jdbc_url=None, jdbc_url_generator=None, jdbc_params=None, dialect=None, paramstyle=None, **kwargs)

Create a QueryEngine from credentials for a JDBC database connection.

Parameters:

jdbc_url (Optional[str]) – The JDBC URL of the database.
jdbc_params (Optional[Dict[str, str]]) – The JDBC parameters to use for the connection.
jdbc_url_generator (Optional[Callable[[], str]]) – A function that returns a JDBC URL. Used to generate a JDBC URL for each connection if required.
**kwargs (Any) – Additional keyword arguments to pass to the QueryEngine constructor.
dialect (sqlalchemy.dialects.Dialect) – The SQL dialect to use to compile statements. Modifies how SQLAlchemy constructs are compiled to their database-specific SQL strings.
paramstyle (str) – The parameter style used to bind parameters to the statement. By default, “qmark” is used, which uses ? placeholders for parameters. Named parameters will be substituted in the SQL string according to the parameter name.

Return type:

QueryEngine

Examples

Create a QueryEngine from a JDBC URL:

>>> from datarobot._experimental.query_engine import QueryEngine
>>> engine = QueryEngine.from_jdbc_connection(
...     jdbc_url="jdbc:postgresql://localhost:5432/mydb",
...     jdbc_params={"user": "postgres", "password": "postgres"},
... )

Create a QueryEngine with a JDBC URL that has to be generated dynamically:

>>> engine = QueryEngine.from_jdbc_connection(
...     jdbc_url_generator=my_function_here",
...     jdbc_params={"user": "postgres", "password": "postgres"},
... )

classmethod from_data_store(data_store_id, credential_id=None, dialect=None, paramstyle=None, **kwargs)

Create a QueryEngine for a DataStore database connection.

Notes

Not all DataStores support statement execution through QueryEngine (e.g. Blob Storage).

Parameters:

data_store_id (str) – The ID of the DataStore to use for the connection.
credential_id (Optional[str]) – The ID of the credential to use for the connection. If not provided, the default credential for the DataStore will be used.
**kwargs (Any) – Additional keyword arguments to pass to the QueryEngine constructor.
dialect (sqlalchemy.dialects.Dialect) – The SQL dialect to use to compile statements. Modifies how SQLAlchemy constructs are compiled to their database-specific SQL strings.
paramstyle (str) – The parameter style used to bind parameters to the statement. By default, “qmark” is used, which uses ? placeholders for parameters. Named parameters will be substituted in the SQL string according to the parameter name.

Return type:

QueryEngine

Examples

Create a QueryEngine for a DataStore:

>>> from datarobot._experimental.query_engine import QueryEngine
>>> engine = QueryEngine.from_data_store(
...     data_store_id="my_data_store_id",
...     credential_id="my_credential_id"
... )

Enums and Helpers

class datarobot._experimental.query_engine.engine.QueryMode

Mode with which to execute a SQL statement. Can be used to override the best-efforts heuristic for determining the mode.

Examples

Override QueryEngine’s guess at determining the mode to ensure you get results back:

>>> from datarobot._experimental.query_engine import QueryMode, QueryEngine
>>> engine = QueryEngine.from_jdbc_connection(jdbc_url="jdbc:postgresql://localhost:5432/mydb")
>>> results = engine.execute(
...     "UPDATE users SET status = 'active' WHERE name = 'John Doe' RETURNING id, name, status",
...     mode=QueryMode.QUERY, # overrides QueryEngine's guess of the mode to ensure you get results back
... )
>>> results.all()
[(1, "John Doe", "active")]

QUERY = 'query': Assumed to return rows.

EXECUTE_UPDATE = 'execute_update': Executes the statement without returning any results.