Merge pull request #10 from Kaligo/DATA-5768/env_vars_in_diffa_command

[DATA-5768][DATA-3562] Support env vars passing into the diff command

Author: quocnguyendinh

src/diffa/config.py

@@ -2,8 +2,7 @@
 import json
 from datetime import date
 from enum import Enum
-
-import dsnparse
+from urllib.parse import urlparse
 
 from diffa.utils import Logger
 
@@ -51,12 +50,15 @@ def update(self, **kwargs):
 
     def _parse_db_info(self):
         try:
-            dns = dsnparse.parse(self.db_uri)
+            # Users can ref an env var in the DB_URI in the command
+            if self.db_uri.startswith("$"):
+                self.db_uri = os.getenv(self.db_uri[1:], "")
+            dns = urlparse(self.db_uri)
             parsed_db_info = self._extract_db_details(dns)
             self._validate_parsed_db_info(parsed_db_info)
             return parsed_db_info
-        except TypeError:
-            logger.error("Invalid db info", exc_info=True)
+        except Exception as e:
+            logger.error(f"Error parsing DB info: {e}")
             raise
 
     def _validate_parsed_db_info(self, parsed_db_info: dict):
@@ -67,17 +69,17 @@ def _validate_parsed_db_info(self, parsed_db_info: dict):
                 )
 
     def _extract_db_details(self, dns):
-        db_database = self.db_name or dns.database
+        db_database = self.db_name or dns.path.lstrip("/")
         return {
-            "host": dns.host,
+            "host": dns.hostname,
             "scheme": dns.scheme,
             "port": dns.port,
             "database": db_database,
             "user": dns.username,
             "password": dns.password,
             "schema": self.db_schema,
             "table": self.db_table,
-            "db_uri": f"{dns.scheme}://{dns.username}:{dns.password}@{dns.host}:{dns.port}/{db_database}",
+            "db_uri": f"{dns.scheme}://{dns.username}:{dns.password}@{dns.hostname}:{dns.port}/{db_database}",
         }
 
     def get_db_config(self):

tests/configs/test_config.py

@@ -14,7 +14,7 @@
     DIFFA_CHECK_RUNS_TABLE,
 )
 
-
+@patch.dict(os.environ, {"SOURCE_URI": "postgresql://user:password@localhost:5432/mydb"})
 @pytest.mark.parametrize(
     "db_uri, db_name, db_schema, db_table, expected_parsed_config",
     [
@@ -72,6 +72,42 @@
                 "db_uri": "postgresql://user:password@localhost:5432/mydb",
             },
         ),
+        # Case 4: Complete db_uri but inputting using env var, missing database parameter
+        (
+            "$SOURCE_URI",
+            None,
+            "myschema",
+            "users",
+            {
+                "host": "localhost",
+                "scheme": "postgresql",
+                "port": 5432,
+                "database": "mydb",
+                "user": "user",
+                "password": "password",
+                "schema": "myschema",
+                "table": "users",
+                "db_uri": "postgresql://user:password@localhost:5432/mydb",
+            },
+        ),
+        # Case 5: Complete db_uri but email-format username, missing database parameter
+        (
+            "postgresql://[email protected]:password@localhost:5432/mydb",
+            None,
+            "myschema",
+            "users",
+            {
+                "host": "localhost",
+                "scheme": "postgresql",
+                "port": 5432,
+                "database": "mydb",
+                "user": "[email protected]",
+                "password": "password",
+                "schema": "myschema",
+                "table": "users",
+                "db_uri": "postgresql://[email protected]:password@localhost:5432/mydb",
+            },
+        ),
     ],
 )
 def test_db_config_parse_db_info(