Use context manager to close connections

README.md

@@ -35,8 +35,45 @@ python postgres/chunk_parquet.py
 
 ## Usage as a module
 
-The module can be installed via `pip` directly from Gtihub:
+The module can be installed via `pip` directly from Github:
 
 ```
 pip install "git+https://github.com/worldbank/DECAT_Space2Stats.git#subdirectory=space2stats_api/src
 ```
+
+It can then be used within Python as such:
+
+```py
+from space2stats import StatsTable
+
+with StatsTable.connect() as stats_table:
+    ...
+```
+
+Connection parameters may be explicitely provided. Otherwise, connection parameters will expected to be available via standard [PostgreSQL Environment Variables](https://www.postgresql.org/docs/current/libpq-envars.html#LIBPQ-ENVARS).
+
+```py
+from space2stats import StatsTable
+
+with StatsTable.connect(
+    PGHOST="localhost",
+    PGPORT="5432",
+    PGUSER="postgres",
+    PGPASSWORD="changeme",
+    PGDATABASE="postgis",
+    PGTABLENAME="space2stats",
+) as stats_table:
+    ...
+
+# alternatively:
+# settings = Settings(
+#     PGHOST="localhost",
+#     PGPORT="5432",
+#     PGUSER="postgres",
+#     PGPASSWORD="changeme",
+#     PGDATABASE="postgis",
+#     PGTABLENAME="space2stats",
+# )
+# with StatsTable.connect(settings):
+#     ...
+```

space2stats_api/src/space2stats/__init__.py

@@ -1,6 +1,7 @@
 """space2stats."""
 
 from .lib import StatsTable
+from .settings import Settings
 
-__all__ = ["StatsTable"]
+__all__ = ["StatsTable", "Settings"]
 __version__ = "1.0.0"

space2stats_api/src/space2stats/lib.py

@@ -19,19 +19,23 @@ class StatsTable:
     def connect(cls, settings: Optional[Settings] = None, **kwargs) -> "StatsTable":
         """
         Helper method to connect to the database and return a StatsTable instance.
-        It is left up to the caller to close the connection when finished, eg:
+
         ```py
-        stats_table = StatsTable.connect()
-        try:
-            stats_table.summaries(aoi, spatial_join_method, fields)
-        finally:
-            stats_table.conn.close()
+        with StatsTable.connect() as stats_table:
+            stats_table.fields()
         ```
         """
-        settings = settings or Settings(**kwargs)
+        settings = settings or Settings(**kwargs, _extra="forbid")
         conn = pg.connect(settings.DB_CONNECTION_STRING)
         return cls(conn=conn, table_name=settings.PGTABLENAME)
 
+    def __enter__(self) -> "StatsTable":
+        return self
+
+    def __exit__(self, exc_type, exc_value, traceback) -> None:
+        if self.conn:
+            self.conn.close()
+
     def _get_summaries(self, fields: List[str], h3_ids: List[str]):
         colnames = ["hex_id"] + fields
         cols = [pg.sql.Identifier(c) for c in colnames]