feat: use .env file directly to get pg params

README.md

@@ -16,7 +16,7 @@ pip install jolteon
 
 ## How to use Jolteon
 
-1. Create a `.env` file like the `.env.example` one you find in this repository and fill it with your Lightdash database connection parameters. After that, run `source .env`.
+1. Create a `.env` file like the `.env.example` one you find in this repository and fill it with your Lightdash database connection parameters.
 
 2. Create a `config.yaml` file like the `config_example.yaml` one you find in this repository. This file should be structured as follows:
 
@@ -26,6 +26,6 @@ pip install jolteon
 
     - `fields_raw_mapping` should be filled with the mapping of the metrics and the dimensions you have changed. If you haven't changed any metric or dimension, you can also leave it empty.
 
-    - `query_ids` should be filled with the ids of the charts you want to affect when updating the database. If you don't known what are the ids of the charts (and you probably won't the first time), you can run `jolteon get-ids --connection-params-from-env`. You will be presented with a table containing the id, the name and the workspace of all the charts of your Lightdash instance.
+    - `query_ids` should be filled with the ids of the charts you want to affect when updating the database. If you don't known what are the ids of the charts (and you probably won't the first time), you can run `jolteon get-ids`. You will be presented with a table containing the id, the name and the workspace of all the charts of your Lightdash instance.
 
-3. Run `jolteon update-db config.yaml --connection-params-from-env`.
+3. Run `jolteon update-db config.yaml`.

jolteon/__main__.py

@@ -3,23 +3,23 @@
 
 from jolteon import Config
 from jolteon.modules import Updater, print_query_ids
-from jolteon.utils import ConnectionParams, get_pg_connection
+from jolteon.utils import get_connection
 
 app = typer.Typer()
 
 
 @app.command()
-def get_ids(connection_params_from_env: bool = False) -> None:
-    conn = get_pg_connection(ConnectionParams(params_from_env=connection_params_from_env))
+def get_ids() -> None:
+    conn = get_connection()
     print_query_ids(conn)
 
 
 @app.command()
-def update_db(config_path: str, connection_params_from_env: bool = False) -> None:
+def update_db(config_path: str) -> None:
     with open(config_path) as f:
         config = Config(**yaml.safe_load(f))
 
-    conn = get_pg_connection(ConnectionParams(params_from_env=connection_params_from_env))
+    conn = get_connection()
     updater = Updater(config, conn)
     updater.overwrite_db()
 

jolteon/utils.py

@@ -1,37 +1,38 @@
-import os
+import sys
 
 import pandas as pd
 import psycopg2
-from pydantic import BaseModel, root_validator
+from pydantic import BaseSettings
 
 
-class ConnectionParams(BaseModel):
-    params_from_env: bool = False
+class ConnectionParams(BaseSettings):
     host: str = "localhost"
     port: int = 5432
     database: str = "postgres"
     user: str = "postgres"
     password: str = "postgres"
 
-    @root_validator(pre=True)
-    def populate_params(cls, values: dict) -> dict:
-        if values["params_from_env"]:
-            for k, v in cls.schema()["properties"].items():
-                if k != "params_from_env":
-                    values[k] = os.getenv(f"PG_{k.upper()}", v["default"])
-        return values
-
-
-def get_pg_connection(
-    connection_params: ConnectionParams,
-) -> psycopg2.extensions.connection:
-    return psycopg2.connect(
-        host=connection_params.host,
-        port=connection_params.port,
-        database=connection_params.database,
-        user=connection_params.user,
-        password=connection_params.password,
-    )
+    class Config:
+        env_file = ".env"
+        env_prefix = "pg_"
+
+
+def get_connection() -> psycopg2.extensions.connection:
+    params = ConnectionParams()
+    while True:
+        ans = input(
+            f"You are connecting to {params.host}:{params.port}/{params.database}. Do you want to continue? y/n "
+        )
+        if ans.lower() == "y":
+            return psycopg2.connect(
+                host=params.host,
+                port=params.port,
+                database=params.database,
+                user=params.user,
+                password=params.password,
+            )
+        if ans.lower() == "n":
+            sys.exit(0)
 
 
 def get_df_from_query(query: str, conn: psycopg2.extensions.connection) -> pd.DataFrame:

pyproject.toml

@@ -12,7 +12,7 @@ jolteon = "jolteon.__main__:app"
 python = "^3.10"
 pandas = "^1.5.3"
 psycopg2-binary = "^2.9.5"
-pydantic = "^1.10.5"
+pydantic = {extras = ["dotenv"], version = "^1.10.5"}
 pyyaml = "^6.0"
 tabulate = "^0.9.0"
 typer = "^0.7.0"