Merge pull request #1 from mixxorz/feature/config-file

Add TOML-based configuration

README.md

@@ -50,59 +50,90 @@ DSLR is 8x faster at taking snapshots and 3x faster at restoring snapshots compa
   ```SQL
   CREATE TABLE large_test (num1 bigint, num2 double precision, num3 double precision);
 
-  INSERT INTO large_test (num1, num2, num3)
-  SELECT round(random() * 10), random(), random() * 142
-  FROM generate_series(1, 20000000) s(i);
-  ```
-
-  I used the following commands to measure the execution time:
-
-  ```
-  time dslr snapshot my-snapshot
-  time dslr restore my-snapshot
-  time pg_dump -Fc -f export.dump
-  time pg_restore --no-acl --no-owner export.dump
-  ```
-
-  I ran each command three times and plotted the mean in the chart.
-
-  Here's the raw data:
-
-  | Command       | Run | Execution time (seconds) |
-  | ------------- | --- | ------------------------ |
-  | dslr snapshot | 1   | 4.797                    |
-  |               | 2   | 4.650                    |
-  |               | 3   | 2.927                    |
-  | dslr restore  | 1   | 5.840                    |
-  |               | 2   | 4.122                    |
-  |               | 3   | 3.331                    |
-  | pg_dump       | 1   | 37.345                   |
-  |               | 2   | 36.227                   |
-  |               | 3   | 36.233                   |
-  | pg_restore    | 1   | 13.304                   |
-  |               | 2   | 13.148                   |
-  |               | 3   | 13.320                   |
+INSERT INTO large*test (num1, num2, num3)
+SELECT round(random() * 10), random(), random() \_ 142
+FROM generate_series(1, 20000000) s(i);
+
+```
+
+I used the following commands to measure the execution time:
+
+```
+
+time dslr snapshot my-snapshot
+time dslr restore my-snapshot
+time pg_dump -Fc -f export.dump
+time pg_restore --no-acl --no-owner export.dump
+
+```
+
+I ran each command three times and plotted the mean in the chart.
+
+Here's the raw data:
+
+| Command       | Run | Execution time (seconds) |
+| ------------- | --- | ------------------------ |
+| dslr snapshot | 1   | 4.797                    |
+|               | 2   | 4.650                    |
+|               | 3   | 2.927                    |
+| dslr restore  | 1   | 5.840                    |
+|               | 2   | 4.122                    |
+|               | 3   | 3.331                    |
+| pg_dump       | 1   | 37.345                   |
+|               | 2   | 36.227                   |
+|               | 3   | 36.233                   |
+| pg_restore    | 1   | 13.304                   |
+|               | 2   | 13.148                   |
+|               | 3   | 13.320                   |
 </details>
 
 ## Install
 
 ```
+
 pip install DSLR
-```
+
+````
 
 DSLR requires that the Postgres client binaries (`psql`, `createdb`, `dropdb`)
 are present in your `PATH`. DSLR uses them to interact with Postgres.
 
-## Usage
+## Configuration
+
+You can tell DSLR which database to take snapshots of in a few ways:
+
+**PG\* environment variables**
+
+If you have the [PG* environment
+variables](https://www.postgresql.org/docs/current/libpq-envars.html) set, DSLR
+will automatically try to use these in a similar way to `psql`.
 
-First you need to point DSLR to the database you want to take snapshots of. The
-easiest way to do this is to export the `DATABASE_URL` environment variable.
+**DATABASE_URL**
+
+If the `DATABASE_URL` environment variable is set, DSLR will use this to connect
+to your target database. DSLR will prefer this over the PG* environment
+variables.
 
 ```bash
 export DATABASE_URL=postgres://username:password@host:port/database_name
+````
+
+**dslr.toml**
+
+If you have a `dslr.toml` file in the same directory where you're running
+`dslr`, DSLR will read its settings from it. DSLR will prefer this over the
+environment variables.
+
+```toml
+url: postgres://username:password@host:port/database_name
 ```
 
-Alternatively, you can pass the connection string via the `--db` option.
+**`--url` option**
+
+Finally, you can explicitly pass the connection string via the `--url` option.
+This will override any of the above settings.
+
+## Usage
 
 You're ready to use DSLR!
 

dslr/cli.py

@@ -3,6 +3,7 @@
 
 import click
 import timeago
+import tomli
 from rich import box
 from rich.table import Table
 
@@ -38,20 +39,39 @@ def complete_snapshot_names(ctx, param, incomplete):
     return []
 
 
+def next_not_none(iterable):
+    """
+    Returns the next item in the iterable that is not None or ""
+    """
+    return next((item for item in iterable if item is not None and item != ""), None)
+
+
 @click.group
 @click.option(
-    "--db",
-    envvar="DATABASE_URL",
-    required=True,
+    "--url",
     help="The database connection string to the database you want to take "
-    "snapshots of. If not provided, DSLR will try to read it from the "
-    "DATABASE_URL environment variable. "
+    "snapshots of."
     "\n\nExample: postgres://username:password@host:port/database_name",
 )
 @click.option("--debug/--no-debug", help="Show additional debugging information.")
