Connection String and api-key authentication support (#566)

* Regenerate rest api client

* Workspace connection params refactoring

* New test recordings

* Remove redundant recording processor

* Add pytest configs for better dev experience

* Fix pytest.ini

* Add script to clear env vars for dev convenience

* Linting and small fixes

* Re-record list jobs test

* Further clean-up and linting

* Small renaming

* Add return type annotations

* Update env var for live tests

* Documentation and clean-up

* Add support for connection string

* Remove QUANTUM_API_KEY env var

* Remove QUANTUM_API_KEY env var

* Update unit tests

* Logic to get current primary key

* Finish test for connection string

* Connection string auth test recording

* Remove unnecessary method

* Upgrade minor version

* Fixes the order of property creation

* Add pytest marks

* Fix order of environment variable loading

* Remove unused code

* Better naming for arm and quantum endpoint urls

* Normalize endpoint urls

* Fix cirq test cases

* Fix qiskit test cases

azure-quantum/azure/quantum/_constants.py

@@ -22,6 +22,7 @@ class EnvironmentVariables:
     AZURE_CLIENT_SECRET = SdkEnvironmentVariables.AZURE_CLIENT_SECRET
     AZURE_TENANT_ID = SdkEnvironmentVariables.AZURE_TENANT_ID
     QUANTUM_TOKEN_FILE = "AZURE_QUANTUM_TOKEN_FILE"
+    CONNECTION_STRING = "AZURE_QUANTUM_CONNECTION_STRING"
     ALL = [
         USER_AGENT_APPID,
         QUANTUM_LOCATION,
@@ -36,6 +37,7 @@ class EnvironmentVariables:
         AZURE_CLIENT_SECRET,
         AZURE_TENANT_ID,
         QUANTUM_TOKEN_FILE,
+        CONNECTION_STRING,
     ]
 
 
@@ -73,6 +75,17 @@ class ConnectionConstants:
         f"Workspaces/{workspace_name}"
     )
 
+    VALID_CONNECTION_STRING = (
+        lambda subscription_id, resource_group, workspace_name, api_key, quantum_endpoint:
+        f"SubscriptionId={subscription_id};" +
+        f"ResourceGroupName={resource_group};" +
+        f"WorkspaceName={workspace_name};" +
+        f"ApiKey={api_key};" +
+        f"QuantumEndpoint={quantum_endpoint};"
+    )
+
+    QUANTUM_API_KEY_HEADER = "x-ms-quantum-api-key"
+
 GUID_REGEX_PATTERN = (
     r"[0-9a-f]{8}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{12}"
 )

azure-quantum/azure/quantum/_workspace_connection_params.py

@@ -12,6 +12,8 @@
     Union,
     Any
 )
