feast-dev · blvp · Aug 10, 2021 · Aug 10, 2021 · Aug 10, 2021 · Aug 10, 2021
@@ -118,6 +118,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 table schema name
+    string schema = 3;
   }
 
   // Defines configuration for custom third-party data sources.

@@ -10,6 +10,14 @@ def __init__(self, path):
         )
 
 
+class UndefinedDatasourceSchemaException(Exception):
+    def __init__(self, schemas_found, path):
+        super().__init__(
+            f"Unable to identify correct schema location from table at '{path}'. \
+                Found schemas: [{', '.join(schemas_found)}]"
+        )
+
+
 class FeastObjectNotFoundException(Exception):
     pass
 

@@ -38,6 +38,9 @@ class RedshiftOfflineStoreConfig(FeastConfigBaseModel):
     database: StrictStr
     """ Redshift database name """
 
+    temp_schema_name: StrictStr
+    """ Redshift schema name to offload temporary tables """
+
     s3_staging_location: StrictStr
     """ S3 path for importing & exporting data to Redshift """
 
@@ -254,13 +257,15 @@ def to_s3(self) -> str:
             )
             return self._s3_path
 
-    def to_redshift(self, table_name: str) -> None:
+    def to_redshift(self, table_name: str, schema: Optional[str] = None) -> None:
         """ Save dataset as a new Redshift table """
         with self._query_generator() as query:
-            query = f'CREATE TABLE "{table_name}" AS ({query});\n'
+            schema_prefix = f"{schema}." if schema is not None else ""
+            full_table_name = f"{schema_prefix}{table_name}"
+            query = f'CREATE TABLE "{full_table_name}" AS ({query});\n'
             if self._drop_columns is not None:
                 for column in self._drop_columns:
-                    query += f"ALTER TABLE {table_name} DROP COLUMN {column};\n"
+                    query += f"ALTER TABLE {full_table_name} DROP COLUMN {column};\n"
 
             aws_utils.execute_redshift_statement(
                 self._redshift_client,
@@ -291,20 +296,25 @@ def _upload_entity_df_and_get_entity_schema(
             config.offline_store.iam_role,
             table_name,
             entity_df,
+            config.offline_store.temp_schema_name,
         )
         return dict(zip(entity_df.columns, entity_df.dtypes))
     elif isinstance(entity_df, str):
         # If the entity_df is a string (SQL query), create a Redshift table out of it,
         # get pandas dataframe consisting of 1 row (LIMIT 1) and generate the schema out of it
+        full_table_name = f"{config.offline_store.temp_schema_name}.{table_name}"
         aws_utils.execute_redshift_statement(
             redshift_client,
             config.offline_store.cluster_id,
             config.offline_store.database,
             config.offline_store.user,
-            f"CREATE TABLE {table_name} AS ({entity_df})",
+            f"CREATE TABLE {full_table_name} AS ({entity_df})",
         )
         limited_entity_df = RedshiftRetrievalJob(
-            f"SELECT * FROM {table_name} LIMIT 1", redshift_client, s3_resource, config
+            f"SELECT * FROM {full_table_name} LIMIT 1",
+            redshift_client,
+            s3_resource,
+            config,
         ).to_df()
         return dict(zip(limited_entity_df.columns, limited_entity_df.dtypes))
     else:

@@ -2,7 +2,11 @@
 
 from feast import type_map
 from feast.data_source import DataSource
-from feast.errors import DataSourceNotFoundException, RedshiftCredentialsError
+from feast.errors import (
+    DataSourceNotFoundException,
+    RedshiftCredentialsError,
+    UndefinedDatasourceSchemaException,
+)
 from feast.protos.feast.core.DataSource_pb2 import DataSource as DataSourceProto
 from feast.repo_config import RepoConfig
 from feast.value_type import ValueType
@@ -13,6 +17,7 @@ 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] = "",
@@ -25,7 +30,9 @@ def __init__(
             date_partition_column,
         )
 
-        self._redshift_options = RedshiftOptions(table=table, query=query)
+        self._redshift_options = RedshiftOptions(
+            table=table, schema=schema, query=query
+        )
 
     @staticmethod
     def from_proto(data_source: DataSourceProto):
