Add DataStores models

geoservercloud/geoservercloud.py

@@ -10,7 +10,13 @@
 from requests import Response
 
 from geoservercloud import utils
-from geoservercloud.models import Workspace, Workspaces
+from geoservercloud.models import (
+    DataStores,
+    KeyDollarListDict,
+    PostGisDataStore,
+    Workspace,
+    Workspaces,
+)
 from geoservercloud.services import (
     AclEndpoints,
     GwcEndpoints,
@@ -156,52 +162,72 @@ def unset_default_locale_for_service(self, workspace_name) -> None:
         """
         self.set_default_locale_for_service(workspace_name, None)
 
+    def get_datastores(self, workspace_name: str) -> dict[str, Any]:
+        """
+        Get all datastores for a given workspace
+        """
+        response = self.get_request(self.rest_endpoints.datastores(workspace_name))
+        return DataStores.from_response(response).datastores
+
     def create_pg_datastore(
         self,
         workspace_name: str,
-        datastore: str,
+        datastore_name: str,
         pg_host: str,
         pg_port: int,
         pg_db: str,
         pg_user: str,
         pg_password: str,
         pg_schema: str = "public",
+        description: str | None = None,
         set_default_datastore: bool = False,
     ) -> Response | None:
         """
         Create a PostGIS datastore from the DB connection parameters, or update it if it already exist.
         """
         response: None | Response = None
-        payload: dict[str, dict[str, Any]] = Templates.postgis_data_store(
-            datastore=datastore,
-            pg_host=pg_host,
-            pg_port=pg_port,
-            pg_db=pg_db,
-            pg_user=pg_user,
-            pg_password=pg_password,
-            namespace=f"http://{workspace_name}",
-            pg_schema=pg_schema,
+        datastore = PostGisDataStore(
+            workspace_name,
+            datastore_name,
+            connection_parameters=KeyDollarListDict(
+                input_dict={
+                    "dbtype": "postgis",
+                    "host": pg_host,
+                    "port": pg_port,
+                    "database": pg_db,
+                    "user": pg_user,
+                    "passwd": pg_password,
+                    "schema": pg_schema,
+                    "namespace": f"http://{workspace_name}",
+                    "Expose primary keys": "true",
+                }
+            ),
+            data_store_type="PostGIS",
+            description=description,
         )
+        payload = datastore.put_payload()
+
         if not self.resource_exists(
-            self.rest_endpoints.datastore(workspace_name, datastore)
+            self.rest_endpoints.datastore(workspace_name, datastore_name)
         ):
             response = self.post_request(
                 self.rest_endpoints.datastores(workspace_name), json=payload
             )
         else:
             response = self.put_request(
-                self.rest_endpoints.datastore(workspace_name, datastore), json=payload
+                self.rest_endpoints.datastore(workspace_name, datastore_name),
+                json=payload,
             )
 
         if set_default_datastore:
-            self.default_datastore = datastore
+            self.default_datastore = datastore_name
 
         return response
 
     def create_jndi_datastore(
         self,
         workspace_name: str,
-        datastore: str,
+        datastore_name: str,
         jndi_reference: str,
         pg_schema: str = "public",
         description: str | None = None,
@@ -211,26 +237,36 @@ def create_jndi_datastore(
         Create a PostGIS datastore from JNDI resource, or update it if it already exist.
         """
         response: None | Response = None
-        payload: dict[str, dict[str, Any]] = Templates.postgis_jndi_data_store(
-            datastore=datastore,
-            jndi_reference=jndi_reference,
-            namespace=f"http://{workspace_name}",
-            pg_schema=pg_schema,
+        datastore = PostGisDataStore(
+            workspace_name,
+            datastore_name,
+            connection_parameters=KeyDollarListDict(
+                input_dict={
+                    "dbtype": "postgis",
+                    "jndiReferenceName": jndi_reference,
+                    "schema": pg_schema,
+                    "namespace": f"http://{workspace_name}",
+                    "Expose primary keys": "true",
+                }
+            ),
+            data_store_type="PostGIS (JNDI)",
             description=description,
         )
+        payload = datastore.put_payload()
         if not self.resource_exists(
-            self.rest_endpoints.datastore(workspace_name, datastore)
+            self.rest_endpoints.datastore(workspace_name, datastore_name)
         ):
             response = self.post_request(
                 self.rest_endpoints.datastores(workspace_name), json=payload
             )
         else:
             response = self.put_request(
-                self.rest_endpoints.datastore(workspace_name, datastore), json=payload
+                self.rest_endpoints.datastore(workspace_name, datastore_name),
+                json=payload,
             )
 
         if set_default_datastore:
-            self.default_datastore = datastore
+            self.default_datastore = datastore_name
 
         return response
 

geoservercloud/models/__init__.py

@@ -1,5 +1,13 @@
 from .common import KeyDollarListDict
+from .dataStore import PostGisDataStore
+from .dataStores import DataStores
 from .workspace import Workspace
 from .workspaces import Workspaces
 
-__all__ = ["KeyDollarListDict", "Workspaces", "Workspace"]
+__all__ = [
+    "DataStores",
+    "KeyDollarListDict",
+    "PostGisDataStore",
+    "Workspaces",
+    "Workspace",
+]

geoservercloud/models/common.py

@@ -5,12 +5,14 @@
 
 
 class KeyDollarListDict(dict):
-    def __init__(self, input_list=None, *args, **kwargs):
+    def __init__(self, input_list=None, input_dict=None, *args, **kwargs):
         super().__init__(*args, **kwargs)
         self.key_prefix = "@key"
         self.value_prefix = "$"
         if input_list:
             self.deserialize(input_list)
+        if input_dict:
+            self.update(input_dict)
         log.debug(self)
 
     def deserialize(self, input_list):
@@ -34,6 +36,6 @@ def __repr__(self) -> str:
     def __str__(self):
         return json.dumps(self.serialize())
 
-    # def update(self, other: dict):
-    #     for key, value in other.items():
-    #         super().__setitem__(key, value)
+    def update(self, other: dict):  # type: ignore
+        for key, value in other.items():
+            super().__setitem__(key, value)

geoservercloud/models/dataStore.py

@@ -0,0 +1,70 @@
+import logging
+
+from . import KeyDollarListDict
+
+log = logging.getLogger()
+
+
+class PostGisDataStore:
+
+    def __init__(
+        self,
+        workspace_name: str,
+        data_store_name: str,
+        connection_parameters: KeyDollarListDict,
+        data_store_type: str = "PostGIS",
+        enabled: bool = True,
+        description: str | None = None,
+    ) -> None:
+        self.workspace_name = workspace_name
+        self.data_store_name = data_store_name
+        self.connection_parameters = connection_parameters
+        self.data_store_type = data_store_type
+        self.description = description
+        self.enabled = enabled
+
+    @property
+    def name(self):
+        return self.data_store_name
+
+    def post_payload(self):
+        payload = {
+            "dataStore": {
+                "name": self.data_store_name,
+                "type": self.data_store_type,
+                "connectionParameters": {
+                    "entry": self.connection_parameters.serialize()
+                },
+            }
+        }
+        if self.description:
+            payload["dataStore"]["description"] = self.description
+        if self.enabled:
+            payload["dataStore"]["enabled"] = self.enabled
+        return payload
+
+    def put_payload(self):
+        payload = self.post_payload()
+        return payload
+
+    @classmethod
+    def from_response(cls, response):
+
+        json_data = response.json()
+        connection_parameters = cls.parse_connection_parameters(json_data)
+        return cls(
+            json_data.get("dataStore", {}).get("workspace", {}).get("name", None),
+            json_data.get("dataStore", {}).get("name", None),
+            connection_parameters,
+            json_data.get("dataStore", {}).get("type", "PostGIS"),
+            json_data.get("dataStore", {}).get("enabled", True),
+            json_data.get("dataStore", {}).get("description", None),
+        )
+
+    @classmethod
+    def parse_connection_parameters(cls, json_data):
+        return KeyDollarListDict(
+            json_data.get("dataStore", {})
+            .get("connectionParameters", {})
+            .get("entry", [])
+        )

geoservercloud/models/dataStores.py

@@ -0,0 +1,30 @@
+import jsonschema
+import requests
+
+import logging
+log = logging.getLogger()
+
+
+class DataStores:
+
+    def __init__(self, workspace_name: str, datastores: list[str] = []) -> None:
+        self.workspace_name = workspace_name
+        self._datastores = datastores
+
+    @property
+    def datastores(self):
+        return self._datastores
+
+    @classmethod
+    def from_response(cls, response):
+        json_data = response.json()
+        datastores = []
+        workspace_name = json_data.get('dataStores', {}).get('workspace', {}).get('name', None)
+        for store in json_data.get('dataStores', {}).get('dataStore', []):
+            datastores.append(store['name'])
+            for data_store_name in datastores:
+                log.debug(f"Name: {data_store_name}")
+        return cls(workspace_name, datastores)
+
+    def __repr__(self):
+        return str(self.datastores)

tests/test_datastore.py

@@ -25,6 +25,8 @@ def pg_payload() -> Generator[dict[str, dict[str, Any]], Any, None]:
     yield {
         "dataStore": {
             "name": STORE,
+            "type": "PostGIS",
+            "enabled": True,
             "connectionParameters": {
                 "entry": [
                     {"@key": "dbtype", "$": "postgis"},
@@ -51,6 +53,8 @@ def jndi_payload() -> Generator[dict[str, dict[str, Any]], Any, None]:
         "dataStore": {
             "name": STORE,
             "description": DESCRIPTION,
+            "type": "PostGIS (JNDI)",
+            "enabled": True,
             "connectionParameters": {
                 "entry": [
                     {"@key": "dbtype", "$": "postgis"},
@@ -73,6 +77,34 @@ def jndi_payload() -> Generator[dict[str, dict[str, Any]], Any, None]:
     }
 
 
+@pytest.fixture(scope="module")
+def datastores_response() -> Generator[dict[str, Any], Any, None]:
+    yield {
+        "dataStores": {
+            "dataStore": [
+                {
+                    "name": STORE,
+                    "href": f"{GEOSERVER_URL}/rest/workspaces/{WORKSPACE}/datastores/{STORE}.json",
+                }
+            ],
+        }
+    }
+
+
+def test_get_datastores(
+    geoserver: GeoServerCloud, datastores_response: dict[str, Any]
+) -> None:
+    with responses.RequestsMock() as rsps:
+        rsps.get(
+            url=f"{GEOSERVER_URL}/rest/workspaces/{WORKSPACE}/datastores.json",
+            status=200,
+            json=datastores_response,
+        )
+
+        datastores = geoserver.get_datastores(workspace_name=WORKSPACE)
+        assert datastores == ["test_store"]
+
+
 def test_create_pg_datastore(
     geoserver: GeoServerCloud, pg_payload: dict[str, dict[str, Any]]
 ) -> None:
@@ -90,7 +122,7 @@ def test_create_pg_datastore(
 
         response = geoserver.create_pg_datastore(
             workspace_name=WORKSPACE,
-            datastore=STORE,
+            datastore_name=STORE,
             pg_host=HOST,
             pg_port=PORT,
             pg_db=DATABASE,
@@ -119,7 +151,7 @@ def test_update_pg_datastore(
 
         response = geoserver.create_pg_datastore(
             workspace_name=WORKSPACE,
-            datastore=STORE,
+            datastore_name=STORE,
             pg_host=HOST,
             pg_port=PORT,
             pg_db=DATABASE,
@@ -148,7 +180,7 @@ def test_create_jndi_datastore(
 
         response = geoserver.create_jndi_datastore(
             workspace_name=WORKSPACE,
-            datastore=STORE,
+            datastore_name=STORE,
             jndi_reference=JNDI,
             pg_schema=SCHEMA,
             description=DESCRIPTION,
@@ -174,7 +206,7 @@ def test_update_jndi_datastore(
 
         response = geoserver.create_jndi_datastore(
             workspace_name=WORKSPACE,
-            datastore=STORE,
+            datastore_name=STORE,
             jndi_reference=JNDI,
             pg_schema=SCHEMA,
             description=DESCRIPTION,