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