| 970 | |
| 971 | |
| 972 | class RedshiftRetrievalJob(RetrievalJob): |
| 973 | def __init__( |
| 974 | self, |
| 975 | query: Union[str, Callable[[], ContextManager[str]]], |
| 976 | redshift_client, |
| 977 | s3_resource, |
| 978 | config: RepoConfig, |
| 979 | full_feature_names: bool, |
| 980 | on_demand_feature_views: Optional[List[OnDemandFeatureView]] = None, |
| 981 | metadata: Optional[RetrievalMetadata] = None, |
| 982 | ): |
| 983 | """Initialize RedshiftRetrievalJob object. |
| 984 | |
| 985 | Args: |
| 986 | query: Redshift SQL query to execute. Either a string, or a generator function that handles the artifact cleanup. |
| 987 | redshift_client: boto3 redshift-data client |
| 988 | s3_resource: boto3 s3 resource object |
| 989 | config: Feast repo config |
| 990 | full_feature_names: Whether to add the feature view prefixes to the feature names |
| 991 | on_demand_feature_views (optional): A list of on demand transforms to apply at retrieval time |
| 992 | """ |
| 993 | if not isinstance(query, str): |
| 994 | self._query_generator = query |
| 995 | else: |
| 996 | |
| 997 | @contextlib.contextmanager |
| 998 | def query_generator() -> Iterator[str]: |
| 999 | assert isinstance(query, str) |
| 1000 | yield query |
| 1001 | |
| 1002 | self._query_generator = query_generator |
| 1003 | self._redshift_client = redshift_client |
| 1004 | self._s3_resource = s3_resource |
| 1005 | self._config = config |
| 1006 | self._s3_path = ( |
| 1007 | self._config.offline_store.s3_staging_location |
| 1008 | + "/unload/" |
| 1009 | + str(uuid.uuid4()) |
| 1010 | ) |
| 1011 | self._full_feature_names = full_feature_names |
| 1012 | self._on_demand_feature_views = on_demand_feature_views or [] |
| 1013 | self._metadata = metadata |
| 1014 | |
| 1015 | @property |
| 1016 | def full_feature_names(self) -> bool: |
| 1017 | return self._full_feature_names |
| 1018 | |
| 1019 | @property |
| 1020 | def on_demand_feature_views(self) -> List[OnDemandFeatureView]: |
| 1021 | return self._on_demand_feature_views |
| 1022 | |
| 1023 | def _to_df_internal(self, timeout: Optional[int] = None) -> pd.DataFrame: |
| 1024 | with self._query_generator() as query: |
| 1025 | return aws_utils.unload_redshift_query_to_df( |
| 1026 | self._redshift_client, |
| 1027 | self._config.offline_store.cluster_id, |
| 1028 | self._config.offline_store.workgroup, |
| 1029 | self._config.offline_store.database, |
no outgoing calls