feat(connections): connect to azure mssql using managed identities

luis-dk · aarthy-dk · commit b2ce30df26b1 · 2025-11-21T12:48:52.000-05:00
diff --git a/deploy/testgen-base.dockerfile b/deploy/testgen-base.dockerfile
@@ -48,7 +48,6 @@ RUN apk del \
     cmake \
     musl-dev \
     gfortran \
-    curl \
     gpg \
     linux-headers \
     openblas-dev \
diff --git a/deploy/testgen.dockerfile b/deploy/testgen.dockerfile
@@ -1,4 +1,4 @@
-ARG TESTGEN_BASE_LABEL=v8
+ARG TESTGEN_BASE_LABEL=v9
 
 FROM datakitchen/dataops-testgen-base:${TESTGEN_BASE_LABEL} AS release-image
 
diff --git a/pyproject.toml b/pyproject.toml
@@ -61,6 +61,7 @@ dependencies = [
     "cron-converter==1.2.1",
     "cron-descriptor==2.0.5",
     "pybars3==0.9.7",
+    "azure-identity==1.25.1",
 
     # Pinned to match the manually compiled libs or for security
     "pyarrow==18.1.0",
diff --git a/testgen/common/database/flavor/flavor_service.py b/testgen/common/database/flavor/flavor_service.py
@@ -22,6 +22,8 @@ class ConnectionParams(TypedDict):
     private_key_passphrase: bytes
     http_path: str
     service_account_key: dict[str,Any]
+    connect_with_identity: bool
+    sql_flavor_code: str
 
 class FlavorService:
 
@@ -49,6 +51,8 @@ def init(self, connection_params: ConnectionParams):
         self.catalog = connection_params.get("catalog") or ""
         self.warehouse = connection_params.get("warehouse") or ""
         self.service_account_key = connection_params.get("service_account_key", None)
+        self.connect_with_identity = connection_params.get("connect_with_identity") or False
+        self.sql_flavor_code = connection_params.get("sql_flavor_code") or self.flavor
 
         password = connection_params.get("project_pw_encrypted", None)
         if isinstance(password, memoryview) or isinstance(password, bytes):
diff --git a/testgen/common/database/flavor/mssql_flavor_service.py b/testgen/common/database/flavor/mssql_flavor_service.py
@@ -1,5 +1,7 @@
 from urllib.parse import quote_plus
 
+from sqlalchemy.engine import URL
+
 from testgen import settings
 from testgen.common.database.flavor.flavor_service import FlavorService
 
@@ -14,14 +16,28 @@ def get_connection_string_head(self):
         return f"mssql+pyodbc://{self.username}:{quote_plus(self.password)}@"
 
     def get_connection_string_from_fields(self):
-        strConnect = (
-            f"mssql+pyodbc://{self.username}:{quote_plus(self.password)}@{self.host}:{self.port}/{self.dbname}?driver=ODBC+Driver+18+for+SQL+Server"
+        connection_url = URL.create(
+            "mssql+pyodbc",
+            username=self.username,
+            password=quote_plus(self.password or ""),
+            host=self.host,
+            port=int(self.port or 1443),
+            database=self.dbname,
+            query={
+                "driver": "ODBC Driver 18 for SQL Server",
+            },
         )
 
-        if "synapse" in self.host:
-            strConnect += "&autocommit=True"
+        if self.connect_with_identity:
+            connection_url = connection_url._replace(username=None, password=None).update_query_dict({
+                "encrypt": "yes",
+                "authentication": "ActiveDirectoryMsi",
+            })
+
+        if self.sql_flavor_code == "synapse_mssql":
+            connection_url = connection_url.update_query_dict({"autocommit": True})
 
-        return strConnect
+        return connection_url.render_as_string(hide_password=False)
 
     def get_pre_connection_queries(self):
         return [
diff --git a/testgen/common/models/connection.py b/testgen/common/models/connection.py
@@ -62,6 +62,7 @@ class Connection(Entity):
     http_path: str = Column(String)
     warehouse: str = Column(String)
     service_account_key: JSON_TYPE = Column(EncryptedJson)
+    connect_with_identity: bool = Column(Boolean, default=False)
 
     _get_by = "connection_id"
     _default_order_by = (asc(func.lower(connection_name)),)
diff --git a/testgen/template/dbsetup/030_initialize_new_schema_structure.sql b/testgen/template/dbsetup/030_initialize_new_schema_structure.sql
@@ -70,6 +70,7 @@ CREATE TABLE connections (
    url VARCHAR(200) default '',
    connect_by_url BOOLEAN default FALSE,
    connect_by_key BOOLEAN DEFAULT FALSE,
+   connect_with_identity BOOLEAN DEFAULT FALSE,
    private_key BYTEA,
    private_key_passphrase BYTEA,
    http_path              VARCHAR(200),
diff --git a/testgen/template/dbupgrade/0159_incremental_upgrade.sql b/testgen/template/dbupgrade/0159_incremental_upgrade.sql
@@ -0,0 +1,2 @@
+SET SEARCH_PATH TO {SCHEMA_NAME};
+ALTER TABLE connections ADD COLUMN connect_with_identity BOOLEAN DEFAULT FALSE;
diff --git a/testgen/ui/components/frontend/js/components/connection_form.js b/testgen/ui/components/frontend/js/components/connection_form.js
@@ -31,6 +31,7 @@
  * @property {boolean} connect_by_url
  * @property {string?} url
  * @property {boolean} connect_by_key
+ * @property {boolean} connect_with_identity
  * @property {string?} private_key
  * @property {string?} private_key_passphrase
  * @property {string?} http_path
@@ -126,6 +127,7 @@ const ConnectionForm = (props, saveButton) => {
         warehouse: connection?.warehouse ?? '',
         url: connection?.url ?? '',
         service_account_key: connection?.service_account_key ?? '',
+        connect_with_identity: connection?.connect_with_identity ?? false,
         sql_flavor_code: connectionFlavor.rawVal ?? '',
         connection_name: connectionName.rawVal ?? '',
         max_threads: connectionMaxThreads.rawVal ?? 4,
@@ -550,7 +552,197 @@ const RedshiftSpectrumForm = RedshiftForm;
 
 const PostgresqlForm = RedshiftForm;
 
-const AzureMSSQLForm = RedshiftForm;
+const AzureMSSQLForm = (
+    connection,
+    flavor,
+    onChange,
+    originalConnection,
+    dynamicConnectionUrl,
+) => {
+    const isValid = van.state(true);
+    const connectByUrl = van.state(connection.rawVal.connect_by_url ?? false);
+    const connectionHost = van.state(connection.rawVal.project_host ?? '');
+    const connectionPort = van.state(connection.rawVal.project_port || defaultPorts[flavor.flavor]);
+    const connectionDatabase = van.state(connection.rawVal.project_db ?? '');
+    const connectionUsername = van.state(connection.rawVal.project_user ?? '');
+    const connectionPassword = van.state(connection.rawVal?.project_pw_encrypted ?? '');
+    const connectionUrl = van.state(connection.rawVal?.url ?? '');
+    const connectWithIdentity = van.state(connection.rawVal?.connect_with_identity ?? '');
+
+    const validityPerField = {};
+
+    van.derive(() => {
+        onChange({
+            project_host: connectionHost.val,
+            project_port: connectionPort.val,
+            project_db: connectionDatabase.val,
+            project_user: connectionUsername.val,
+            project_pw_encrypted: connectionPassword.val,
+            connect_by_url: connectByUrl.val,
+            url: connectByUrl.val ? connectionUrl.val : connectionUrl.rawVal,
+            connect_by_key: false,
+            connect_with_identity: connectWithIdentity.val,
+        }, isValid.val);
+    });
+
+    van.derive(() => {
+        const newUrlValue = (dynamicConnectionUrl.val ?? '').replace(extractPrefix(dynamicConnectionUrl.rawVal), '');
+        if (!connectByUrl.rawVal) {
+            connectionUrl.val = newUrlValue;
+        }
+    });
+
+    return div(
+        {class: 'flex-column fx-gap-3 fx-flex'},
+        div(
+            { class: 'flex-column border border-radius-1 p-3 mt-1 fx-gap-1', style: 'position: relative;' },
+            Caption({content: 'Server', style: 'position: absolute; top: -10px; background: var(--app-background-color); padding: 0px 8px;' }),
+            RadioGroup({
+                label: 'Connect by',
+                options: [
+                    {
+                        label: 'Host',
+                        value: false,
+                    },
+                    {
+                        label: 'URL',
+                        value: true,
+                    },
+                ],
+                value: connectByUrl,
+                onChange: (value) => connectByUrl.val = value,
+                layout: 'inline',
+            }),
+            div(
+                { class: 'flex-row fx-gap-3 fx-flex' },
+                Input({
+                    name: 'db_host',
+                    label: 'Host',
+                    value: connectionHost,
+                    class: 'fx-flex',
+                    disabled: connectByUrl,
+                    onChange: (value, state) => {
+                        connectionHost.val = value;
+                        validityPerField['db_host'] = state.valid;
+                        isValid.val = Object.values(validityPerField).every(v => v);
+                    },
+                    validators: [
+                        maxLength(250),
+                        requiredIf(() => !connectByUrl.val),
+                    ],
+                }),
+                Input({
+                    name: 'db_port',
+                    label: 'Port',
+                    value: connectionPort,
+                    type: 'number',
+                    disabled: connectByUrl,
+                    onChange: (value, state) => {
+                        connectionPort.val = value;
+                        validityPerField['db_port'] = state.valid;
+                        isValid.val = Object.values(validityPerField).every(v => v);
+                    },
+                    validators: [
+                        minLength(3),
+                        maxLength(5),
+                        requiredIf(() => !connectByUrl.val),
+                    ],
+                })
+            ),
+            Input({
+                name: 'db_name',
+                label: 'Database',
+                value: connectionDatabase,
+                disabled: connectByUrl,
+                onChange: (value, state) => {
+                    connectionDatabase.val = value;
+                    validityPerField['db_name'] = state.valid;
+                    isValid.val = Object.values(validityPerField).every(v => v);
+                },
+                validators: [
+                    maxLength(100),
+                    requiredIf(() => !connectByUrl.val),
+                ],
+            }),
+            () => div(
+                { class: 'flex-row fx-gap-3 fx-align-stretch', style: 'position: relative;' },
+                Input({
+                    label: 'URL',
+                    value: connectionUrl,
+                    class: 'fx-flex',
+                    name: 'url_suffix',
+                    prefix: span({ style: 'white-space: nowrap; color: var(--disabled-text-color)' }, extractPrefix(dynamicConnectionUrl.val)),
+                    disabled: !connectByUrl.val,
+                    onChange: (value, state) => {
+                        connectionUrl.val = value;
+                        validityPerField['url_suffix'] = state.valid;
+                        isValid.val = Object.values(validityPerField).every(v => v);
+                    },
+                    validators: [
+                        requiredIf(() => connectByUrl.val),
+                    ],
+                }),
+            ),
+        ),
+
+        div(
+            { class: 'flex-column border border-radius-1 p-3 mt-1 fx-gap-1', style: 'position: relative;' },
+            Caption({content: 'Authentication', style: 'position: absolute; top: -10px; background: var(--app-background-color); padding: 0px 8px;' }),
+
+            RadioGroup({
+                label: 'Connection Strategy',
+                options: [
+                    {label: 'Connect By Password', value: false},
+                    {label: 'Connect with Managed Identity', value: true},
+                ],
+                value: connectWithIdentity,
+                onChange: (value) => connectWithIdentity.val = value,
+                layout: 'inline',
+            }),
+
+            () => {
+                const _connectWithIdentity = connectWithIdentity.val;
+                if (_connectWithIdentity) {
+                    return div(
+                        {class: 'flex-row p-4 fx-justify-center text-secondary'},
+                        'Configured Microsoft Entra ID credentials will be used',
+                    );
+                }
+
+                return div(
+                    {class: 'flex-column fx-gap-1'},
+                    Input({
+                        name: 'db_user',
+                        label: 'Username',
+                        value: connectionUsername,
+                        onChange: (value, state) => {
+                            connectionUsername.val = value;
+                            validityPerField['db_user'] = state.valid;
+                            isValid.val = Object.values(validityPerField).every(v => v);
+                        },
+                        validators: [
+                            requiredIf(() => !connectWithIdentity.val),
+                            maxLength(50),
+                        ],
+                    }),
+                    Input({
+                        name: 'password',
+                        label: 'Password',
+                        value: connectionPassword,
+                        type: 'password',
+                        passwordSuggestions: false,
+                        placeholder: (originalConnection?.connection_id && originalConnection?.project_pw_encrypted) ? secretsPlaceholder : '',
+                        onChange: (value, state) => {
+                            connectionPassword.val = value;
+                            validityPerField['password'] = state.valid;
+                            isValid.val = Object.values(validityPerField).every(v => v);
+                        },
+                    }),
+                )
+            },
+        ),
+    );
+};
 
 const SynapseMSSQLForm = RedshiftForm;
 
@@ -1110,19 +1302,27 @@ const BigqueryForm = (
 };
 
 function extractPrefix(url) {
-    const parts = (url ?? '').split('@');
-    if (!parts[0]) {
+    if (!url) {
         return '';
     }
-    return `${parts[0]}@`;
+
+    if (url.includes('@')) {
+        const parts = url.split('@');
+        if (!parts[0]) {
+            return '';
+        }
+        return `${parts[0]}@`;
+    }
+
+    return url.slice(0, url.indexOf('://') + 3);
 }
 
 function shouldRefreshUrl(previous, current) {
     if (current.connect_by_url) {
         return false;
     }
 
-    const fields = ['sql_flavor', 'project_host', 'project_port', 'project_db', 'project_user', 'connect_by_key', 'http_path', 'warehouse'];
+    const fields = ['sql_flavor', 'project_host', 'project_port', 'project_db', 'project_user', 'connect_by_key', 'http_path', 'warehouse', 'connect_with_identity'];
     return fields.some((fieldName) => previous[fieldName] !== current[fieldName]);
 }
 
diff --git a/testgen/ui/views/connections.py b/testgen/ui/views/connections.py
@@ -119,6 +119,10 @@ def on_save_connection_clicked(updated_connection):
                 elif updated_connection.get("project_pw_encrypted") == CLEAR_SENTINEL:
                     updated_connection["project_pw_encrypted"] = ""
 
+            if updated_connection.get("connect_with_identity"):
+                updated_connection["project_user"] = ""
+                updated_connection["project_pw_encrypted"] = ""
+
             updated_connection["sql_flavor"] = self._get_sql_flavor_from_value(updated_connection["sql_flavor_code"]).flavor
 
             set_save(True)
@@ -143,6 +147,10 @@ def on_test_connection_clicked(updated_connection: dict) -> None:
             elif updated_connection.get("private_key_passphrase") == CLEAR_SENTINEL:
                 updated_connection["private_key_passphrase"] = ""
 
+            if updated_connection.get("connect_with_identity"):
+                updated_connection["project_user"] = ""
+                updated_connection["project_pw_encrypted"] = ""
+
             updated_connection["sql_flavor"] = self._get_sql_flavor_from_value(updated_connection["sql_flavor_code"]).flavor
 
             set_check_status(True)

Original file line number	Diff line number	Diff line change
`@@ -1,4 +1,4 @@`
`1`		`-ARG TESTGEN_BASE_LABEL=v8`
	`1`	`+ARG TESTGEN_BASE_LABEL=v9`
`2`	`2`
`3`	`3`	`FROM datakitchen/dataops-testgen-base:${TESTGEN_BASE_LABEL} AS release-image`
`4`	`4`
Original file line number	Diff line number	Diff line change
`@@ -0,0 +1,2 @@`
	`1`	`+SET SEARCH_PATH TO {SCHEMA_NAME};`
	`2`	`+ALTER TABLE connections ADD COLUMN connect_with_identity BOOLEAN DEFAULT FALSE;`