+from azure.core.credentials import AzureKeyCredential
+from azure.core.pipeline.policies import AzureKeyCredentialPolicy
 from azure.quantum._authentication import _DefaultAzureCredential
 from azure.quantum._constants import (
     EnvironmentKind,
@@ -37,6 +39,17 @@ class WorkspaceConnectionParams:
         """,
         re.VERBOSE | re.IGNORECASE)
 
+    CONNECTION_STRING_REGEX = re.compile(
+        fr"""
+            ^
+            SubscriptionId=(?P<subscription_id>{GUID_REGEX_PATTERN});
+            ResourceGroupName=(?P<resource_group>[^\s;]+);
+            WorkspaceName=(?P<workspace_name>[^\s;]+);
+            ApiKey=(?P<api_key>[^\s;]+);
+            QuantumEndpoint=(?P<quantum_endpoint>https://(?P<location>[^\s\.]+).quantum(?:-test)?.azure.com/);
+        """,
+        re.VERBOSE | re.IGNORECASE)
+
     def __init__(
         self,
         subscription_id: Optional[str] = None,
@@ -53,6 +66,7 @@ def __init__(
         tenant_id: Optional[str] = None,
         client_id: Optional[str] = None,
         api_version: Optional[str] = None,
+        connection_string: Optional[str] = None,
         on_new_client_request: Optional[Callable] = None,
     ):
         # fields are used for these properties since
@@ -75,6 +89,9 @@ def __init__(
         # for example, when changing the user agent
         self.on_new_client_request = on_new_client_request
         # merge the connection parameters passed
+        # connection_string is set first as it
+        # should be overridden by other parameters
+        self.apply_connection_string(connection_string)
         self.merge(
             api_version=api_version,
             arm_endpoint=arm_endpoint,
@@ -162,6 +179,21 @@ def arm_endpoint(self):
     def arm_endpoint(self, value: str):
         self._arm_endpoint = value
 
+    @property
+    def api_key(self):
+        """
+        The api-key stored in a AzureKeyCredential.
+        """
+        return (self.credential.key
+                if isinstance(self.credential, AzureKeyCredential)
+                else None)
+
+    @api_key.setter
+    def api_key(self, value: str):
+        if value:
+            self.credential = AzureKeyCredential(value)
+        self._api_key = value
+
     def __repr__(self):
         """
         Print all fields and properties.
@@ -191,6 +223,19 @@ def apply_resource_id(self, resource_id: str):
                 raise ValueError("Invalid resource id")
             self._merge_re_match(match)
 
+    def apply_connection_string(self, connection_string: str):
+        """
+        Parses the connection_string and set the connection
+        parameters obtained from it.
+        """
+        if connection_string:
+            match = re.search(
+                WorkspaceConnectionParams.CONNECTION_STRING_REGEX,
+                connection_string)
+            if not match:
+                raise ValueError("Invalid connection string")
+            self._merge_re_match(match)
+
     def merge(
         self,
         subscription_id: Optional[str] = None,
@@ -206,6 +251,7 @@ def merge(
         tenant_id: Optional[str] = None,
         client_id: Optional[str] = None,
         api_version: Optional[str] = None,
+        api_key: Optional[str] = None,
     ):
         """
         Set all fields/properties with `not None` values
@@ -226,6 +272,7 @@ def merge(
             user_agent=user_agent,
             user_agent_app_id=user_agent_app_id,
             workspace_name=workspace_name,
+            api_key=api_key,
             merge_default_mode=False,
         )
         return self
@@ -245,6 +292,7 @@ def apply_defaults(
         tenant_id: Optional[str] = None,
         client_id: Optional[str] = None,
         api_version: Optional[str] = None,
+        api_key: Optional[str] = None,
     ) -> WorkspaceConnectionParams:
         """
         Set all fields/properties with `not None` values
@@ -266,6 +314,7 @@ def apply_defaults(
             user_agent=user_agent,
             user_agent_app_id=user_agent_app_id,
             workspace_name=workspace_name,
+            api_key=api_key,
             merge_default_mode=True,
         )
         return self
@@ -286,6 +335,7 @@ def _merge(
         tenant_id: Optional[str] = None,
         client_id: Optional[str] = None,
         api_version: Optional[str] = None,
+        api_key: Optional[str] = None,
     ):
         """
         Set all fields/properties with `not None` values
@@ -313,6 +363,7 @@ def _get_value_or_default(old_value, new_value):
         self.client_id = _get_value_or_default(self.client_id, client_id)
         self.tenant_id = _get_value_or_default(self.tenant_id, tenant_id)
         self.api_version = _get_value_or_default(self.api_version, api_version)
+        self.api_key = _get_value_or_default(self.api_key, api_key)
         # for these properties that have a default value in the getter, we use
         # the private field as the old_value
         self.quantum_endpoint = _get_value_or_default(self._quantum_endpoint, quantum_endpoint)
@@ -360,6 +411,16 @@ def get_credential_or_default(self) -> Any:
                     arm_endpoint=self.arm_endpoint,
                     tenant_id=self.tenant_id))
 
+    def get_auth_policy(self) -> Any:
+        """
+        Returns a AzureKeyCredentialPolicy if using an AzureKeyCredential.
+        Defaults to None.
+        """
+        if isinstance(self.credential, AzureKeyCredential):
+            return AzureKeyCredentialPolicy(self.credential,
+                                            ConnectionConstants.QUANTUM_API_KEY_HEADER)
+        return None
+
     def append_user_agent(self, value: str):
         """
         Append a new value to the Workspace's UserAgent and re-initialize the
@@ -417,6 +478,7 @@ def assert_complete(self):
                     1) A valid combination of location and resource ID.
                     2) A valid combination of location, subscription ID,
                     resource group name, and workspace name.
+                    3) A valid connection string (via Workspace.from_connection_string()).
                 """)
 
     def default_from_env_vars(self) -> WorkspaceConnectionParams:
@@ -445,6 +507,18 @@ def default_from_env_vars(self) -> WorkspaceConnectionParams:
         # because the getter return default values
         self.environment = (self._environment
                             or os.environ.get(EnvironmentVariables.QUANTUM_ENV))
+        # only try to use the connection string from env var if
+        # we really need it
+        if (not self.location
+            or not self.subscription_id
+            or not self.resource_group
+            or not self.workspace_name
+            or not self.credential
+        ):
+            self._merge_connection_params(
+                connection_params=WorkspaceConnectionParams(
+                    connection_string=os.environ.get(EnvironmentVariables.CONNECTION_STRING)),
+                merge_default_mode=True)
         return self
 
     @classmethod
@@ -466,5 +540,6 @@ def get_value(group_name):
             workspace_name=get_value('workspace_name'),
             location=get_value('location'),
             quantum_endpoint=get_value('quantum_endpoint'),
+            api_key=get_value('api_key'),
             arm_endpoint=get_value('arm_endpoint'),
         )

azure-quantum/azure/quantum/workspace.py

@@ -58,6 +58,10 @@ class Workspace:
     2. specify a valid location, subscription ID,
        resource group, and workspace name.
 
+    You can also use a connection string to specify the connection parameters
+    to a Azure Quantum Workspace by calling:
+        Workspace.from_connection_string()
+
     If the Azure Quantum workspace does not have linked storage, the caller
     must also pass a valid Azure storage account connection string.
 
@@ -172,6 +176,7 @@ def _create_client(self) -> QuantumClient:
             user_agent=connection_params.get_full_user_agent(),
             credential_scopes = [ConnectionConstants.DATA_PLANE_CREDENTIAL_SCOPE],
             endpoint=connection_params.quantum_endpoint,
+            authentication_policy=connection_params.get_auth_policy(),
             **kwargs
         )
         return client
@@ -193,6 +198,20 @@ def append_user_agent(self, value: str):
         """
         self._connection_params.append_user_agent(value=value)
 
