sqlit/tests/helpers.py at 68b5a2e5d5d02f3c930b72bf5bca528995e1332c · Maxteabag/sqlit · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
"""Test helpers for building ConnectionConfig instances with legacy-style inputs."""

from __future__ import annotations

from collections.abc import Mapping
from typing import Any

from sqlit.domains.connections.domain.config import ConnectionConfig as _ConnectionConfig
from sqlit.domains.connections.domain.config import FileEndpoint, TcpEndpoint, TunnelConfig


class _ConnectionConfigFactory:
    def __call__(self, **kwargs: Any) -> _ConnectionConfig:
        name = str(kwargs.pop("name", ""))
        db_type = str(kwargs.pop("db_type", "mssql"))
        source = kwargs.pop("source", None)
        connection_url = kwargs.pop("connection_url", None)
        folder_path = str(kwargs.pop("folder_path", "")) if "folder_path" in kwargs else ""

        extra_options = dict(kwargs.pop("extra_options", {}) or {})
        options = dict(kwargs.pop("options", {}) or {})

        # Capture legacy top-level auth fields into options.
        for key in ("auth_type", "trusted_connection"):
            if key in kwargs and key not in options:
                options[key] = kwargs.pop(key)

        endpoint = kwargs.pop("endpoint", None)
        tunnel = kwargs.pop("tunnel", None)

        file_path = kwargs.pop("file_path", None)
        if file_path is None:
            file_path = options.pop("file_path", None)

        if endpoint is None:
            if file_path:
                endpoint = FileEndpoint(path=str(file_path))
            else:
                endpoint = TcpEndpoint(
                    host=str(kwargs.pop("server", "")),
                    port=str(kwargs.pop("port", "")),
                    database=str(kwargs.pop("database", "")),
                    username=str(kwargs.pop("username", "")),
                    password=kwargs.pop("password", None),
                    password_command=kwargs.pop("password_command", None),
                )

        if tunnel is None:
            ssh_enabled = kwargs.pop("ssh_enabled", None)
            ssh_host = str(kwargs.pop("ssh_host", ""))
            ssh_port = str(kwargs.pop("ssh_port", "22"))
            ssh_username = str(kwargs.pop("ssh_username", ""))
            ssh_auth_type = str(kwargs.pop("ssh_auth_type", "key"))
            ssh_password = kwargs.pop("ssh_password", None)
            ssh_password_command = kwargs.pop("ssh_password_command", None)
            ssh_key_path = str(kwargs.pop("ssh_key_path", ""))

            enabled_flag = str(ssh_enabled).lower() if ssh_enabled is not None else ""
            if ssh_host or enabled_flag in {"enabled", "true", "1", "yes"}:
                tunnel = TunnelConfig(
                    enabled=True,
                    host=ssh_host,
                    port=ssh_port or "22",
                    username=ssh_username,
                    auth_type=ssh_auth_type or "key",
                    password=ssh_password,
                    password_command=ssh_password_command,
                    key_path=ssh_key_path,
                )

        # Any remaining kwargs get folded into options (legacy behavior).
        for key, value in list(kwargs.items()):
            if key not in options:
                options[key] = value

        return _ConnectionConfig(
            name=name,
            db_type=db_type,
            endpoint=endpoint,
            tunnel=tunnel,
            source=source,
            connection_url=connection_url,
            folder_path=folder_path,
            extra_options=extra_options,
            options=options,
        )

    def from_dict(self, data: Mapping[str, Any]) -> _ConnectionConfig:
        if "endpoint" in data:
            return _ConnectionConfig.from_dict(data)
        return self(**dict(data))

    def __getattr__(self, name: str) -> Any:
        return getattr(_ConnectionConfig, name)


ConnectionConfig = _ConnectionConfigFactory()