def whitelist_for_serdes(
__cls: Optional[Type] = None,
*,
serializer: Optional[Type["Serializer"]] = None,
storage_name: Optional[str] = None,
):
"""
Decorator to whitelist a NamedTuple or enum to be serializable. If a `storage_name` is provided
for a NamedTuple, then serialized instances of the NamedTuple will be stored with under the
`storage_name` instead of the class name. This is primarily useful for maintaining backwards
compatibility. If a serialized object undergoes a name change, then setting `storage_name` to
the old name will (a) allow the object to be deserialized by versions of Dagster prior to the
name change; (b) allow Dagster to load objects stored using the old name.
@whitelist_for_serdes
class
"""
check.invariant(
not storage_name
or (serializer is None or issubclass(serializer, DefaultNamedTupleSerializer)),
"storage_name can only be used with DefaultNamedTupleSerializer",
)
if __cls is not None: # decorator invoked directly on class
check.class_param(__cls, "__cls")
return _whitelist_for_serdes(whitelist_map=_WHITELIST_MAP)(__cls)
else: # decorator passed params
check.opt_class_param(serializer, "serializer", superclass=Serializer)
serializer = cast(Type[Serializer], serializer)
return _whitelist_for_serdes(
whitelist_map=_WHITELIST_MAP, serializer=serializer, storage_name=storage_name
)
def execute_queries(self,
queries,
fetch_results=False,
cursor_factory=None,
error_callback=None):
"""Fake for execute_queries; returns [self.QUERY_RESULT] * 3
Args:
queries (List[str]): The queries to execute.
fetch_results (Optional[bool]): Whether to return the results of executing the query.
Defaults to False, in which case the query will be executed without retrieving the
results.
cursor_factory (Optional[:py:class:`psycopg2.extensions.cursor`]): An alternative
cursor_factory; defaults to None. Will be used when constructing the cursor.
error_callback (Optional[Callable[[Exception, Cursor, DagsterLogManager], None]]): A
callback function, invoked when an exception is encountered during query execution;
this is intended to support executing additional queries to provide diagnostic
information, e.g. by querying ``stl_load_errors`` using ``pg_last_copy_id()``. If no
function is provided, exceptions during query execution will be raised directly.
Returns:
Optional[List[List[Tuple[Any, ...]]]]: Results of the query, as a list of list of
tuples, when fetch_results is set. Otherwise return None.
"""
check.list_param(queries, "queries", of_type=str)
check.bool_param(fetch_results, "fetch_results")
check.opt_class_param(cursor_factory,
"cursor_factory",
superclass=psycopg2.extensions.cursor)
check.opt_callable_param(error_callback, "error_callback")
for query in queries:
self.log.info("Executing query '{query}'".format(query=query))
if fetch_results:
return [self.QUERY_RESULT] * 3
def test_class_param():
class Bar:
pass
assert check.class_param(int, "foo")
assert check.class_param(Bar, "foo")
with pytest.raises(CheckError):
check.class_param(None, "foo")
with pytest.raises(CheckError):
check.class_param(check, "foo")
with pytest.raises(CheckError):
check.class_param(234, "foo")
with pytest.raises(CheckError):
check.class_param("bar", "foo")
with pytest.raises(CheckError):
check.class_param(Bar(), "foo")
class Super:
pass
class Sub(Super):
pass
class Alone:
pass
assert check.class_param(Sub, "foo", superclass=Super)
with pytest.raises(CheckError):
assert check.class_param(Alone, "foo", superclass=Super)
with pytest.raises(CheckError):
assert check.class_param("value", "foo", superclass=Super)
assert check.opt_class_param(Sub, "foo", superclass=Super)
assert check.opt_class_param(None, "foo", superclass=Super) is None
with pytest.raises(CheckError):
assert check.opt_class_param(Alone, "foo", superclass=Super)
with pytest.raises(CheckError):
assert check.opt_class_param("value", "foo", superclass=Super)
def _get_cursor(self, conn, cursor_factory=None):
check.opt_class_param(cursor_factory,
"cursor_factory",
superclass=psycopg2.extensions.cursor)
# Could be none, in which case we should respect the connection default. Otherwise
# explicitly set to true/false.
if self.autocommit is not None:
conn.autocommit = self.autocommit
with conn:
with conn.cursor(cursor_factory=cursor_factory) as cursor:
yield cursor
# If autocommit is set, we'll commit after each and every query execution. Otherwise, we
# want to do a final commit after we're wrapped up executing the full set of one or more
# queries.
if not self.autocommit:
conn.commit()
def test_opt_class_param():
class Foo:
pass
assert check.opt_class_param(int, "foo")
assert check.opt_class_param(Foo, "foo")
assert check.opt_class_param(None, "foo") is None
assert check.opt_class_param(None, "foo", Foo) is Foo
with pytest.raises(CheckError):
check.opt_class_param(check, "foo")
with pytest.raises(CheckError):
check.opt_class_param(234, "foo")
with pytest.raises(CheckError):
check.opt_class_param("bar", "foo")
with pytest.raises(CheckError):
check.opt_class_param(Foo(), "foo")
def execute_query(self,
query,
fetch_results=False,
cursor_factory=None,
error_callback=None):
"""Synchronously execute a single query against Redshift. Will return a list of rows, where
each row is a tuple of values, e.g. SELECT 1 will return [(1,)].
Args:
query (str): The query to execute.
fetch_results (Optional[bool]): Whether to return the results of executing the query.
Defaults to False, in which case the query will be executed without retrieving the
results.
cursor_factory (Optional[:py:class:`psycopg2.extensions.cursor`]): An alternative
cursor_factory; defaults to None. Will be used when constructing the cursor.
error_callback (Optional[Callable[[Exception, Cursor, DagsterLogManager], None]]): A
callback function, invoked when an exception is encountered during query execution;
this is intended to support executing additional queries to provide diagnostic
information, e.g. by querying ``stl_load_errors`` using ``pg_last_copy_id()``. If no
function is provided, exceptions during query execution will be raised directly.
Returns:
Optional[List[Tuple[Any, ...]]]: Results of the query, as a list of tuples, when
fetch_results is set. Otherwise return None.
"""
check.str_param(query, "query")
check.bool_param(fetch_results, "fetch_results")
check.opt_class_param(cursor_factory,
"cursor_factory",
superclass=psycopg2.extensions.cursor)
check.opt_callable_param(error_callback, "error_callback")
with self._get_conn() as conn:
with self._get_cursor(conn,
cursor_factory=cursor_factory) as cursor:
try:
self.log.info(
"Executing query '{query}'".format(query=query))
cursor.execute(query)
if fetch_results and cursor.rowcount > 0:
return cursor.fetchall()
else:
self.log.info("Empty result from query")
except Exception as e:
# If autocommit is disabled or not set (it is disabled by default), Redshift
# will be in the middle of a transaction at exception time, and because of
# the failure the current transaction will not accept any further queries.
#
# This conn.commit() call closes the open transaction before handing off
# control to the error callback, so that the user can issue additional
# queries. Notably, for e.g. pg_last_copy_id() to work, it requires you to
# use the same conn/cursor, so you have to do this conn.commit() to ensure
# things are in a usable state in the error callback.
if not self.autocommit:
conn.commit()
if error_callback is not None:
error_callback(e, cursor, self.log)
else:
raise