@@ -47,6 +54,7 @@ def __eq__(self, other):
         return (
             self.redshift_options.table == other.redshift_options.table
             and self.redshift_options.query == other.redshift_options.query
+            and self.redshift_options.schema == other.redshift_options.schema
             and self.event_timestamp_column == other.event_timestamp_column
             and self.created_timestamp_column == other.created_timestamp_column
             and self.field_mapping == other.field_mapping
@@ -60,6 +68,10 @@ def table(self):
     def query(self):
         return self._redshift_options.query
 
+    @property
+    def schema(self):
+        return self._redshift_options.schema
+
     @property
     def redshift_options(self):
         """
@@ -95,7 +107,8 @@ def validate(self, config: RepoConfig):
     def get_table_query_string(self) -> str:
         """Returns a string that can directly be used to reference this table in SQL"""
         if self.table:
-            return f'"{self.table}"'
+            schema_prefix = f"{self.schema}." if self.schema is not None else ""
+            return f'"{schema_prefix}{self.table}"'
         else:
             return f"({self.query})"
 
@@ -117,12 +130,16 @@ def get_table_column_names_and_types(
 
         if self.table is not None:
             try:
-                table = client.describe_table(
-                    ClusterIdentifier=config.offline_store.cluster_id,
-                    Database=config.offline_store.database,
-                    DbUser=config.offline_store.user,
-                    Table=self.table,
-                )
+                desribe_table_req = {
+                    "ClusterIdentifier": config.offline_store.cluster_id,
+                    "Database": config.offline_store.database,
+                    "DbUser": config.offline_store.user,
+                    "Table": self.table,
+                }
+                if self.schema is not None:
+                    desribe_table_req["Schema"] = self.schema
+
+                table = client.describe_table(**desribe_table_req)
             except ClientError as e:
                 if e.response["Error"]["Code"] == "ValidationException":
                     raise RedshiftCredentialsError() from e
@@ -132,6 +149,10 @@ def get_table_column_names_and_types(
             if len(table["ColumnList"]) == 0:
                 raise DataSourceNotFoundException(self.table)
 
+            unique_schemas = {col["schemaName"] for col in table["ColumnList"]}
+            if len(unique_schemas) > 1:
+                raise UndefinedDatasourceSchemaException(unique_schemas, self.table)
+
             columns = table["ColumnList"]
         else:
             statement_id = aws_utils.execute_redshift_statement(
@@ -153,9 +174,21 @@ class RedshiftOptions:
     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], query: Optional[str], schema: Optional[str]
+    ):
+        """Redshift options to encapsulate logic for parsing and working with 2 kinds of source creation
+        table + schema or query
+
+        Args:
+            table (Optional[str]): Redshift table to be looked for in redshift cluster to form datasource
+            query (Optional[str]): Query to run to gather datasource
+            schema (Optional[str]): Schema in redshift cluster to lookup a table.
+                Has to be provided in case of tables with same name.
+        """
         self._table = table
         self._query = query
+        self._schema = schema
 
     @property
     def query(self):
@@ -185,6 +218,20 @@ def table(self, table_name):
         """
         self._table = table_name
 
+    @property
+    def schema(self):
+        """
+        Returns the schema name of this Redshift table schema
+        """
+        return self._schema
+
+    @schema.setter
+    def schema(self, schema_name):
+        """
+        Sets the schema ref of this Redshift table schema
+        """
+        self._schema = schema_name
+
     @classmethod
     def from_proto(cls, redshift_options_proto: DataSourceProto.RedshiftOptions):
         """
@@ -198,7 +245,9 @@ def from_proto(cls, redshift_options_proto: DataSourceProto.RedshiftOptions):
         """
 
         redshift_options = cls(
-            table=redshift_options_proto.table, query=redshift_options_proto.query,
+            table=redshift_options_proto.table,
+            query=redshift_options_proto.query,
+            schema=redshift_options_proto.schema,
         )
 
         return redshift_options
@@ -212,7 +261,7 @@ def to_proto(self) -> DataSourceProto.RedshiftOptions:
         """
 
         redshift_options_proto = DataSourceProto.RedshiftOptions(
-            table=self.table, query=self.query,
+            table=self.table, query=self.query, schema=self.schema
         )
 
         return redshift_options_proto
@@ -146,6 +146,7 @@ def upload_df_to_redshift(
     iam_role: str,
     table_name: str,
     df: pd.DataFrame,
+    schema_name: Optional[str] = None,
 ) -> None:
     """Uploads a Pandas DataFrame to Redshift as a new table.
 
@@ -204,9 +205,11 @@ def upload_df_to_redshift(
 
     # Create the table with the desired schema and
     # copy the Parquet file contents to the Redshift table
+    schema_prefix = f"{schema_name}." if schema_name is not None else ""
+    full_table_name = f"{schema_prefix}{table_name}"
     create_and_copy_query = (
-        f"CREATE TABLE {table_name}({column_query_list}); "
-        + f"COPY {table_name} FROM '{s3_path}' IAM_ROLE '{iam_role}' FORMAT AS PARQUET"
+        f"CREATE TABLE {full_table_name}({column_query_list}); "
+        + f"COPY {full_table_name} FROM '{s3_path}' IAM_ROLE '{iam_role}' FORMAT AS PARQUET"
     )
     execute_redshift_statement(
         redshift_data_client, cluster_id, database, user, create_and_copy_query
@@ -227,6 +230,7 @@ def temporarily_upload_df_to_redshift(
     iam_role: str,
     table_name: str,
     df: pd.DataFrame,
+    schema_name: Optional[str] = None,
 ) -> Iterator[None]:
     """Uploads a Pandas DataFrame to Redshift as a new table with cleanup logic.
 
@@ -249,6 +253,7 @@ def temporarily_upload_df_to_redshift(
         iam_role,
         table_name,
         df,
+        schema_name,
     )
 
     yield