-def cli(db, debug):
+def cli(url, debug):
+    toml_params = {}
+    try:
+        with open("dslr.toml", "rb") as tomlf:
+            toml_params = tomli.load(tomlf)
+    except FileNotFoundError:
+        pass
+
+    config = {
+        "url": next_not_none(
+            [url, toml_params.get("url"), os.environ.get("DATABASE_URL")]
+        )
+        or "",
+        "debug": next_not_none([debug, toml_params.get("debug"), False]),
+    }
+
     # Update the settings singleton
-    settings.initialize(url=db, debug=debug)
+    settings.initialize(**config)
 
 
 @cli.command()

pyproject.toml

@@ -24,6 +24,7 @@ python = "^3.7"
 click = "^8.1.3"
 rich = "^12.5.1"
 timeago = "^1.0.15"
+tomli = "^2.0.1"
 
 [tool.poetry.dev-dependencies]
 isort = "^5.10.1"

tests/test_cli.py

@@ -166,3 +166,112 @@ def test_import_overwrite(self):
             "Imported snapshot existing-snapshot-1 from pyproject.toml",
             result.output,
         )
+
+
+@mock.patch("dslr.operations.exec", new=stub_exec)
+class ConfigTest(TestCase):
+    @mock.patch.dict(
+        os.environ, {"DATABASE_URL": "postgres://envvar:pw@test:5432/my_db"}
+    )
+    @mock.patch("dslr.cli.settings")
+    @mock.patch("dslr.operations.settings")
+    def test_database_url(self, mock_operations_settings, mock_cli_settings):
+        runner = CliRunner()
+        result = runner.invoke(cli.cli, ["snapshot", "my-snapshot"])
+
+        self.assertEqual(result.exit_code, 0)
+
+        mock_cli_settings.initialize.assert_called_once_with(
+            debug=False,
+            url="postgres://envvar:pw@test:5432/my_db",
+        )
+
+    @mock.patch("dslr.cli.settings")
+    @mock.patch("dslr.operations.settings")
+    def test_toml(self, mock_operations_settings, mock_cli_settings):
+        with mock.patch(
+            "builtins.open",
+            mock.mock_open(read_data=b"url = 'postgres://toml:pw@test:5432/my_db'"),
+        ):
+            runner = CliRunner()
+            result = runner.invoke(cli.cli, ["snapshot", "my-snapshot"])
+
+        self.assertEqual(result.exit_code, 0)
+
+        mock_cli_settings.initialize.assert_called_once_with(
+            debug=False,
+            url="postgres://toml:pw@test:5432/my_db",
+        )
+
+    @mock.patch("dslr.cli.settings")
+    @mock.patch("dslr.operations.settings")
+    def test_db_option(self, mock_operations_settings, mock_cli_settings):
+        runner = CliRunner()
+        result = runner.invoke(
+            cli.cli,
+            ["--url", "postgres://cli:pw@test:5432/my_db", "snapshot", "my-snapshot"],
+        )
+
+        self.assertEqual(result.exit_code, 0)
+
+        mock_cli_settings.initialize.assert_called_once_with(
+            debug=False,
+            url="postgres://cli:pw@test:5432/my_db",
+        )
+
+    @mock.patch("dslr.cli.settings")
+    @mock.patch("dslr.operations.settings")
+    def test_settings_preference_order(
+        self, mock_operations_settings, mock_cli_settings
+    ):
+        # No options passed (e.g. PG environment variables are used)
+        runner = CliRunner()
+        result = runner.invoke(cli.cli, ["snapshot", "my-snapshot"])
+        self.assertEqual(result.exit_code, 0)
+
+        # DATABASE_URL environment variable is used
+        with mock.patch.dict(
+            os.environ, {"DATABASE_URL": "postgres://envvar:pw@test:5432/my_db"}
+        ):
+            runner = CliRunner()
+            result = runner.invoke(cli.cli, ["snapshot", "my-snapshot"])
+            self.assertEqual(result.exit_code, 0)
+
+            # TOML file is used
+            with mock.patch(
+                "builtins.open",
+                mock.mock_open(read_data=b"url = 'postgres://toml:pw@test:5432/my_db'"),
+            ):
+                runner = CliRunner()
+                result = runner.invoke(cli.cli, ["snapshot", "my-snapshot"])
+                self.assertEqual(result.exit_code, 0)
+
+                # --url option is used
+                runner = CliRunner()
+                result = runner.invoke(
+                    cli.cli,
+                    [
+                        "--url",
+                        "postgres://cli:pw@test:5432/my_db",
+                        "snapshot",
+                        "my-snapshot",
+                    ],
+                )
+                self.assertEqual(result.exit_code, 0)
+
+        # Check that the correct order of settings is used
+        self.assertEqual(4, mock_cli_settings.initialize.call_count)
+
+        self.assertEqual(
+            mock_cli_settings.initialize.call_args_list,
+            [
+                # Nothing is passed
+                mock.call(debug=False, url=""),
+                # DATABASE_URL is present so use that
+                mock.call(debug=False, url="postgres://envvar:pw@test:5432/my_db"),
+                # TOML is present, so use that over DATABASE_URL
+                mock.call(debug=False, url="postgres://toml:pw@test:5432/my_db"),
+                # --url is present, so use that over everything
+                mock.call(debug=False, url="postgres://cli:pw@test:5432/my_db"),
+            ],
+        )