Merge pull request #2 from rog555/azure-cosmosdb-client

use a singleton Azure Cosmos DB client for the lifetime of the application

CHANGELOG.md

@@ -1,6 +1,9 @@
 Changelog
 =========
 
+## [v0.0.21] - 2024-02-14
+- use a singleton Azure Cosmos DB client for the lifetime of the application
+
 ## [v0.0.20] - 2023-10-09
 - abnosql_check_exists=False on cosmos delete bug
 

abnosql/plugins/table/cosmos.py

@@ -35,6 +35,42 @@
     MISSING_DEPS = True
 
 
+def _get_database_client():
+    cf = {}
+    pc = parse_connstr()
+    if pc is not None:
+        cf.update({
+            'account': pc.username,
+            'database': pc.hostname,
+            'credential': (
+                None if (
+                    pc.password == 'DefaultAzureCredential'
+                    or pc.password == ''
+                )
+                else pc.password
+            )
+        })
+    for attr in ['account', 'credential', 'endpoint', 'database']:
+        cf[attr] = os.environ.get(
+            'ABNOSQL_COSMOS_' + attr.upper(),
+            cf.get(attr)
+        )
+    if cf['endpoint'] is None and cf['account'] is not None:
+        cf['endpoint'] = 'https://%s.documents.azure.com' % cf['account']
+    if cf['endpoint'] is None or cf['database'] is None:
+        return None
+    if cf['credential'] is None:
+        cf['credential'] = DefaultAzureCredential()
+    return CosmosClient(
+        url=cf['endpoint'], credential=cf['credential']
+    ).get_database_client(cf['database'])
+
+
+# Azure recommends using a singleton client for the lifetime of your app
+# https://learn.microsoft.com/en-us/azure/azure-functions/manage-connections
+DATABASE_CLIENT = _get_database_client()
+
+
 def cosmos_ex_handler(raise_not_found: t.Optional[bool] = True):
 
     def get_message(e):
@@ -89,7 +125,9 @@ def __init__(
         self.pm = pm
         self.name = name
         self.set_config(config)
-        self.database_client = None
+        self.database_client = DATABASE_CLIENT
+        if os.environ.get('ABNOSQL_DISABLE_GLOBAL_CACHE', 'FALSE') == 'TRUE':
+            self.database_client = None
         self.key_attrs = get_key_attrs(self.config)
         self.check_exists = check_exists_enabled(self.config)
         # enabled by default
@@ -111,35 +149,44 @@ def _database_client(self):
         _client = self.config.get('database_client', self.database_client)
         if _client is not None:
             return _client
-        pc = parse_connstr()
         cf = {}
-        if pc is not None:
-            cf.update({
-                'account': pc.username,
-                'database': pc.hostname,
-                'credential': (
-                    None if (
-                        pc.password == 'DefaultAzureCredential'
-                        or pc.password == ''
+        required = ['endpoint', 'database']
+        additional = ['account', 'credential']
+        local = [
+            _ for _ in required + additional if self.config.get(_) is not None
+        ]
+        disable_cache = os.environ.get('ABNOSQL_DISABLE_GLOBAL_CACHE', 'FALSE')
+        if len(local) or disable_cache == 'TRUE':
+            # prefer local config over env var
+            pc = parse_connstr()
+            if pc is not None:
+                cf.update({
+                    'account': pc.username,
+                    'database': pc.hostname,
+                    'credential': (
+                        None if (
+                            pc.password == 'DefaultAzureCredential'
+                            or pc.password == ''
+                        )
+                        else pc.password
                     )
-                    else pc.password
+                })
+            for attr in required + additional:
+                val = self.config.get(
+                    attr, os.environ.get('ABNOSQL_COSMOS_' + attr.upper())
                 )
-            })
-        required = ['endpoint', 'database']
-        for attr in ['account', 'credential'] + required:
-            val = self.config.get(
-                attr, os.environ.get('ABNOSQL_COSMOS_' + attr.upper())
-            )
-            # override
-            if val is not None:
-                cf[attr] = val
-        if cf.get('endpoint') is None and cf['account'] is not None:
+                # override
+                if val is not None:
+                    cf[attr] = val
+        if cf.get('endpoint') is None and cf.get('account') is not None:
             cf['endpoint'] = 'https://%s.documents.azure.com' % cf['account']
         # use managed identity if no credential supplied
         if cf.get('credential') is None:
             cf['credential'] = DefaultAzureCredential()
-        missing = [_ for _ in required if cf[_] is None]
+        missing = [_ for _ in required if cf.get(_) is None]
         if len(missing):
+            if self.database_client is not None:
+                return self.database_client
             raise ex.ConfigException('missing config: ' + ', '.join(missing))
         self.database_client = CosmosClient(
             url=cf['endpoint'], credential=cf['credential']

abnosql/version.py

@@ -1,4 +1,4 @@
-__version__ = '0.0.20'
+__version__ = '0.0.21'
 
 
 if __name__ == '__main__':

setup.py

@@ -58,8 +58,8 @@ def get_version():
 )
 test_deps = all_deps + [
     'coverage',
-    'moto[dynamodb]',
-    'moto[kms]',
+    'moto[dynamodb]==4.2.13',
+    'moto[kms]==4.2.13',
     'mypy',
     'pytest',
     'pytest-cov',

tests/test_azure_kms.py

@@ -31,6 +31,7 @@ def setup_cosmos():
     os.environ['AZURE_TENANT_ID'] = 'ddb6e3bf-68ca-4216-a84f-d3a4b05c8314'
     os.environ['AZURE_CLIENT_SECRET'] = 'foobar12345'
     os.environ['ABNOSQL_KMS_KEYS'] = KEY_ID
+    os.environ['ABNOSQL_DISABLE_GLOBAL_CACHE'] = 'TRUE'
     return {
         'kms': {
             'key_attrs': ['hk', 'rk'],

tests/test_cosmos.py

@@ -9,6 +9,7 @@
 import abnosql.exceptions as ex
 from abnosql.mocks import mock_cosmos
 from abnosql.mocks.mock_cosmos import set_keyattrs
+from abnosql.plugins.table import cosmos
 from abnosql.plugins.table.memory import clear_tables
 from tests import common as cmn
 
@@ -25,6 +26,29 @@ def setup_cosmos():
         'mycredential'.encode('utf-8')
     ).decode()
     os.environ['ABNOSQL_COSMOS_DATABASE'] = 'bar'
+    os.environ['ABNOSQL_DISABLE_GLOBAL_CACHE'] = 'TRUE'
+
+
+@mock_cosmos
+@responses.activate
+def test_config_exception():
+    setup_cosmos()
+    os.environ['ABNOSQL_DISABLE_GLOBAL_CACHE'] = 'FALSE'
+    with pytest.raises(ex.ConfigException) as e:
+        cmn.test_get_item()
+    os.environ['ABNOSQL_DISABLE_GLOBAL_CACHE'] = 'TRUE'
+    assert str(e.value) == 'missing config: endpoint, database'
+
+
+@mock_cosmos
+@responses.activate
+def test_config_exception1():
+    setup_cosmos()
+    os.environ['ABNOSQL_DISABLE_GLOBAL_CACHE'] = 'FALSE'
+    from importlib import reload
+    reload(cosmos)
+    cmn.test_get_item()
+    os.environ['ABNOSQL_DISABLE_GLOBAL_CACHE'] = 'TRUE'
 
 
 @mock_cosmos