From a8e63bd71a43866846143881cded6f9a130b64da Mon Sep 17 00:00:00 2001 From: Kirill Karmadonov Date: Tue, 12 Nov 2024 09:29:47 +0100 Subject: [PATCH 1/7] Add support for Microsoft Entra ID Authentication --- bcpandas/main.py | 8 +++++++- bcpandas/utils.py | 2 ++ 2 files changed, 9 insertions(+), 1 deletion(-) diff --git a/bcpandas/main.py b/bcpandas/main.py index 64a37de..67baabe 100644 --- a/bcpandas/main.py +++ b/bcpandas/main.py @@ -39,8 +39,9 @@ class SqlCreds: engine that uses `pyodbc` as the DBAPI, and store it in the `self.engine` attribute. If `username` and `password` are not provided, `with_krb_auth` will be `True`. + If `entra_id_token` are provided uses Microsoft Entra ID Authentication. - Only supports SQL based logins, not Active Directory or Azure AD. + Only supports SQL based logins and Microsoft Entra ID, not Active Directory. Parameters ---------- @@ -54,6 +55,8 @@ class SqlCreds: odbc_kwargs : dict of {str, str or int}, optional additional keyword arguments, to pass into ODBC connection string, such as Encrypted='yes' + entra_id_token: str, optional + Microsoft Entra ID Authentication token Returns ------- @@ -69,6 +72,7 @@ def __init__( driver_version: Optional[int] = None, port: int = 1433, odbc_kwargs: Optional[Dict[str, Union[str, int]]] = None, + entra_id_token: Optional[str] = None, ): self.server = server self.database = database @@ -107,6 +111,8 @@ def __init__( self.with_krb_auth = True db_url += "Trusted_Connection=yes;" + self.entra_id_token = entra_id_token + self_msg = sub(r"password=\'.*\'", "password=[REDACTED]", str(self)) logger.info(f"Created creds:\t{self_msg}") diff --git a/bcpandas/utils.py b/bcpandas/utils.py index 9111b52..5b3250f 100644 --- a/bcpandas/utils.py +++ b/bcpandas/utils.py @@ -71,6 +71,8 @@ def bcp( # auth if creds.with_krb_auth: auth = ["-T"] + elif creds.entra_id_token: + auth = ["-G", "-P", quote_this(creds.entra_id_token)] else: auth = ["-U", quote_this(creds.username), "-P", quote_this(creds.password)] if creds.odbc_kwargs: From 4913a0efacff9ee5b76783274c98af9d4f2fe810 Mon Sep 17 00:00:00 2001 From: Kirill Karmadonov Date: Tue, 12 Nov 2024 10:12:54 +0100 Subject: [PATCH 2/7] Update a few tests --- tests/test_utils.py | 38 ++++++++++++++++++++++++++++++++++++++ 1 file changed, 38 insertions(+) diff --git a/tests/test_utils.py b/tests/test_utils.py index 4e1558f..ca3ba21 100644 --- a/tests/test_utils.py +++ b/tests/test_utils.py @@ -28,6 +28,7 @@ def test_bcpandas_creates_command_without_port_if_default(run_cmd): username="me", password="secret", odbc_kwargs=None, + entra_id_token=None, ) utils.bcp("table", "in", "", creds, True) assert run_cmd.call_args == mock.call( @@ -60,6 +61,7 @@ def test_bcpandas_creates_command_with_port_if_not_default(run_cmd): username="me", password="secret", odbc_kwargs=None, + entra_id_token=None, ) utils.bcp("table", "in", "", creds, True) assert run_cmd.call_args == mock.call( @@ -92,6 +94,7 @@ def test_bcpandas_creates_command_with_encrypt_no(run_cmd): username="me", password="secret", odbc_kwargs=dict(encrypt="no"), + entra_id_token=None, ) utils.bcp("table", "in", "", creds, True) assert run_cmd.call_args == mock.call( @@ -125,6 +128,7 @@ def test_bcpandas_creates_command_with_encrypt_yes(run_cmd): username="me", password="secret", odbc_kwargs=dict(Encrypt="1"), + entra_id_token=None, ) utils.bcp("table", "in", "", creds, True) assert run_cmd.call_args == mock.call( @@ -148,6 +152,40 @@ def test_bcpandas_creates_command_with_encrypt_yes(run_cmd): ) +def test_bcpandas_creates_command_with_entra_id_token(run_cmd): + Creds = namedtuple("Creds", "server port database with_krb_auth username password odbc_kwargs") + creds = Creds( + server="localhost", + port=1433, + database="DB", + with_krb_auth=False, + username=None, + password=None, + odbc_kwargs=dict(Encrypt="1"), + entra_id_token="secret_token", + ) + utils.bcp("table", "in", "", creds, True) + assert run_cmd.call_args == mock.call( + [ + "bcp", + "dbo.table", + "in", + "", + "-S", + "localhost", + "-d", + "DB", + "-q", + "", + "-G" + "-P", + "secret_token", + ] + + (["-Ym"] if sys.platform != "win32" else []), + print_output=True, + ) + + @pytest.mark.usefixtures("database") def test_bcp_login_failure(sql_creds: SqlCreds): wrong_sql_creds = SqlCreds( From 532ab3f817177adcadc83a0cdfc4e645094412d1 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Tue, 12 Nov 2024 09:13:23 +0000 Subject: [PATCH 3/7] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- tests/test_utils.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/tests/test_utils.py b/tests/test_utils.py index ca3ba21..92d8d33 100644 --- a/tests/test_utils.py +++ b/tests/test_utils.py @@ -177,8 +177,7 @@ def test_bcpandas_creates_command_with_entra_id_token(run_cmd): "DB", "-q", "", - "-G" - "-P", + "-G" "-P", "secret_token", ] + (["-Ym"] if sys.platform != "win32" else []), From 7fd3f2f6345ea525188caeaa2dbcdcf05d1612fb Mon Sep 17 00:00:00 2001 From: Kirill Karmadonov Date: Tue, 12 Nov 2024 10:25:02 +0100 Subject: [PATCH 4/7] Update a few tests --- tests/test_utils.py | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/tests/test_utils.py b/tests/test_utils.py index ca3ba21..4911653 100644 --- a/tests/test_utils.py +++ b/tests/test_utils.py @@ -19,7 +19,7 @@ def fixture_run_cmd_capture(monkeypatch): def test_bcpandas_creates_command_without_port_if_default(run_cmd): - Creds = namedtuple("Creds", "server port database with_krb_auth username password odbc_kwargs") + Creds = namedtuple("Creds", "server port database with_krb_auth username password odbc_kwargs entra_id_token") creds = Creds( server="localhost", port=1433, @@ -52,7 +52,7 @@ def test_bcpandas_creates_command_without_port_if_default(run_cmd): def test_bcpandas_creates_command_with_port_if_not_default(run_cmd): - Creds = namedtuple("Creds", "server port database with_krb_auth username password odbc_kwargs") + Creds = namedtuple("Creds", "server port database with_krb_auth username password odbc_kwargs entra_id_token") creds = Creds( server="localhost", port=1234, @@ -85,7 +85,7 @@ def test_bcpandas_creates_command_with_port_if_not_default(run_cmd): def test_bcpandas_creates_command_with_encrypt_no(run_cmd): - Creds = namedtuple("Creds", "server port database with_krb_auth username password odbc_kwargs") + Creds = namedtuple("Creds", "server port database with_krb_auth username password odbc_kwargs entra_id_token") creds = Creds( server="localhost", port=1433, @@ -119,7 +119,7 @@ def test_bcpandas_creates_command_with_encrypt_no(run_cmd): def test_bcpandas_creates_command_with_encrypt_yes(run_cmd): - Creds = namedtuple("Creds", "server port database with_krb_auth username password odbc_kwargs") + Creds = namedtuple("Creds", "server port database with_krb_auth username password odbc_kwargs entra_id_token") creds = Creds( server="localhost", port=1433, @@ -153,7 +153,7 @@ def test_bcpandas_creates_command_with_encrypt_yes(run_cmd): def test_bcpandas_creates_command_with_entra_id_token(run_cmd): - Creds = namedtuple("Creds", "server port database with_krb_auth username password odbc_kwargs") + Creds = namedtuple("Creds", "server port database with_krb_auth username password odbc_kwargs entra_id_token") creds = Creds( server="localhost", port=1433, From c6ea66c7bc63c4cb3dcb0855909ae33f3ba362ba Mon Sep 17 00:00:00 2001 From: Kirill Karmadonov Date: Tue, 12 Nov 2024 10:27:31 +0100 Subject: [PATCH 5/7] Update a few tests --- tests/test_utils.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/tests/test_utils.py b/tests/test_utils.py index b6c45f3..fe26206 100644 --- a/tests/test_utils.py +++ b/tests/test_utils.py @@ -177,7 +177,8 @@ def test_bcpandas_creates_command_with_entra_id_token(run_cmd): "DB", "-q", "", - "-G" "-P", + "-G", + "-P", "secret_token", ] + (["-Ym"] if sys.platform != "win32" else []), From 4614101b578d82c031806305ff8ba06262aed07c Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Tue, 12 Nov 2024 09:30:04 +0000 Subject: [PATCH 6/7] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- tests/test_utils.py | 20 +++++++++++++++----- 1 file changed, 15 insertions(+), 5 deletions(-) diff --git a/tests/test_utils.py b/tests/test_utils.py index fe26206..e0c5523 100644 --- a/tests/test_utils.py +++ b/tests/test_utils.py @@ -19,7 +19,9 @@ def fixture_run_cmd_capture(monkeypatch): def test_bcpandas_creates_command_without_port_if_default(run_cmd): - Creds = namedtuple("Creds", "server port database with_krb_auth username password odbc_kwargs entra_id_token") + Creds = namedtuple( + "Creds", "server port database with_krb_auth username password odbc_kwargs entra_id_token" + ) creds = Creds( server="localhost", port=1433, @@ -52,7 +54,9 @@ def test_bcpandas_creates_command_without_port_if_default(run_cmd): def test_bcpandas_creates_command_with_port_if_not_default(run_cmd): - Creds = namedtuple("Creds", "server port database with_krb_auth username password odbc_kwargs entra_id_token") + Creds = namedtuple( + "Creds", "server port database with_krb_auth username password odbc_kwargs entra_id_token" + ) creds = Creds( server="localhost", port=1234, @@ -85,7 +89,9 @@ def test_bcpandas_creates_command_with_port_if_not_default(run_cmd): def test_bcpandas_creates_command_with_encrypt_no(run_cmd): - Creds = namedtuple("Creds", "server port database with_krb_auth username password odbc_kwargs entra_id_token") + Creds = namedtuple( + "Creds", "server port database with_krb_auth username password odbc_kwargs entra_id_token" + ) creds = Creds( server="localhost", port=1433, @@ -119,7 +125,9 @@ def test_bcpandas_creates_command_with_encrypt_no(run_cmd): def test_bcpandas_creates_command_with_encrypt_yes(run_cmd): - Creds = namedtuple("Creds", "server port database with_krb_auth username password odbc_kwargs entra_id_token") + Creds = namedtuple( + "Creds", "server port database with_krb_auth username password odbc_kwargs entra_id_token" + ) creds = Creds( server="localhost", port=1433, @@ -153,7 +161,9 @@ def test_bcpandas_creates_command_with_encrypt_yes(run_cmd): def test_bcpandas_creates_command_with_entra_id_token(run_cmd): - Creds = namedtuple("Creds", "server port database with_krb_auth username password odbc_kwargs entra_id_token") + Creds = namedtuple( + "Creds", "server port database with_krb_auth username password odbc_kwargs entra_id_token" + ) creds = Creds( server="localhost", port=1433, From 01fa6ee404ed58dc776005cde8bfb38724a574d7 Mon Sep 17 00:00:00 2001 From: Kirill Karmadonov Date: Tue, 12 Nov 2024 10:39:17 +0100 Subject: [PATCH 7/7] Update a few tests --- tests/test_utils.py | 1 - 1 file changed, 1 deletion(-) diff --git a/tests/test_utils.py b/tests/test_utils.py index fe26206..e087cb9 100644 --- a/tests/test_utils.py +++ b/tests/test_utils.py @@ -176,7 +176,6 @@ def test_bcpandas_creates_command_with_entra_id_token(run_cmd): "-d", "DB", "-q", - "", "-G", "-P", "secret_token",