+    @classmethod
+    def from_connection_string(cls, connection_string: str, **kwargs) -> Workspace:
+        """
+        Creates a new Azure Quantum Workspace client from a connection string.
+        """
+        connection_params = WorkspaceConnectionParams(connection_string=connection_string)
+        return cls(
+            subscription_id=connection_params.subscription_id,
+            resource_group=connection_params.resource_group,
+            name=connection_params.workspace_name,
+            location=connection_params.location,
+            credential=connection_params.get_credential_or_default(),
+            **kwargs)
+
     def _get_top_level_items_client(self) -> TopLevelItemsOperations:
         return self._client.top_level_items
 

azure-quantum/eng/Clear-Env-Vars.ps1

@@ -10,3 +10,4 @@ $env:AZURE_QUANTUM_WORKSPACE_LOCATION = $null
 $env:AZURE_SUBSCRIPTION_ID = $env:SUBSCRIPTION_ID 
 $env:AZURE_RESOURCE_GROUP = $env:AZURE_QUANTUM_WORKSPACE_RG
 $env:QUANTUM_TOKEN_FILE = $null
+$env:AZURE_QUANTUM_CONNECTION_STRING = $null

azure-quantum/tests/unit/common.py

@@ -283,6 +283,9 @@ class CustomRecordingProcessor(RecordingProcessor):
         "user-agent",
         "www-authenticate",
     ]
+    ALLOW_SANITIZED_HEADERS = [
+        ConnectionConstants.QUANTUM_API_KEY_HEADER,
+    ]
 
     def __init__(self, quantumTest: QuantumTestBase):
         self._regexes = []
@@ -405,6 +408,8 @@ def process_request(self, request):
         for key in request.headers:
             if key.lower() in self.ALLOW_HEADERS:
                 headers[key] = self._regex_replace_all(request.headers[key])
+            if key.lower() in self.ALLOW_SANITIZED_HEADERS:
+                headers[key] = PLACEHOLDER
         request.headers = headers
 
         request.uri = self._regex_replace_all(request.uri)

azure-quantum/tests/unit/recordings/test_workspace_auth_connection_string_api_key.yaml

