Improvements to database connection

src/too/__main__.py

@@ -53,7 +53,7 @@
 )
 @server_options.option(
     "--user",
-    default=None,
+    default="sdss",
     help="The user to connect to the database.",
 )
 @click.option(
@@ -74,7 +74,7 @@ def too_cli(
     dbname: str = "sdss5db",
     host: str = "localhost",
     port: int | None = None,
-    user: str | None = None,
+    user: str = "sdss",
     verbose: bool = False,
     write_log: str | None = None,
 ):

src/too/database.py

@@ -8,6 +8,7 @@
 
 from __future__ import annotations
 
+import os
 import pathlib
 
 import polars
@@ -30,16 +31,22 @@
     "get_active_targets",
 ]
 
+DEFAULT_USER: str = "sdss"
+DEFAULT_HOST = "localhost"
+
 
 def connect_to_database(
     dbname: str,
-    host: str = "localhost",
+    host: str | None = None,
     port: int | None = None,
     user: str | None = None,
     password: str | None = None,
 ):
     """Connects the ``sdssdb`` ``sdss5db`` models to the database."""
 
+    if port is None:
+        port = int(os.environ.get("PGPORT", 5432))
+
     catalogdb.database.connect(
         dbname,
         host=host,
@@ -56,21 +63,24 @@ def connect_to_database(
 
 def get_database_uri(
     dbname: str,
-    host: str = "localhost",
+    host: str | None = None,
     port: int | None = None,
     user: str | None = None,
     password: str | None = None,
 ):
     """Returns the URI to the database."""
 
-    if user is None and password is None:
+    user = user or str(os.environ.get("PGUSER", DEFAULT_USER))
+    host = host or str(os.environ.get("PGHOST", DEFAULT_HOST))
+    port = port or int(os.environ.get("PGPORT", 5432))
+
+    if user is None and password is None:  # pragma: no cover
+        # This should never happen. user is always set.
         auth: str = ""
-    elif user is not None and password is None:
+    elif password is None:
         auth: str = f"{user}@"
-    elif user is not None and password is not None:
-        auth: str = f"{user}:{password}@"
     else:
-        raise ValueError("Passing a password requires also passing a user.")
+        auth: str = f"{user}:{password}@"
 
     host_port: str = f"{host or ''}" if port is None else f"{host or ''}:{port}"
 

tests/scripts/sdss5db_too_test.sql

@@ -247,14 +247,14 @@ CREATE TABLE targetdb.cadence (
     label_version TEXT DEFAULT ''::TEXT
 );
 
-\copy catalogdb.catalog FROM PROGRAM 'gzip -dc catalog.csv.gz' WITH CSV HEADER;
-\copy catalogdb.sdss_id_stacked FROM PROGRAM 'gzip -dc sdss_id_stacked.csv.gz' WITH CSV HEADER;
-\copy catalogdb.catalog_to_gaia_dr3_source FROM PROGRAM 'gzip -dc catalog_to_gaia_dr3_source.csv.gz' WITH CSV HEADER;
--- \copy catalogdb.catalog_to_sdss_dr13_photoobj_primary FROM PROGRAM 'gzip -dc catalog_to_sdss_dr13_photoobj_primary.csv.gz' WITH CSV HEADER;
-\copy catalogdb.catalog_to_twomass_psc FROM PROGRAM 'gzip -dc catalog_to_twomass_psc.csv.gz' WITH CSV HEADER;
-\copy catalogdb.gaia_dr3_source FROM PROGRAM 'gzip -dc gaia_dr3_source.csv.gz' WITH CSV HEADER;
--- \copy catalogdb.sdss_dr13_photoobj FROM PROGRAM 'gzip -dc sdss_dr13_photoobj.csv.gz' WITH CSV HEADER;
-\copy catalogdb.twomass_psc FROM PROGRAM 'gzip -dc twomass_psc.csv.gz' WITH CSV HEADER;
+\copy catalogdb.catalog FROM PROGRAM '/usr/bin/gzip -dc catalog.csv.gz' WITH CSV HEADER;
+\copy catalogdb.sdss_id_stacked FROM PROGRAM '/usr/bin/gzip -dc sdss_id_stacked.csv.gz' WITH CSV HEADER;
+\copy catalogdb.catalog_to_gaia_dr3_source FROM PROGRAM '/usr/bin/gzip -dc catalog_to_gaia_dr3_source.csv.gz' WITH CSV HEADER;
+-- \copy catalogdb.catalog_to_sdss_dr13_photoobj_primary FROM PROGRAM '/usr/bin/gzip -dc catalog_to_sdss_dr13_photoobj_primary.csv.gz' WITH CSV HEADER;
+\copy catalogdb.catalog_to_twomass_psc FROM PROGRAM '/usr/bin/gzip -dc catalog_to_twomass_psc.csv.gz' WITH CSV HEADER;
+\copy catalogdb.gaia_dr3_source FROM PROGRAM '/usr/bin/gzip -dc gaia_dr3_source.csv.gz' WITH CSV HEADER;
+-- \copy catalogdb.sdss_dr13_photoobj FROM PROGRAM '/usr/bin/gzip -dc sdss_dr13_photoobj.csv.gz' WITH CSV HEADER;
+\copy catalogdb.twomass_psc FROM PROGRAM '/usr/bin/gzip -dc twomass_psc.csv.gz' WITH CSV HEADER;
 
 ALTER TABLE catalogdb.catalog ADD PRIMARY KEY (catalogid);
 ALTER TABLE catalogdb.sdss_id_stacked ADD PRIMARY KEY (sdss_id);

tests/test_database.py

@@ -8,6 +8,8 @@
 
 from __future__ import annotations
 
+import os
+
 import polars
 import pytest
 from conftest import DBNAME
@@ -25,6 +27,9 @@
 from too.exceptions import ValidationError
 
 
+PGPORT = os.environ.get("PGPORT", 5432)
+
+
 def test_database_exists():
     assert catalogdb.Catalog.table_exists()
     assert catalogdb.database.dbname == DBNAME
@@ -61,8 +66,8 @@ def test_models_exist():
 @pytest.mark.parametrize(
     "dbname,user,password,host,port,expected",
     [
-        ("testdb", None, None, "localhost", None, "localhost/testdb"),
-        ("testdb", "user", "1234", "localhost", None, "user:1234@localhost/testdb"),
+        ("testdb", None, None, "localhost", PGPORT, f"sdss@localhost:{PGPORT}/testdb"),
+        ("db", "user", "1234", "10.1.1.1", PGPORT, f"user:1234@10.1.1.1:{PGPORT}/db"),
         ("testdb", "user", None, "localhost", 5432, "user@localhost:5432/testdb"),
     ],
 )
@@ -86,11 +91,6 @@ def test_get_database_uri(
     assert uri == f"postgresql://{expected}"
 
 
-def test_get_database_uri_password_fails():
-    with pytest.raises(ValueError):
-        get_database_uri("testdb", password="1234")
-
-
 def test_validate_too_target_passes(too_mock: polars.DataFrame):
     assert isinstance(validate_too_targets(too_mock), polars.DataFrame)