From 96651a1f46c23a31017ee14f9a5f5874f775b5aa Mon Sep 17 00:00:00 2001
From: Dirk Kulawiak <dirk@semi.technology>
Date: Thu, 9 Mar 2023 11:08:27 +0100
Subject: [PATCH] Add support for API keys

---
 ci/docker-compose-wcs.yml          |  3 +++
 integration/test_authentication.py | 23 +++++++++++++++++++++--
 weaviate/auth.py                   |  9 ++++++++-
 weaviate/connect/connection.py     |  7 +++++--
 4 files changed, 37 insertions(+), 5 deletions(-)

diff --git a/ci/docker-compose-wcs.yml b/ci/docker-compose-wcs.yml
index e53f36781..80e917136 100644
--- a/ci/docker-compose-wcs.yml
+++ b/ci/docker-compose-wcs.yml
@@ -25,4 +25,7 @@ services:
       AUTHORIZATION_ADMINLIST_ENABLED: 'true'
       AUTHORIZATION_ADMINLIST_USERS: 'ms_2d0e007e7136de11d5f29fce7a53dae219a51458@existiert.net'
       AUTHENTICATION_OIDC_SCOPES: 'openid,email'
+      AUTHENTICATION_APIKEY_ENABLED: 'true'
+      AUTHENTICATION_APIKEY_ALLOWED_KEYS: 'my-secret-key'
+      AUTHENTICATION_APIKEY_USERS: 'ms_2d0e007e7136de11d5f29fce7a53dae219a51458@existiert.net'
 ...
diff --git a/integration/test_authentication.py b/integration/test_authentication.py
index a04997ed2..f53a3a471 100644
--- a/integration/test_authentication.py
+++ b/integration/test_authentication.py
@@ -1,13 +1,13 @@
 import os
+import time
 import warnings
 from typing import Optional, Dict
 
 import pytest
 import requests
-import time
-from weaviate.exceptions import WeaviateStartUpError
 from requests.exceptions import ConnectionError as RequestsConnectionError
 from requests.exceptions import HTTPError as RequestsHTTPError
+
 import weaviate
 from weaviate import (
     AuthenticationFailedException,
@@ -15,6 +15,8 @@
     AuthClientPassword,
     AuthBearerToken,
 )
+from weaviate.auth import AuthApiKey
+from weaviate.exceptions import WeaviateStartUpError, UnexpectedStatusCodeException
 
 ANON_PORT = "8080"
 AZURE_PORT = "8081"
@@ -253,3 +255,20 @@ def test_bearer_token_without_refresh(recwarn):
     w = recwarn.pop()
     assert issubclass(w.category, UserWarning)
     assert str(w.message).startswith("Auth002")
+
+
+def test_api_key():
+    url = "http://127.0.0.1:" + WCS_PORT
+    assert is_auth_enabled(url)
+
+    client = weaviate.Client(url, auth_client_secret=AuthApiKey(api_key="my-secret-key"))
+    client.schema.delete_all()  # no exception, client works
+
+
+def test_api_key_wrong_key():
+    url = "http://127.0.0.1:" + WCS_PORT
+    assert is_auth_enabled(url)
+
+    with pytest.raises(UnexpectedStatusCodeException) as e:
+        weaviate.Client(url, auth_client_secret=AuthApiKey(api_key="wrong_key"))
+        assert e.value.status_code == 401
diff --git a/weaviate/auth.py b/weaviate/auth.py
index 62d007cc3..8d970c9dd 100644
--- a/weaviate/auth.py
+++ b/weaviate/auth.py
@@ -68,4 +68,11 @@ def __post_init__(self):
             _Warnings.auth_negative_expiration_time(self.expires_in)
 
 
-AuthCredentials = Union[AuthBearerToken, AuthClientPassword, AuthClientCredentials]
+@dataclass
+class AuthApiKey:
+    """Using the given API key to authenticate with weaviate."""
+
+    api_key: str
+
+
+AuthCredentials = Union[AuthBearerToken, AuthClientPassword, AuthClientCredentials, AuthApiKey]
diff --git a/weaviate/connect/connection.py b/weaviate/connect/connection.py
index 3a81d0271..e587f65b1 100644
--- a/weaviate/connect/connection.py
+++ b/weaviate/connect/connection.py
@@ -16,7 +16,7 @@
 from requests.exceptions import HTTPError as RequestsHTTPError
 from requests.exceptions import JSONDecodeError
 
-from weaviate.auth import AuthCredentials, AuthClientCredentials
+from weaviate.auth import AuthCredentials, AuthClientCredentials, AuthApiKey
 from weaviate.connect.authentication import _Auth
 from weaviate.exceptions import (
     AuthenticationFailedException,
@@ -141,7 +141,7 @@ def _create_session(self, auth_client_secret: Optional[AuthCredentials]) -> None
                 self._session = requests.Session()
                 return
 
-            if auth_client_secret is not None:
+            if auth_client_secret is not None and not isinstance(auth_client_secret, AuthApiKey):
                 _auth = _Auth(resp, auth_client_secret, self)
                 self._session = _auth.get_auth_session()
 
@@ -150,6 +150,9 @@ def _create_session(self, auth_client_secret: Optional[AuthCredentials]) -> None
                     self._create_background_token_refresh(_auth)
                 else:
                     self._create_background_token_refresh()
+            elif auth_client_secret is not None and isinstance(auth_client_secret, AuthApiKey):
+                self._headers["authorization"] = "Bearer " + auth_client_secret.api_key
+                self._session = requests.Session()
             else:
                 msg = f""""No login credentials provided. The weaviate instance at {self.url} requires login credentials.