We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN ISSUE ---
Table in Postgres OnlineStore is not populated after calling `materialize`
## Expected Behavior
When calling the `materialize` functionality to materialize data from a `SnowflakeSource` offline store to a local `PostgreSQLOnlineStore`, the table is not populated with the data.

## Current Behavior
The feature table in the local Postgres instance is not populated, while no exception is raised, and from the logs it seems like the data should be pushed to Postgres.

## Steps to reproduce

1) Use this feature_store.yaml file:

```
project: my_project
provider: local
registry:
    registry_type: sql
    path: postgresql://postgres:test@0.0.0.0:5432/feature_store
    cache_ttl_seconds: 60
online_store:
   type: postgres
   host: 0.0.0.0
   port: 5432
   database: feature_store
   db_schema: public
   user: postgres
   password: test
offline_store:
    <SNOWFLAKE_INFORMATION>
entity_key_serialization_version: 2
```

2) Spin up this docker-compose file:

```
---
version: "3"
services:
  db:
    restart: always
    image: postgres:15-alpine
    container_name: feast_db
    ports:
      - "5432:5432"
    volumes:
      - ~/feast_postgres_data:/var/lib/postgresql/data
    environment:
      - POSTGRES_DB=feature_store
      - POSTGRES_USER=postgres
      - POSTGRES_PASSWORD=test
volumes:
  feast_postgres_data: null
```

3) Initialize the Entities, SnowflakeSource (or another source), FeatureView, and FeatureService, and apply these. All using the Python SDK.

```
from datetime import timedelta

from feast import (
    Entity,
    FeatureService,
    FeatureView,
    Field,
    SnowflakeSource,
    ValueType,
    FeatureStore,
)
from feast.types import Float32

feature_store = FeatureStore()

entity = Entity(
    name="entity",
    join_keys=["entity_ID"],
    value_type=ValueType.STRING,
)

source = SnowflakeSource(
    name="snowflake_source_name",
    timestamp_field="EVENT_TIMESTAMP",
    schema="TEMP",
    table="TABLE"
)

feature_view = FeatureView(
    name="feature_view_name",
    entities=[entity],
    ttl=timedelta(days=0),
    schema=[
        Field(name="feature_1", dtype=Float32),
        Field(name="feature_2", dtype=Float32),
    ],
    online=True,
    source=source,
    tags={"team": "team"},
)

feature_service = FeatureService(
    name="feature_service",
    features=[feature_view],
)

feature_store.apply(
    [
        entity,
        source,
        feature_view,
        feature_service,
    ]
)
```

4) Run materialize commands using the Python SDK

```
feature_store = FeatureStore()
feature_store.materialize(
    start_date=datetime.utcnow() - timedelta(weeks=52),
    end_date=datetime.utcnow(),
    feature_views=["feature_view_name"],
)
```

### Specifications
- Version: 0.35.0
- Platform: Local MacBook M1

