feat: Added support for reading from Reader Endpoints for AWS Aurora usecases

Bhargav Dodla · Bhargav Dodla · commit c111de71f64a · 2024-09-06T01:21:39.000-07:00
Signed-off-by: Bhargav Dodla &lt;bdodla@expediagroup.com&gt;
diff --git a/sdk/python/feast/infra/registry/sql.py b/sdk/python/feast/infra/registry/sql.py
@@ -202,6 +202,10 @@ class SqlRegistryConfig(RegistryConfig):
     """ str: Path to metadata store.
     If registry_type is 'sql', then this is a database URL as expected by SQLAlchemy """
 
+    read_path: Optional[StrictStr] = None
+    """ str: Read Path to metadata store if different from path.
+    If registry_type is 'sql', then this is a Read Endpoint for database URL. If not set, path will be used for read and write. """
+
     sqlalchemy_config_kwargs: Dict[str, Any] = {"echo": False}
     """ Dict[str, Any]: Extra arguments to pass to SQLAlchemy.create_engine. """
 
@@ -223,13 +227,20 @@ def __init__(
             registry_config, SqlRegistryConfig
         ), "SqlRegistry needs a valid registry_config"
 
-        self.engine: Engine = create_engine(
+        self.write_engine: Engine = create_engine(
             registry_config.path, **registry_config.sqlalchemy_config_kwargs
         )
+        if registry_config.read_path:
+            self.read_engine: Engine = create_engine(
+                registry_config.read_path,
+                **registry_config.sqlalchemy_config_kwargs,
+            )
+        else:
+            self.read_engine = self.write_engine
+        metadata.create_all(self.write_engine)
         self.thread_pool_executor_worker_count = (
             registry_config.thread_pool_executor_worker_count
         )
-        metadata.create_all(self.engine)
         self.purge_feast_metadata = registry_config.purge_feast_metadata
         # Sync feast_metadata to projects table
         # when purge_feast_metadata is set to True, Delete data from
@@ -246,7 +257,7 @@ def __init__(
     def _sync_feast_metadata_to_projects_table(self):
         feast_metadata_projects: set = []
         projects_set: set = []
-        with self.engine.begin() as conn:
+        with self.write_engine.begin() as conn:
             stmt = select(feast_metadata).where(
                 feast_metadata.c.metadata_key == FeastMetadataKeys.PROJECT_UUID.value
             )
@@ -255,7 +266,7 @@ def _sync_feast_metadata_to_projects_table(self):
                 feast_metadata_projects.append(row._mapping["project_id"])
 
         if len(feast_metadata_projects) > 0:
-            with self.engine.begin() as conn:
+            with self.write_engine.begin() as conn:
                 stmt = select(projects)
                 rows = conn.execute(stmt).all()
                 for row in rows:
@@ -267,7 +278,7 @@ def _sync_feast_metadata_to_projects_table(self):
                 self.apply_project(Project(name=project_name), commit=True)
 
             if self.purge_feast_metadata:
-                with self.engine.begin() as conn:
+                with self.write_engine.begin() as conn:
                     for project_name in feast_metadata_projects:
                         stmt = delete(feast_metadata).where(
                             feast_metadata.c.project_id == project_name
@@ -285,7 +296,7 @@ def teardown(self):
             validation_references,
             permissions,
         }:
-            with self.engine.begin() as conn:
+            with self.write_engine.begin() as conn:
                 stmt = delete(t)
                 conn.execute(stmt)
 
@@ -494,7 +505,7 @@ def apply_feature_service(
         )
 
     def delete_data_source(self, name: str, project: str, commit: bool = True):
-        with self.engine.begin() as conn:
+        with self.write_engine.begin() as conn:
             stmt = delete(data_sources).where(
                 data_sources.c.data_source_name == name,
                 data_sources.c.project_id == project,
@@ -552,7 +563,7 @@ def _list_on_demand_feature_views(
         )
 
     def _list_project_metadata(self, project: str) -> List[ProjectMetadata]:
-        with self.engine.begin() as conn:
+        with self.read_engine.begin() as conn:
             stmt = select(feast_metadata).where(
                 feast_metadata.c.project_id == project,
             )
@@ -671,7 +682,7 @@ def apply_user_metadata(
         table = self._infer_fv_table(feature_view)
 
         name = feature_view.name
-        with self.engine.begin() as conn:
+        with self.write_engine.begin() as conn:
             stmt = select(table).where(
                 getattr(table.c, "feature_view_name") == name,
                 table.c.project_id == project,
@@ -726,7 +737,7 @@ def get_user_metadata(
         table = self._infer_fv_table(feature_view)
 
         name = feature_view.name
-        with self.engine.begin() as conn:
+        with self.read_engine.begin() as conn:
             stmt = select(table).where(getattr(table.c, "feature_view_name") == name)
             row = conn.execute(stmt).first()
             if row:
@@ -830,7 +841,7 @@ def _apply_object(
         name = name or (obj.name if hasattr(obj, "name") else None)
         assert name, f"name needs to be provided for {obj}"
 
-        with self.engine.begin() as conn:
+        with self.write_engine.begin() as conn:
             update_datetime = _utc_now()
             update_time = int(update_datetime.timestamp())
             stmt = select(table).where(
@@ -906,7 +917,7 @@ def _apply_object(
 
     def _maybe_init_project_metadata(self, project):
         # Initialize project metadata if needed
-        with self.engine.begin() as conn:
+        with self.write_engine.begin() as conn:
             update_datetime = _utc_now()
             update_time = int(update_datetime.timestamp())
             stmt = select(feast_metadata).where(
@@ -933,7 +944,7 @@ def _delete_object(
         id_field_name: str,
         not_found_exception: Optional[Callable],
     ):
-        with self.engine.begin() as conn:
+        with self.write_engine.begin() as conn:
             stmt = delete(table).where(
                 getattr(table.c, id_field_name) == name, table.c.project_id == project
             )
@@ -959,7 +970,7 @@ def _get_object(
         proto_field_name: str,
         not_found_exception: Optional[Callable],
     ):
-        with self.engine.begin() as conn:
+        with self.read_engine.begin() as conn:
             stmt = select(table).where(
                 getattr(table.c, id_field_name) == name, table.c.project_id == project
             )
@@ -981,7 +992,7 @@ def _list_objects(
         proto_field_name: str,
         tags: Optional[dict[str, str]] = None,
     ):
-        with self.engine.begin() as conn:
+        with self.read_engine.begin() as conn:
             stmt = select(table).where(table.c.project_id == project)
             rows = conn.execute(stmt).all()
             if rows:
@@ -996,7 +1007,7 @@ def _list_objects(
         return []
 
     def _set_last_updated_metadata(self, last_updated: datetime, project: str):
-        with self.engine.begin() as conn:
+        with self.write_engine.begin() as conn:
             stmt = select(feast_metadata).where(
                 feast_metadata.c.metadata_key
                 == FeastMetadataKeys.LAST_UPDATED_TIMESTAMP.value,
@@ -1030,7 +1041,7 @@ def _set_last_updated_metadata(self, last_updated: datetime, project: str):
                 conn.execute(insert_stmt)
 
     def _get_last_updated_metadata(self, project: str):
-        with self.engine.begin() as conn:
+        with self.read_engine.begin() as conn:
             stmt = select(feast_metadata).where(
                 feast_metadata.c.metadata_key
                 == FeastMetadataKeys.LAST_UPDATED_TIMESTAMP.value,
@@ -1075,7 +1086,7 @@ def apply_permission(
         )
 
     def delete_permission(self, name: str, project: str, commit: bool = True):
-        with self.engine.begin() as conn:
+        with self.write_engine.begin() as conn:
             stmt = delete(permissions).where(
                 permissions.c.permission_name == name,
                 permissions.c.project_id == project,
@@ -1088,7 +1099,7 @@ def _list_projects(
         self,
         tags: Optional[dict[str, str]],
     ) -> List[Project]:
-        with self.engine.begin() as conn:
+        with self.read_engine.begin() as conn:
             stmt = select(projects)
             rows = conn.execute(stmt).all()
             if rows:
@@ -1133,7 +1144,7 @@ def delete_project(
     ):
         project = self.get_project(name, allow_cache=False)
         if project:
-            with self.engine.begin() as conn:
+            with self.write_engine.begin() as conn:
                 for t in {
                     managed_infra,
                     saved_datasets,
diff --git a/sdk/python/tests/integration/registration/test_universal_registry.py b/sdk/python/tests/integration/registration/test_universal_registry.py
@@ -123,6 +123,8 @@ def minio_registry() -> Registry:
 
 POSTGRES_USER = "test"
 POSTGRES_PASSWORD = "test"
+POSTGRES_READONLY_USER = "read_only_user"
+POSTGRES_READONLY_PASSWORD = "readonly_password"
 POSTGRES_DB = "test"
 
 logger = logging.getLogger(__name__)
@@ -166,6 +168,38 @@ def pg_registry_async():
     container.stop()
 
 
+def add_pg_read_only_user(container_host, container_port):
+    # Connect to PostgreSQL as an admin
+    import psycopg
+
+    conn_string = f"dbname={POSTGRES_DB} user={POSTGRES_USER} password={POSTGRES_PASSWORD} host={container_host} port={container_port}"
+    print(f"conn_string: {conn_string}")
+
+    with psycopg.connect(conn_string) as conn:
+        # Create the read-only user (if not exists)
+        conn.execute(
+            f"CREATE USER {POSTGRES_READONLY_USER} WITH PASSWORD '{POSTGRES_READONLY_PASSWORD}';"
+        )
+
+        # Revoke any privileges just in case (to ensure it's clean)
+        conn.execute(
+            f"REVOKE ALL PRIVILEGES ON DATABASE {POSTGRES_DB} FROM {POSTGRES_READONLY_USER};"
+        )
+
+        # Grant connect and select privileges on all tables in the database
+        conn.execute(
+            f"GRANT CONNECT ON DATABASE {POSTGRES_DB} TO {POSTGRES_READONLY_USER};"
+        )
+        conn.execute(
+            f"GRANT SELECT ON ALL TABLES IN SCHEMA public TO {POSTGRES_READONLY_USER};"
+        )
+
+        # Optionally, ensure read-only access to any new tables added in the future
+        conn.execute(
+            f"ALTER DEFAULT PRIVILEGES IN SCHEMA public GRANT SELECT ON TABLES TO {POSTGRES_READONLY_USER};"
+        )
+
+
 def _given_registry_config_for_pg_sql(
     container,
     cache_ttl_seconds=2,
@@ -184,13 +218,16 @@ def _given_registry_config_for_pg_sql(
     container_port = container.get_exposed_port(5432)
     container_host = container.get_container_host_ip()
 
+    add_pg_read_only_user(container_host, container_port)
+
     return SqlRegistryConfig(
         registry_type="sql",
         cache_ttl_seconds=cache_ttl_seconds,
         cache_mode=cache_mode,
         # The `path` must include `+psycopg` in order for `sqlalchemy.create_engine()`
         # to understand that we are using psycopg3.
         path=f"postgresql+psycopg://{POSTGRES_USER}:{POSTGRES_PASSWORD}@{container_host}:{container_port}/{POSTGRES_DB}",
+        read_path=f"postgresql+psycopg://{POSTGRES_READONLY_USER}:{POSTGRES_READONLY_PASSWORD}@{container_host}:{container_port}/{POSTGRES_DB}",
         sqlalchemy_config_kwargs={"echo": False, "pool_pre_ping": True},
         thread_pool_executor_worker_count=thread_pool_executor_worker_count,
         purge_feast_metadata=purge_feast_metadata,
