From 809ce38157537ef6477fc91c2d3f012ed231f03a Mon Sep 17 00:00:00 2001 From: Felix Wang Date: Wed, 8 Sep 2021 10:23:30 -0700 Subject: [PATCH] Add schema parameter to RedshiftSource Signed-off-by: Felix Wang --- protos/feast/core/DataSource.proto | 3 + .../infra/offline_stores/redshift_source.py | 113 +++++++++++++----- 2 files changed, 86 insertions(+), 30 deletions(-) diff --git a/protos/feast/core/DataSource.proto b/protos/feast/core/DataSource.proto index 2e43ac9a233..7bdd2ec298d 100644 --- a/protos/feast/core/DataSource.proto +++ b/protos/feast/core/DataSource.proto @@ -121,6 +121,9 @@ message DataSource { // SQL query that returns a table containing feature data. Must contain an event_timestamp column, and respective // entity columns string query = 2; + + // Redshift schema name + string schema = 3; } // Defines configuration for custom third-party data sources. diff --git a/sdk/python/feast/infra/offline_stores/redshift_source.py b/sdk/python/feast/infra/offline_stores/redshift_source.py index 6e3a8e0f293..e7e88a54ef7 100644 --- a/sdk/python/feast/infra/offline_stores/redshift_source.py +++ b/sdk/python/feast/infra/offline_stores/redshift_source.py @@ -13,11 +13,27 @@ def __init__( self, event_timestamp_column: Optional[str] = "", table: Optional[str] = None, + schema: Optional[str] = None, created_timestamp_column: Optional[str] = "", field_mapping: Optional[Dict[str, str]] = None, date_partition_column: Optional[str] = "", query: Optional[str] = None, ): + """ + Creates a RedshiftSource object. + + Args: + event_timestamp_column (optional): Event timestamp column used for point in + time joins of feature values. + table (optional): Redshift table where the features are stored. + schema (optional): Redshift schema in which the table is located. + created_timestamp_column (optional): Timestamp column indicating when the + row was created, used for deduplicating rows. + field_mapping (optional): A dictionary mapping of column names in this data + source to column names in a feature table or view. + date_partition_column (optional): Timestamp column used for partitioning. + query (optional): The query to be executed to obtain the features. + """ super().__init__( event_timestamp_column, created_timestamp_column, @@ -25,13 +41,28 @@ def __init__( date_partition_column, ) - self._redshift_options = RedshiftOptions(table=table, query=query) + # The default Redshift schema is named "public". + _schema = "public" if table and not schema else schema + + self._redshift_options = RedshiftOptions( + table=table, schema=_schema, query=query + ) @staticmethod def from_proto(data_source: DataSourceProto): + """ + Creates a RedshiftSource from a protobuf representation of a RedshiftSource. + + Args: + data_source: A protobuf representation of a RedshiftSource + + Returns: + A RedshiftSource object based on the data_source protobuf. + """ return RedshiftSource( field_mapping=dict(data_source.field_mapping), table=data_source.redshift_options.table, + schema=data_source.redshift_options.schema, event_timestamp_column=data_source.event_timestamp_column, created_timestamp_column=data_source.created_timestamp_column, date_partition_column=data_source.date_partition_column, @@ -46,6 +77,7 @@ def __eq__(self, other): return ( self.redshift_options.table == other.redshift_options.table + and self.redshift_options.schema == other.redshift_options.schema and self.redshift_options.query == other.redshift_options.query and self.event_timestamp_column == other.event_timestamp_column and self.created_timestamp_column == other.created_timestamp_column @@ -54,27 +86,36 @@ def __eq__(self, other): @property def table(self): + """Returns the table of this Redshift source.""" return self._redshift_options.table + @property + def schema(self): + """Returns the schema of this Redshift source.""" + return self._redshift_options.schema + @property def query(self): + """Returns the Redshift options of this Redshift source.""" return self._redshift_options.query @property def redshift_options(self): - """ - Returns the Redshift options of this data source - """ + """Returns the Redshift options of this Redshift source.""" return self._redshift_options @redshift_options.setter def redshift_options(self, _redshift_options): - """ - Sets the Redshift options of this data source - """ + """Sets the Redshift options of this Redshift source.""" self._redshift_options = _redshift_options def to_proto(self) -> DataSourceProto: + """ + Converts a RedshiftSource object to its protobuf representation. + + Returns: + A DataSourceProto object. + """ data_source_proto = DataSourceProto( type=DataSourceProto.BATCH_REDSHIFT, field_mapping=self.field_mapping, @@ -93,9 +134,9 @@ def validate(self, config: RepoConfig): self.get_table_column_names_and_types(config) def get_table_query_string(self) -> str: - """Returns a string that can directly be used to reference this table in SQL""" + """Returns a string that can directly be used to reference this table in SQL.""" if self.table: - return f'"{self.table}"' + return f'"{self.schema}"."{self.table}"' else: return f"({self.query})" @@ -106,6 +147,12 @@ def source_datatype_to_feast_value_type() -> Callable[[str], ValueType]: def get_table_column_names_and_types( self, config: RepoConfig ) -> Iterable[Tuple[str, str]]: + """ + Returns a mapping of column names to types for this Redshift source. + + Args: + config: A RepoConfig describing the feature repo + """ from botocore.exceptions import ClientError from feast.infra.offline_stores.redshift import RedshiftOfflineStoreConfig @@ -122,6 +169,7 @@ def get_table_column_names_and_types( Database=config.offline_store.database, DbUser=config.offline_store.user, Table=self.table, + Schema=self.schema, ) except ClientError as e: if e.response["Error"]["Code"] == "ValidationException": @@ -150,55 +198,61 @@ def get_table_column_names_and_types( class RedshiftOptions: """ - DataSource Redshift options used to source features from Redshift query + DataSource Redshift options used to source features from Redshift query. """ - def __init__(self, table: Optional[str], query: Optional[str]): + def __init__( + self, table: Optional[str], schema: Optional[str], query: Optional[str] + ): self._table = table + self._schema = schema self._query = query @property def query(self): - """ - Returns the Redshift SQL query referenced by this source - """ + """Returns the Redshift SQL query referenced by this source.""" return self._query @query.setter def query(self, query): - """ - Sets the Redshift SQL query referenced by this source - """ + """Sets the Redshift SQL query referenced by this source.""" self._query = query @property def table(self): - """ - Returns the table name of this Redshift table - """ + """Returns the table name of this Redshift table.""" return self._table @table.setter def table(self, table_name): - """ - Sets the table ref of this Redshift table - """ + """Sets the table ref of this Redshift table.""" self._table = table_name + @property + def schema(self): + """Returns the schema name of this Redshift table.""" + return self._schema + + @schema.setter + def schema(self, schema): + """Sets the schema of this Redshift table.""" + self._schema = schema + @classmethod def from_proto(cls, redshift_options_proto: DataSourceProto.RedshiftOptions): """ - Creates a RedshiftOptions from a protobuf representation of a Redshift option + Creates a RedshiftOptions from a protobuf representation of a Redshift option. Args: redshift_options_proto: A protobuf representation of a DataSource Returns: - Returns a RedshiftOptions object based on the redshift_options protobuf + A RedshiftOptions object based on the redshift_options protobuf. """ - redshift_options = cls( - table=redshift_options_proto.table, query=redshift_options_proto.query, + table=redshift_options_proto.table, + schema=redshift_options_proto.schema, + query=redshift_options_proto.query, ) return redshift_options @@ -208,11 +262,10 @@ def to_proto(self) -> DataSourceProto.RedshiftOptions: Converts an RedshiftOptionsProto object to its protobuf representation. Returns: - RedshiftOptionsProto protobuf + A RedshiftOptionsProto protobuf. """ - redshift_options_proto = DataSourceProto.RedshiftOptions( - table=self.table, query=self.query, + table=self.table, schema=self.schema, query=self.query, ) return redshift_options_proto