@@ -0,0 +1,66 @@
+interactions:
+- request:
+    body: null
+    headers:
+      Accept:
+      - application/json
+      Accept-Encoding:
+      - gzip, deflate
+      Connection:
+      - keep-alive
+      User-Agent:
+      - azsdk-python-quantum/0.0.1.0a1 Python/3.9.18 (Windows-10-10.0.22631-SP0)
+      x-ms-quantum-api-key:
+      - PLACEHOLDER
+    method: GET
+    uri: https://eastus.quantum.azure.com/subscriptions/00000000-0000-0000-0000-000000000000/resourceGroups/myresourcegroup/providers/Microsoft.Quantum/workspaces/myworkspace/providerStatus?api-version=2022-09-12-preview&test-sequence-id=1
+  response:
+    body:
+      string: '{"value": [{"id": "microsoft-elements", "currentAvailability": "Available",
+        "targets": [{"id": "microsoft.dft", "currentAvailability": "Available", "averageQueueTime":
+        0, "statusPage": null}]}, {"id": "ionq", "currentAvailability": "Degraded",
+        "targets": [{"id": "ionq.qpu", "currentAvailability": "Available", "averageQueueTime":
+        465550, "statusPage": "https://status.ionq.co"}, {"id": "ionq.qpu.aria-1",
+        "currentAvailability": "Available", "averageQueueTime": 763856, "statusPage":
+        "https://status.ionq.co"}, {"id": "ionq.qpu.aria-2", "currentAvailability":
+        "Unavailable", "averageQueueTime": 431102, "statusPage": "https://status.ionq.co"},
+        {"id": "ionq.simulator", "currentAvailability": "Available", "averageQueueTime":
+        5, "statusPage": "https://status.ionq.co"}]}, {"id": "microsoft-qc", "currentAvailability":
+        "Available", "targets": [{"id": "microsoft.estimator", "currentAvailability":
+        "Available", "averageQueueTime": 0, "statusPage": null}]}, {"id": "pasqal",
+        "currentAvailability": "Degraded", "targets": [{"id": "pasqal.sim.emu-tn",
+        "currentAvailability": "Available", "averageQueueTime": 0, "statusPage": "https://pasqal.com"},
+        {"id": "pasqal.qpu.fresnel", "currentAvailability": "Degraded", "averageQueueTime":
+        0, "statusPage": "https://pasqal.com"}]}, {"id": "quantinuum", "currentAvailability":
+        "Degraded", "targets": [{"id": "quantinuum.qpu.h1-1", "currentAvailability":
+        "Unavailable", "averageQueueTime": 0, "statusPage": "https://www.quantinuum.com/hardware/h1"},
+        {"id": "quantinuum.sim.h1-1sc", "currentAvailability": "Available", "averageQueueTime":
+        0, "statusPage": "https://www.quantinuum.com/hardware/h1"}, {"id": "quantinuum.qpu.h1-2",
+        "currentAvailability": "Unavailable", "averageQueueTime": 0, "statusPage":
+        "https://www.quantinuum.com/hardware/h1"}, {"id": "quantinuum.sim.h1-2sc",
+        "currentAvailability": "Unavailable", "averageQueueTime": 0, "statusPage":
+        "https://www.quantinuum.com/hardware/h1"}, {"id": "quantinuum.sim.h1-1e",
+        "currentAvailability": "Available", "averageQueueTime": 16, "statusPage":
+        "https://www.quantinuum.com/hardware/h1"}, {"id": "quantinuum.sim.h1-2e",
+        "currentAvailability": "Unavailable", "averageQueueTime": 0, "statusPage":
+        "https://www.quantinuum.com/hardware/h1"}]}, {"id": "rigetti", "currentAvailability":
+        "Degraded", "targets": [{"id": "rigetti.sim.qvm", "currentAvailability": "Available",
+        "averageQueueTime": 5, "statusPage": "https://rigetti.statuspage.io/"}, {"id":
+        "rigetti.qpu.aspen-m-3", "currentAvailability": "Unavailable", "averageQueueTime":
+        0, "statusPage": null}, {"id": "rigetti.qpu.ankaa-9q-1", "currentAvailability":
+        "Available", "averageQueueTime": 5, "statusPage": "https://rigetti.statuspage.io/"},
+        {"id": "rigetti.qpu.ankaa-2", "currentAvailability": "Available", "averageQueueTime":
+        5, "statusPage": "https://rigetti.statuspage.io/"}]}], "nextLink": null}'
+    headers:
+      connection:
+      - keep-alive
+      content-length:
+      - '2'
+      content-type:
+      - application/json; charset=utf-8
+      transfer-encoding:
+      - chunked
+    status:
+      code: 200
+      message: OK
+version: 1