## Possible Solution
It seems like a `conn.commit()` statement is missing in the `online_write_batch` method of the `PostgreSQLOnlineStore`. Specifically, on [this line](https://github.com/feast-dev/feast/blob/master/sdk/python/feast/infra/online_stores/contrib/postgres.py#L102).
After adding this, the table is populated.

The PR implementing this proposed fix can be found [here](https://github.com/feast-dev/feast/pull/3904).

## Additional notes
When replacing the the postgres online store with the following sqlite online store in the config file, everything works without any code changes

```
online_store:
    type: sqlite
    path: data/online_store.db
```

--- END ISSUE ---

Below are some code segments, each from a relevant file. One or more of these files may contain bugs.
--- BEGIN FILES ---
Path: `sdk/python/feast/infra/online_stores/contrib/postgres.py`
Content:
```
1 import contextlib
2 import logging
3 from collections import defaultdict
4 from datetime import datetime
5 from typing import Any, Callable, Dict, List, Optional, Sequence, Tuple
6 
7 import psycopg2
8 import pytz
9 from psycopg2 import sql
10 from psycopg2.extras import execute_values
11 from psycopg2.pool import SimpleConnectionPool
12 from pydantic.schema import Literal
13 
14 from feast import Entity
15 from feast.feature_view import FeatureView
16 from feast.infra.key_encoding_utils import serialize_entity_key
17 from feast.infra.online_stores.online_store import OnlineStore
18 from feast.infra.utils.postgres.connection_utils import _get_conn, _get_connection_pool
19 from feast.infra.utils.postgres.postgres_config import ConnectionType, PostgreSQLConfig
20 from feast.protos.feast.types.EntityKey_pb2 import EntityKey as EntityKeyProto
21 from feast.protos.feast.types.Value_pb2 import Value as ValueProto
22 from feast.repo_config import RepoConfig
23 from feast.usage import log_exceptions_and_usage
24 
25 
26 class PostgreSQLOnlineStoreConfig(PostgreSQLConfig):
27     type: Literal["postgres"] = "postgres"
28 
29 
30 class PostgreSQLOnlineStore(OnlineStore):
31     _conn: Optional[psycopg2._psycopg.connection] = None
32     _conn_pool: Optional[SimpleConnectionPool] = None
33 
34     @contextlib.contextmanager
35     def _get_conn(self, config: RepoConfig):
36         assert config.online_store.type == "postgres"
37         if config.online_store.conn_type == ConnectionType.pool:
38             if not self._conn_pool:
39                 self._conn_pool = _get_connection_pool(config.online_store)
40             connection = self._conn_pool.getconn()
41             yield connection
42             self._conn_pool.putconn(connection)
43         else:
44             if not self._conn:
45                 self._conn = _get_conn(config.online_store)
46             yield self._conn
47 
48     @log_exceptions_and_usage(online_store="postgres")
49     def online_write_batch(
50         self,
51         config: RepoConfig,
52         table: FeatureView,
53         data: List[
54             Tuple[EntityKeyProto, Dict[str, ValueProto], datetime, Optional[datetime]]
55         ],
56         progress: Optional[Callable[[int], Any]],
57     ) -> None:
58         project = config.project
59 
60         with self._get_conn(config) as conn, conn.cursor() as cur:
61             insert_values = []
62             for entity_key, values, timestamp, created_ts in data:
63                 entity_key_bin = serialize_entity_key(
64                     entity_key,
65                     entity_key_serialization_version=config.entity_key_serialization_version,
66                 )
67                 timestamp = _to_naive_utc(timestamp)
68                 if created_ts is not None:
69                     created_ts = _to_naive_utc(created_ts)
70 
71                 for feature_name, val in values.items():
72                     insert_values.append(
73                         (
74                             entity_key_bin,
75                             feature_name,
76                             val.SerializeToString(),
77                             timestamp,
78                             created_ts,
79                         )
80                     )
81             # Control the batch so that we can update the progress
82             batch_size = 5000
83             for i in range(0, len(insert_values), batch_size):
84                 cur_batch = insert_values[i : i + batch_size]
85                 execute_values(
86                     cur,
87                     sql.SQL(
88                         """
89                         INSERT INTO {}
90                         (entity_key, feature_name, value, event_ts, created_ts)
91                         VALUES %s
92                         ON CONFLICT (entity_key, feature_name) DO
93                         UPDATE SET
94                             value = EXCLUDED.value,
95                             event_ts = EXCLUDED.event_ts,
96                             created_ts = EXCLUDED.created_ts;
97                         """,
98                     ).format(sql.Identifier(_table_id(project, table))),
99                     cur_batch,
100                     page_size=batch_size,
101                 )
102                 if progress:
103                     progress(len(cur_batch))
104 
105     @log_exceptions_and_usage(online_store="postgres")
106     def online_read(
107         self,
108         config: RepoConfig,
109         table: FeatureView,
110         entity_keys: List[EntityKeyProto],
111         requested_features: Optional[List[str]] = None,
112     ) -> List[Tuple[Optional[datetime], Optional[Dict[str, ValueProto]]]]:
113         result: List[Tuple[Optional[datetime], Optional[Dict[str, ValueProto]]]] = []
114 
115         project = config.project
116         with self._get_conn(config) as conn, conn.cursor() as cur:
117             # Collecting all the keys to a list allows us to make fewer round trips
118             # to PostgreSQL
119             keys = []
120             for entity_key in entity_keys:
121                 keys.append(
122                     serialize_entity_key(
123                         entity_key,
124                         entity_key_serialization_version=config.entity_key_serialization_version,
125                     )
126                 )
127 
128             if not requested_features:
129                 cur.execute(
130                     sql.SQL(
131                         """
132                         SELECT entity_key, feature_name, value, event_ts
133                         FROM {} WHERE entity_key = ANY(%s);
134                         """
135                     ).format(
136                         sql.Identifier(_table_id(project, table)),
137                     ),
138                     (keys,),
139                 )
140             else:
141                 cur.execute(
142                     sql.SQL(
143                         """
144                         SELECT entity_key, feature_name, value, event_ts
145                         FROM {} WHERE entity_key = ANY(%s) and feature_name = ANY(%s);
146                         """
147                     ).format(
148                         sql.Identifier(_table_id(project, table)),
149                     ),
150                     (keys, requested_features),
151                 )
152 
153             rows = cur.fetchall()
154 
155             # Since we don't know the order returned from PostgreSQL we'll need
156             # to construct a dict to be able to quickly look up the correct row
157             # when we iterate through the keys since they are in the correct order
158             values_dict = defaultdict(list)
159             for row in rows if rows is not None else []:
160                 values_dict[row[0].tobytes()].append(row[1:])
161 
162             for key in keys:
163                 if key in values_dict:
164                     value = values_dict[key]
165                     res = {}
166                     for feature_name, value_bin, event_ts in value:
167                         val = ValueProto()
168                         val.ParseFromString(bytes(value_bin))
169                         res[feature_name] = val
170                     result.append((event_ts, res))
171                 else:
172                     result.append((None, None))
173 
174         return result
175 
176     @log_exceptions_and_usage(online_store="postgres")
177     def update(
178         self,
179         config: RepoConfig,
180         tables_to_delete: Sequence[FeatureView],
181         tables_to_keep: Sequence[FeatureView],
182         entities_to_delete: Sequence[Entity],
183         entities_to_keep: Sequence[Entity],
184         partial: bool,
185     ):
186         project = config.project
187         schema_name = config.online_store.db_schema or config.online_store.user
188         with self._get_conn(config) as conn, conn.cursor() as cur:
189             # If a db_schema is provided, then that schema gets created if it doesn't
190             # exist. Else a schema is created for the feature store user.
191 
192             cur.execute(
193                 """
194                 SELECT schema_name
195                 FROM information_schema.schemata
196                 WHERE schema_name = %s
197                 """,
198                 (schema_name,),
199             )
200             schema_exists = cur.fetchone()
201             if not schema_exists:
202                 cur.execute(
203                     sql.SQL("CREATE SCHEMA IF NOT EXISTS {} AUTHORIZATION {}").format(
204                         sql.Identifier(schema_name),
205                         sql.Identifier(config.online_store.user),
206                     ),
207                 )
208 
209             for table in tables_to_delete:
210                 table_name = _table_id(project, table)
211                 cur.execute(_drop_table_and_index(table_name))
212 
213             for table in tables_to_keep:
214                 table_name = _table_id(project, table)
215                 cur.execute(
216                     sql.SQL(
217                         """
218                         CREATE TABLE IF NOT EXISTS {}
219                         (
220                             entity_key BYTEA,
221                             feature_name TEXT,
222                             value BYTEA,
223                             event_ts TIMESTAMPTZ,
224                             created_ts TIMESTAMPTZ,
225                             PRIMARY KEY(entity_key, feature_name)
226                         );
227                         CREATE INDEX IF NOT EXISTS {} ON {} (entity_key);
228                         """
229                     ).format(
230                         sql.Identifier(table_name),
231                         sql.Identifier(f"{table_name}_ek"),
232                         sql.Identifier(table_name),
233                     )
234                 )
235 
236             conn.commit()
237 
238     def teardown(
239         self,
240         config: RepoConfig,
241         tables: Sequence[FeatureView],
242         entities: Sequence[Entity],
243     ):
244         project = config.project
245         try:
246             with self._get_conn(config) as conn, conn.cursor() as cur:
247                 for table in tables:
248                     table_name = _table_id(project, table)
249                     cur.execute(_drop_table_and_index(table_name))
250         except Exception:
251             logging.exception("Teardown failed")
252             raise
253 
254 
255 def _table_id(project: str, table: FeatureView) -> str:
256     return f"{project}_{table.name}"
257 
258 
259 def _drop_table_and_index(table_name):
260     return sql.SQL(
261         """
262         DROP TABLE IF EXISTS {};
263         DROP INDEX IF EXISTS {};
264         """
265     ).format(
266         sql.Identifier(table_name),
267         sql.Identifier(f"{table_name}_ek"),
268     )
269 
270 
271 def _to_naive_utc(ts: datetime):
272     if ts.tzinfo is None:
273         return ts
274     else:
275         return ts.astimezone(pytz.utc).replace(tzinfo=None)
276
```
--- END FILES ---

Please first localize the bug based on the issue statement, and then generate a patch according to the `git diff` format fenced by three backticks.

Here is an example:

```diff
diff --git a/examples/server_async.py b/examples/server_async.py
--- a/examples/server_async.py
+++ b/examples/server_async.py
@@ -313,4 +313,4 @@
 
 
 if __name__ == "__main__":
-    asyncio.run(run_async_server("."), debug=True)
+    asyncio.run(run_async_server(), debug=True)
diff --git a/examples/server_sync.py b/examples/server_sync.py
--- a/examples/server_sync.py
+++ b/examples/server_sync.py
@@ -313,5 +313,5 @@
 
 
 if __name__ == "__main__":
-    server = run_sync_server(".")
+    server = run_sync_server()
     server.shutdown()

```

diff --git a/sdk/python/feast/infra/online_stores/contrib/postgres.py b/sdk/python/feast/infra/online_stores/contrib/postgres.py
--- a/sdk/python/feast/infra/online_stores/contrib/postgres.py
+++ b/sdk/python/feast/infra/online_stores/contrib/postgres.py
@@ -99,6 +99,7 @@
                     cur_batch,
                     page_size=batch_size,
                 )
+                conn.commit()
                 if progress:
                     progress(len(cur_batch))

{"golden_diff": "diff --git a/sdk/python/feast/infra/online_stores/contrib/postgres.py b/sdk/python/feast/infra/online_stores/contrib/postgres.py\n--- a/sdk/python/feast/infra/online_stores/contrib/postgres.py\n+++ b/sdk/python/feast/infra/online_stores/contrib/postgres.py\n@@ -99,6 +99,7 @@\n                     cur_batch,\n                     page_size=batch_size,\n                 )\n+                conn.commit()\n                 if progress:\n                     progress(len(cur_batch))\n", "issue": "Table in Postgres OnlineStore is not populated after calling `materialize`\n## Expected Behavior\r\nWhen calling the `materialize` functionality to materialize data from a `SnowflakeSource` offline store to a local `PostgreSQLOnlineStore`, the table is not populated with the data.\r\n\r\n## Current Behavior\r\nThe feature table in the local Postgres instance is not populated, while no exception is raised, and from the logs it seems like the data should be pushed to Postgres.\r\n\r\n## Steps to reproduce\r\n\r\n1) Use this feature_store.yaml file:\r\n\r\n```\r\nproject: my_project\r\nprovider: local\r\nregistry:\r\n    registry_type: sql\r\n    path: postgresql://postgres:test@0.0.0.0:5432/feature_store\r\n    cache_ttl_seconds: 60\r\nonline_store:\r\n   type: postgres\r\n   host: 0.0.0.0\r\n   port: 5432\r\n   database: feature_store\r\n   db_schema: public\r\n   user: postgres\r\n   password: test\r\noffline_store:\r\n    <SNOWFLAKE_INFORMATION>\r\nentity_key_serialization_version: 2\r\n```\r\n\r\n2) Spin up this docker-compose file:\r\n\r\n```\r\n---\r\nversion: \"3\"\r\nservices:\r\n  db:\r\n    restart: always\r\n    image: postgres:15-alpine\r\n    container_name: feast_db\r\n    ports:\r\n      - \"5432:5432\"\r\n    volumes:\r\n      - ~/feast_postgres_data:/var/lib/postgresql/data\r\n    environment:\r\n      - POSTGRES_DB=feature_store\r\n      - POSTGRES_USER=postgres\r\n      - POSTGRES_PASSWORD=test\r\nvolumes:\r\n  feast_postgres_data: null\r\n```\r\n\r\n3) Initialize the Entities, SnowflakeSource (or another source), FeatureView, and FeatureService, and apply these. All using the Python SDK.\r\n\r\n```\r\nfrom datetime import timedelta\r\n\r\nfrom feast import (\r\n    Entity,\r\n    FeatureService,\r\n    FeatureView,\r\n    Field,\r\n    SnowflakeSource,\r\n    ValueType,\r\n    FeatureStore,\r\n)\r\nfrom feast.types import Float32\r\n\r\nfeature_store = FeatureStore()\r\n\r\nentity = Entity(\r\n    name=\"entity\",\r\n    join_keys=[\"entity_ID\"],\r\n    value_type=ValueType.STRING,\r\n)\r\n\r\nsource = SnowflakeSource(\r\n    name=\"snowflake_source_name\",\r\n    timestamp_field=\"EVENT_TIMESTAMP\",\r\n    schema=\"TEMP\",\r\n    table=\"TABLE\"\r\n)\r\n\r\nfeature_view = FeatureView(\r\n    name=\"feature_view_name\",\r\n    entities=[entity],\r\n    ttl=timedelta(days=0),\r\n    schema=[\r\n        Field(name=\"feature_1\", dtype=Float32),\r\n        Field(name=\"feature_2\", dtype=Float32),\r\n    ],\r\n    online=True,\r\n    source=source,\r\n    tags={\"team\": \"team\"},\r\n)\r\n\r\nfeature_service = FeatureService(\r\n    name=\"feature_service\",\r\n    features=[feature_view],\r\n)\r\n\r\nfeature_store.apply(\r\n    [\r\n        entity,\r\n        source,\r\n        feature_view,\r\n        feature_service,\r\n    ]\r\n)\r\n```\r\n\r\n4) Run materialize commands using the Python SDK\r\n\r\n```\r\nfeature_store = FeatureStore()\r\nfeature_store.materialize(\r\n    start_date=datetime.utcnow() - timedelta(weeks=52),\r\n    end_date=datetime.utcnow(),\r\n    feature_views=[\"feature_view_name\"],\r\n)\r\n```\r\n\r\n### Specifications\r\n- Version: 0.35.0\r\n- Platform: Local MacBook M1\r\n\r\n## Possible Solution\r\nIt seems like a `conn.commit()` statement is missing in the `online_write_batch` method of the `PostgreSQLOnlineStore`. Specifically, on [this line](https://github.com/feast-dev/feast/blob/master/sdk/python/feast/infra/online_stores/contrib/postgres.py#L102).\r\nAfter adding this, the table is populated.\r\n\r\nThe PR implementing this proposed fix can be found [here](https://github.com/feast-dev/feast/pull/3904).\r\n\r\n## Additional notes\r\nWhen replacing the the postgres online store with the following sqlite online store in the config file, everything works without any code changes\r\n\r\n```\r\nonline_store:\r\n    type: sqlite\r\n    path: data/online_store.db\r\n```\n", "before_files": [{"content": "import contextlib\nimport logging\nfrom collections import defaultdict\nfrom datetime import datetime\nfrom typing import Any, Callable, Dict, List, Optional, Sequence, Tuple\n\nimport psycopg2\nimport pytz\nfrom psycopg2 import sql\nfrom psycopg2.extras import execute_values\nfrom psycopg2.pool import SimpleConnectionPool\nfrom pydantic.schema import Literal\n\nfrom feast import Entity\nfrom feast.feature_view import FeatureView\nfrom feast.infra.key_encoding_utils import serialize_entity_key\nfrom feast.infra.online_stores.online_store import OnlineStore\nfrom feast.infra.utils.postgres.connection_utils import _get_conn, _get_connection_pool\nfrom feast.infra.utils.postgres.postgres_config import ConnectionType, PostgreSQLConfig\nfrom feast.protos.feast.types.EntityKey_pb2 import EntityKey as EntityKeyProto\nfrom feast.protos.feast.types.Value_pb2 import Value as ValueProto\nfrom feast.repo_config import RepoConfig\nfrom feast.usage import log_exceptions_and_usage\n\n\nclass PostgreSQLOnlineStoreConfig(PostgreSQLConfig):\n    type: Literal[\"postgres\"] = \"postgres\"\n\n\nclass PostgreSQLOnlineStore(OnlineStore):\n    _conn: Optional[psycopg2._psycopg.connection] = None\n    _conn_pool: Optional[SimpleConnectionPool] = None\n\n    @contextlib.contextmanager\n    def _get_conn(self, config: RepoConfig):\n        assert config.online_store.type == \"postgres\"\n        if config.online_store.conn_type == ConnectionType.pool:\n            if not self._conn_pool:\n                self._conn_pool = _get_connection_pool(config.online_store)\n            connection = self._conn_pool.getconn()\n            yield connection\n            self._conn_pool.putconn(connection)\n        else:\n            if not self._conn:\n                self._conn = _get_conn(config.online_store)\n            yield self._conn\n\n    @log_exceptions_and_usage(online_store=\"postgres\")\n    def online_write_batch(\n        self,\n        config: RepoConfig,\n        table: FeatureView,\n        data: List[\n            Tuple[EntityKeyProto, Dict[str, ValueProto], datetime, Optional[datetime]]\n        ],\n        progress: Optional[Callable[[int], Any]],\n    ) -> None:\n        project = config.project\n\n        with self._get_conn(config) as conn, conn.cursor() as cur:\n            insert_values = []\n            for entity_key, values, timestamp, created_ts in data:\n                entity_key_bin = serialize_entity_key(\n                    entity_key,\n                    entity_key_serialization_version=config.entity_key_serialization_version,\n                )\n                timestamp = _to_naive_utc(timestamp)\n                if created_ts is not None:\n                    created_ts = _to_naive_utc(created_ts)\n\n                for feature_name, val in values.items():\n                    insert_values.append(\n                        (\n                            entity_key_bin,\n                            feature_name,\n                            val.SerializeToString(),\n                            timestamp,\n                            created_ts,\n                        )\n                    )\n            # Control the batch so that we can update the progress\n            batch_size = 5000\n            for i in range(0, len(insert_values), batch_size):\n                cur_batch = insert_values[i : i + batch_size]\n                execute_values(\n                    cur,\n                    sql.SQL(\n                        \"\"\"\n                        INSERT INTO {}\n                        (entity_key, feature_name, value, event_ts, created_ts)\n                        VALUES %s\n                        ON CONFLICT (entity_key, feature_name) DO\n                        UPDATE SET\n                            value = EXCLUDED.value,\n                            event_ts = EXCLUDED.event_ts,\n                            created_ts = EXCLUDED.created_ts;\n                        \"\"\",\n                    ).format(sql.Identifier(_table_id(project, table))),\n                    cur_batch,\n                    page_size=batch_size,\n                )\n                if progress:\n                    progress(len(cur_batch))\n\n    @log_exceptions_and_usage(online_store=\"postgres\")\n    def online_read(\n        self,\n        config: RepoConfig,\n        table: FeatureView,\n        entity_keys: List[EntityKeyProto],\n        requested_features: Optional[List[str]] = None,\n    ) -> List[Tuple[Optional[datetime], Optional[Dict[str, ValueProto]]]]:\n        result: List[Tuple[Optional[datetime], Optional[Dict[str, ValueProto]]]] = []\n\n        project = config.project\n        with self._get_conn(config) as conn, conn.cursor() as cur:\n            # Collecting all the keys to a list allows us to make fewer round trips\n            # to PostgreSQL\n            keys = []\n            for entity_key in entity_keys:\n                keys.append(\n                    serialize_entity_key(\n                        entity_key,\n                        entity_key_serialization_version=config.entity_key_serialization_version,\n                    )\n                )\n\n            if not requested_features:\n                cur.execute(\n                    sql.SQL(\n                        \"\"\"\n                        SELECT entity_key, feature_name, value, event_ts\n                        FROM {} WHERE entity_key = ANY(%s);\n                        \"\"\"\n                    ).format(\n                        sql.Identifier(_table_id(project, table)),\n                    ),\n                    (keys,),\n                )\n            else:\n                cur.execute(\n                    sql.SQL(\n                        \"\"\"\n                        SELECT entity_key, feature_name, value, event_ts\n                        FROM {} WHERE entity_key = ANY(%s) and feature_name = ANY(%s);\n                        \"\"\"\n                    ).format(\n                        sql.Identifier(_table_id(project, table)),\n                    ),\n                    (keys, requested_features),\n                )\n\n            rows = cur.fetchall()\n\n            # Since we don't know the order returned from PostgreSQL we'll need\n            # to construct a dict to be able to quickly look up the correct row\n            # when we iterate through the keys since they are in the correct order\n            values_dict = defaultdict(list)\n            for row in rows if rows is not None else []:\n                values_dict[row[0].tobytes()].append(row[1:])\n\n            for key in keys:\n                if key in values_dict:\n                    value = values_dict[key]\n                    res = {}\n                    for feature_name, value_bin, event_ts in value:\n                        val = ValueProto()\n                        val.ParseFromString(bytes(value_bin))\n                        res[feature_name] = val\n                    result.append((event_ts, res))\n                else:\n                    result.append((None, None))\n\n        return result\n\n    @log_exceptions_and_usage(online_store=\"postgres\")\n    def update(\n        self,\n        config: RepoConfig,\n        tables_to_delete: Sequence[FeatureView],\n        tables_to_keep: Sequence[FeatureView],\n        entities_to_delete: Sequence[Entity],\n        entities_to_keep: Sequence[Entity],\n        partial: bool,\n    ):\n        project = config.project\n        schema_name = config.online_store.db_schema or config.online_store.user\n        with self._get_conn(config) as conn, conn.cursor() as cur:\n            # If a db_schema is provided, then that schema gets created if it doesn't\n            # exist. Else a schema is created for the feature store user.\n\n            cur.execute(\n                \"\"\"\n                SELECT schema_name\n                FROM information_schema.schemata\n                WHERE schema_name = %s\n                \"\"\",\n                (schema_name,),\n            )\n            schema_exists = cur.fetchone()\n            if not schema_exists:\n                cur.execute(\n                    sql.SQL(\"CREATE SCHEMA IF NOT EXISTS {} AUTHORIZATION {}\").format(\n                        sql.Identifier(schema_name),\n                        sql.Identifier(config.online_store.user),\n                    ),\n                )\n\n            for table in tables_to_delete:\n                table_name = _table_id(project, table)\n                cur.execute(_drop_table_and_index(table_name))\n\n            for table in tables_to_keep:\n                table_name = _table_id(project, table)\n                cur.execute(\n                    sql.SQL(\n                        \"\"\"\n                        CREATE TABLE IF NOT EXISTS {}\n                        (\n                            entity_key BYTEA,\n                            feature_name TEXT,\n                            value BYTEA,\n                            event_ts TIMESTAMPTZ,\n                            created_ts TIMESTAMPTZ,\n                            PRIMARY KEY(entity_key, feature_name)\n                        );\n                        CREATE INDEX IF NOT EXISTS {} ON {} (entity_key);\n                        \"\"\"\n                    ).format(\n                        sql.Identifier(table_name),\n                        sql.Identifier(f\"{table_name}_ek\"),\n                        sql.Identifier(table_name),\n                    )\n                )\n\n            conn.commit()\n\n    def teardown(\n        self,\n        config: RepoConfig,\n        tables: Sequence[FeatureView],\n        entities: Sequence[Entity],\n    ):\n        project = config.project\n        try:\n            with self._get_conn(config) as conn, conn.cursor() as cur:\n                for table in tables:\n                    table_name = _table_id(project, table)\n                    cur.execute(_drop_table_and_index(table_name))\n        except Exception:\n            logging.exception(\"Teardown failed\")\n            raise\n\n\ndef _table_id(project: str, table: FeatureView) -> str:\n    return f\"{project}_{table.name}\"\n\n\ndef _drop_table_and_index(table_name):\n    return sql.SQL(\n        \"\"\"\n        DROP TABLE IF EXISTS {};\n        DROP INDEX IF EXISTS {};\n        \"\"\"\n    ).format(\n        sql.Identifier(table_name),\n        sql.Identifier(f\"{table_name}_ek\"),\n    )\n\n\ndef _to_naive_utc(ts: datetime):\n    if ts.tzinfo is None:\n        return ts\n    else:\n        return ts.astimezone(pytz.utc).replace(tzinfo=None)\n", "path": "sdk/python/feast/infra/online_stores/contrib/postgres.py"}], "after_files": [{"content": "import contextlib\nimport logging\nfrom collections import defaultdict\nfrom datetime import datetime\nfrom typing import Any, Callable, Dict, List, Optional, Sequence, Tuple\n\nimport psycopg2\nimport pytz\nfrom psycopg2 import sql\nfrom psycopg2.extras import execute_values\nfrom psycopg2.pool import SimpleConnectionPool\nfrom pydantic.schema import Literal\n\nfrom feast import Entity\nfrom feast.feature_view import FeatureView\nfrom feast.infra.key_encoding_utils import serialize_entity_key\nfrom feast.infra.online_stores.online_store import OnlineStore\nfrom feast.infra.utils.postgres.connection_utils import _get_conn, _get_connection_pool\nfrom feast.infra.utils.postgres.postgres_config import ConnectionType, PostgreSQLConfig\nfrom feast.protos.feast.types.EntityKey_pb2 import EntityKey as EntityKeyProto\nfrom feast.protos.feast.types.Value_pb2 import Value as ValueProto\nfrom feast.repo_config import RepoConfig\nfrom feast.usage import log_exceptions_and_usage\n\n\nclass PostgreSQLOnlineStoreConfig(PostgreSQLConfig):\n    type: Literal[\"postgres\"] = \"postgres\"\n\n\nclass PostgreSQLOnlineStore(OnlineStore):\n    _conn: Optional[psycopg2._psycopg.connection] = None\n    _conn_pool: Optional[SimpleConnectionPool] = None\n\n    @contextlib.contextmanager\n    def _get_conn(self, config: RepoConfig):\n        assert config.online_store.type == \"postgres\"\n        if config.online_store.conn_type == ConnectionType.pool:\n            if not self._conn_pool:\n                self._conn_pool = _get_connection_pool(config.online_store)\n            connection = self._conn_pool.getconn()\n            yield connection\n            self._conn_pool.putconn(connection)\n        else:\n            if not self._conn:\n                self._conn = _get_conn(config.online_store)\n            yield self._conn\n\n    @log_exceptions_and_usage(online_store=\"postgres\")\n    def online_write_batch(\n        self,\n        config: RepoConfig,\n        table: FeatureView,\n        data: List[\n            Tuple[EntityKeyProto, Dict[str, ValueProto], datetime, Optional[datetime]]\n        ],\n        progress: Optional[Callable[[int], Any]],\n    ) -> None:\n        project = config.project\n\n        with self._get_conn(config) as conn, conn.cursor() as cur:\n            insert_values = []\n            for entity_key, values, timestamp, created_ts in data:\n                entity_key_bin = serialize_entity_key(\n                    entity_key,\n                    entity_key_serialization_version=config.entity_key_serialization_version,\n                )\n                timestamp = _to_naive_utc(timestamp)\n                if created_ts is not None:\n                    created_ts = _to_naive_utc(created_ts)\n\n                for feature_name, val in values.items():\n                    insert_values.append(\n                        (\n                            entity_key_bin,\n                            feature_name,\n                            val.SerializeToString(),\n                            timestamp,\n                            created_ts,\n                        )\n                    )\n            # Control the batch so that we can update the progress\n            batch_size = 5000\n            for i in range(0, len(insert_values), batch_size):\n                cur_batch = insert_values[i : i + batch_size]\n                execute_values(\n                    cur,\n                    sql.SQL(\n                        \"\"\"\n                        INSERT INTO {}\n                        (entity_key, feature_name, value, event_ts, created_ts)\n                        VALUES %s\n                        ON CONFLICT (entity_key, feature_name) DO\n                        UPDATE SET\n                            value = EXCLUDED.value,\n                            event_ts = EXCLUDED.event_ts,\n                            created_ts = EXCLUDED.created_ts;\n                        \"\"\",\n                    ).format(sql.Identifier(_table_id(project, table))),\n                    cur_batch,\n                    page_size=batch_size,\n                )\n                conn.commit()\n                if progress:\n                    progress(len(cur_batch))\n\n    @log_exceptions_and_usage(online_store=\"postgres\")\n    def online_read(\n        self,\n        config: RepoConfig,\n        table: FeatureView,\n        entity_keys: List[EntityKeyProto],\n        requested_features: Optional[List[str]] = None,\n    ) -> List[Tuple[Optional[datetime], Optional[Dict[str, ValueProto]]]]:\n        result: List[Tuple[Optional[datetime], Optional[Dict[str, ValueProto]]]] = []\n\n        project = config.project\n        with self._get_conn(config) as conn, conn.cursor() as cur:\n            # Collecting all the keys to a list allows us to make fewer round trips\n            # to PostgreSQL\n            keys = []\n            for entity_key in entity_keys:\n                keys.append(\n                    serialize_entity_key(\n                        entity_key,\n                        entity_key_serialization_version=config.entity_key_serialization_version,\n                    )\n                )\n\n            if not requested_features:\n                cur.execute(\n                    sql.SQL(\n                        \"\"\"\n                        SELECT entity_key, feature_name, value, event_ts\n                        FROM {} WHERE entity_key = ANY(%s);\n                        \"\"\"\n                    ).format(\n                        sql.Identifier(_table_id(project, table)),\n                    ),\n                    (keys,),\n                )\n            else:\n                cur.execute(\n                    sql.SQL(\n                        \"\"\"\n                        SELECT entity_key, feature_name, value, event_ts\n                        FROM {} WHERE entity_key = ANY(%s) and feature_name = ANY(%s);\n                        \"\"\"\n                    ).format(\n                        sql.Identifier(_table_id(project, table)),\n                    ),\n                    (keys, requested_features),\n                )\n\n            rows = cur.fetchall()\n\n            # Since we don't know the order returned from PostgreSQL we'll need\n            # to construct a dict to be able to quickly look up the correct row\n            # when we iterate through the keys since they are in the correct order\n            values_dict = defaultdict(list)\n            for row in rows if rows is not None else []:\n                values_dict[row[0].tobytes()].append(row[1:])\n\n            for key in keys:\n                if key in values_dict:\n                    value = values_dict[key]\n                    res = {}\n                    for feature_name, value_bin, event_ts in value:\n                        val = ValueProto()\n                        val.ParseFromString(bytes(value_bin))\n                        res[feature_name] = val\n                    result.append((event_ts, res))\n                else:\n                    result.append((None, None))\n\n        return result\n\n    @log_exceptions_and_usage(online_store=\"postgres\")\n    def update(\n        self,\n        config: RepoConfig,\n        tables_to_delete: Sequence[FeatureView],\n        tables_to_keep: Sequence[FeatureView],\n        entities_to_delete: Sequence[Entity],\n        entities_to_keep: Sequence[Entity],\n        partial: bool,\n    ):\n        project = config.project\n        schema_name = config.online_store.db_schema or config.online_store.user\n        with self._get_conn(config) as conn, conn.cursor() as cur:\n            # If a db_schema is provided, then that schema gets created if it doesn't\n            # exist. Else a schema is created for the feature store user.\n\n            cur.execute(\n                \"\"\"\n                SELECT schema_name\n                FROM information_schema.schemata\n                WHERE schema_name = %s\n                \"\"\",\n                (schema_name,),\n            )\n            schema_exists = cur.fetchone()\n            if not schema_exists:\n                cur.execute(\n                    sql.SQL(\"CREATE SCHEMA IF NOT EXISTS {} AUTHORIZATION {}\").format(\n                        sql.Identifier(schema_name),\n                        sql.Identifier(config.online_store.user),\n                    ),\n                )\n\n            for table in tables_to_delete:\n                table_name = _table_id(project, table)\n                cur.execute(_drop_table_and_index(table_name))\n\n            for table in tables_to_keep:\n                table_name = _table_id(project, table)\n                cur.execute(\n                    sql.SQL(\n                        \"\"\"\n                        CREATE TABLE IF NOT EXISTS {}\n                        (\n                            entity_key BYTEA,\n                            feature_name TEXT,\n                            value BYTEA,\n                            event_ts TIMESTAMPTZ,\n                            created_ts TIMESTAMPTZ,\n                            PRIMARY KEY(entity_key, feature_name)\n                        );\n                        CREATE INDEX IF NOT EXISTS {} ON {} (entity_key);\n                        \"\"\"\n                    ).format(\n                        sql.Identifier(table_name),\n                        sql.Identifier(f\"{table_name}_ek\"),\n                        sql.Identifier(table_name),\n                    )\n                )\n\n            conn.commit()\n\n    def teardown(\n        self,\n        config: RepoConfig,\n        tables: Sequence[FeatureView],\n        entities: Sequence[Entity],\n    ):\n        project = config.project\n        try:\n            with self._get_conn(config) as conn, conn.cursor() as cur:\n                for table in tables:\n                    table_name = _table_id(project, table)\n                    cur.execute(_drop_table_and_index(table_name))\n        except Exception:\n            logging.exception(\"Teardown failed\")\n            raise\n\n\ndef _table_id(project: str, table: FeatureView) -> str:\n    return f\"{project}_{table.name}\"\n\n\ndef _drop_table_and_index(table_name):\n    return sql.SQL(\n        \"\"\"\n        DROP TABLE IF EXISTS {};\n        DROP INDEX IF EXISTS {};\n        \"\"\"\n    ).format(\n        sql.Identifier(table_name),\n        sql.Identifier(f\"{table_name}_ek\"),\n    )\n\n\ndef _to_naive_utc(ts: datetime):\n    if ts.tzinfo is None:\n        return ts\n    else:\n        return ts.astimezone(pytz.utc).replace(tzinfo=None)\n", "path": "sdk/python/feast/infra/online_stores/contrib/postgres.py"}]}