Repository: dscc-admin-ch/lomas (language: python)

Issue:
fixing python versions



Discussion / hints:

Fix patch:
diff --git a/.github/workflows/check_pr_to_master.yml b/.github/workflows/check_pr_to_master.yml
index f9c72ce7..54003715 100644
--- a/.github/workflows/check_pr_to_master.yml
+++ b/.github/workflows/check_pr_to_master.yml
@@ -13,10 +13,10 @@ jobs:
     steps:
     - name: Checkout
       uses: actions/checkout@v3
-    - name: Set up Python 11
+    - name: Set up Python 12
       uses: actions/setup-python@v5
       with:
-        python-version: "3.11"
+        python-version: "3.12"
     - name: Get current branch name
       run: |
         echo ${{ github.head_ref }}
diff --git a/.github/workflows/check_security_codeQL.yml b/.github/workflows/check_security_codeQL.yml
index 396fb0de..b8611d95 100644
--- a/.github/workflows/check_security_codeQL.yml
+++ b/.github/workflows/check_security_codeQL.yml
@@ -1,4 +1,4 @@
-name: "CodeQL Extended"
+name: CodeQL Extended
 
 on:
   pull_request:
@@ -19,27 +19,25 @@ jobs:
     strategy:
       fail-fast: false
       matrix:
-        include:
-        - language: python
-          build-mode: none
+        python-version: ["3.11", "3.12"]
 
     steps:
     - name: Checkout repository
       uses: actions/checkout@v4
 
-    # Initializes the CodeQL tools for scanning.
+    # Set up Python for each version specified in the matrix
+    - name: Set up Python ${{ matrix.python-version }}
+      uses: actions/setup-python@v4
+      with:
+        python-version: ${{ matrix.python-version }}
+
+    # Initialize CodeQL with the current Python version
     - name: Initialize CodeQL
       uses: github/codeql-action/init@v3
       with:
         languages: python
         build-mode: none
         queries: security-extended
-        # If you wish to specify custom queries, you can do so here or in a config file.
-        # By default, queries listed here will override any specified in a config file.
-        # Prefix the list here with "+" to use these queries and those in the config file.
-
-        # For more details on CodeQL's query packs, refer to: https://docs.github.com/en/code-security/code-scanning/automatically-scanning-your-code-for-vulnerabilities-and-errors/configuring-code-scanning#using-queries-in-ql-packs
-        # queries: security-extended,security-and-quality
 
     - name: Perform CodeQL Analysis
       uses: github/codeql-action/analyze@v3
diff --git a/.github/workflows/push_pypi.yml b/.github/workflows/push_pypi.yml
index 7dc5f6f9..6b174c95 100644
--- a/.github/workflows/push_pypi.yml
+++ b/.github/workflows/push_pypi.yml
@@ -8,10 +8,10 @@ jobs:
     steps:
       - name: Checkout
         uses: actions/checkout@v3
-      - name: Set up Python 11
+      - name: Set up Python 12
         uses: actions/setup-python@v5
         with:
-          python-version: "3.11" 
+          python-version: "3.12" 
       - name: Install dependencies
         run: |
           python -m pip install --upgrade pip
diff --git a/.github/workflows/push_readthedocs.yml b/.github/workflows/push_readthedocs.yml
index 5e87bc39..0d0c5f46 100644
--- a/.github/workflows/push_readthedocs.yml
+++ b/.github/workflows/push_readthedocs.yml
@@ -17,10 +17,10 @@ jobs:
     steps:
     - name: Checkout
       uses: actions/checkout@v3
-    - name: Set up Python 11
+    - name: Set up Python 12
       uses: actions/setup-python@v5
       with:
-        python-version: "3.11"
+        python-version: "3.12"
     - name: Install dependencies
       run: |
         cd client
diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md
index c7f19ad8..aeac2a4c 100644
--- a/CONTRIBUTING.md
+++ b/CONTRIBUTING.md
@@ -117,13 +117,6 @@ Some examples are avalaible here:
 
 Do not forget to add these libraries in the `requirements.txt` files.
 
-## Adding a Dataset Store
-Here is the explanation of how to add a new dataset store named `NewDatasetStore` for the example.
-
-1. Add the new dataset store the the `DatasetStoreType` StrEnum class in `lomas/lomas_server/constants.py` and add the `NewDatasetStore` option in the `dataset_store_factory` function (in `lomas/lomas_server/dataset_store/factory.py`).
-2. Create a file for your dataset store in the folder `lomas/lomas_server/dataset_store/new_dataset_store.py`. Inside, create a class `NewDatasetStore` that inherits from `DatasetStore` (`lomas/lomas_server/dataset_store/dataset_store.py`), your class must contain a `_add_dataset` method that handle adding a dataset in memory and a `get_querier` method that the querier for the given dataset and library.
-3. Add tests in `lomas/lomas_server/tests/` to test all functionnalities of the new dataset store.
-
 ## Adding a Data Connector (for private dataset in various databases)
 Here is the explanation of how to add a new data connector named `NewDataConnector` for the example.
 
diff --git a/README.md b/README.md
index 3a9ac81b..eb797b94 100644
--- a/README.md
+++ b/README.md
@@ -7,7 +7,7 @@
 </h1><br>
 
 [![License: MIT](https://img.shields.io/badge/License-MIT-yellow.svg)](https://opensource.org/licenses/MIT)
-[![Python](https://img.shields.io/badge/python-3.11-blue)](https://www.python.org/)
+[![Python versions](https://img.shields.io/pypi/pyversions/lomas_client.svg)](https://pypi.org/project/lomas_client/)
 [![Documentation](https://img.shields.io/badge/docs-Read%20the%20Docs-blue)](https://dscc-admin-ch.github.io/lomas-docs/index.html)
 [![Tests](https://github.com/dscc-admin-ch/lomas/actions/workflows/test_and_coverage_server.yml/badge.svg)](https://github.com/dscc-admin-ch/lomas/actions/workflows/test_and_coverage_server.yml)
 [![Coverage badge](https://raw.githubusercontent.com/dscc-admin-ch/lomas/python-coverage-comment-action-data/badge.svg)](https://htmlpreview.github.io/?https://raw.githubusercontent.com/dscc-admin-ch/lomas/python-coverage-comment-action-data/htmlcov/index.html)
diff --git a/client/Dockerfile b/client/Dockerfile
index 316183e0..e2525be7 100644
--- a/client/Dockerfile
+++ b/client/Dockerfile
@@ -1,4 +1,4 @@
-FROM python:3.11 AS lomas_client_base
+FROM python:3.12 AS lomas_client_base
 
 WORKDIR /code
 
diff --git a/client/README.md b/client/README.md
index 63d114a4..8a2d6dd7 100644
--- a/client/README.md
+++ b/client/README.md
@@ -1,11 +1,12 @@
 [![PyPi version](https://img.shields.io/pypi/v/lomas_client.svg)](https://pypi.org/project/lomas_client/)
 [![PyPi status](https://img.shields.io/pypi/status/lomas_client.svg)](https://pypi.org/project/lomas_client/)
+[![Python versions](https://img.shields.io/pypi/pyversions/lomas_client.svg)](https://pypi.org/project/lomas_client/)
 
 <h1 align="center">
 <picture>
   <source media="(prefers-color-scheme: dark)" srcset="https://github.com/dscc-admin-ch/lomas/blob/wip_322_darkmode-logo/images/lomas_logo_darkmode_txt.png"  width="300">
   <source media="(prefers-color-scheme: light)" srcset="https://github.com/dscc-admin-ch/lomas/blob/wip_322_darkmode-logo/images/lomas_logo_txt.png"  width="300">
-  <img alt="I don't know why but this is needed." src="https://user-images.githubusercontent.com/25423296/163456779-a8556205-d0a5-45e2-ac17-42d089e3c3f8.png">
+  <img alt="This is what is displayed on Pypi" src="https://github.com/dscc-admin-ch/lomas/blob/wip_322_darkmode-logo/images/lomas_logo_txt.png">
 </picture>
 </h1><br>
 
diff --git a/client/deploy/helm/charts/lomas_client/Chart.yaml b/client/deploy/helm/charts/lomas_client/Chart.yaml
index c3a9eebf..fd78019c 100644
--- a/client/deploy/helm/charts/lomas_client/Chart.yaml
+++ b/client/deploy/helm/charts/lomas_client/Chart.yaml
@@ -4,6 +4,6 @@ description: Lomas's Secure Data Disclosure deployment chart for the client envi
 
 type: application
 
-version: 0.3.2
+version: 0.3.3
 
-appVersion: "0.3.2"
+appVersion: "0.3.3"
diff --git a/client/notebooks/Demo_Client_Notebook.ipynb b/client/notebooks/Demo_Client_Notebook.ipynb
index 7dd7394f..f4c8e3f7 100644
--- a/client/notebooks/Demo_Client_Notebook.ipynb
+++ b/client/notebooks/Demo_Client_Notebook.ipynb
@@ -685,7 +685,7 @@
      "name": "stdout",
      "output_type": "stream",
      "text": [
-      "Average bill length in remote dummy: 46.54mm.\n"
+      "Average bill length in remote dummy: 46.59mm.\n"
      ]
     }
    ],
@@ -758,7 +758,7 @@
      "name": "stdout",
      "output_type": "stream",
      "text": [
-      "Average bill length of penguins in real data: 43.83mm.\n"
+      "Average bill length of penguins in real data: 45.02mm.\n"
      ]
     }
    ],
@@ -1030,7 +1030,7 @@
      "name": "stdout",
      "output_type": "stream",
      "text": [
-      "Dummy result for variance: 35.5\n"
+      "Dummy result for variance: 30.28\n"
      ]
     }
    ],
@@ -1120,7 +1120,7 @@
      "name": "stdout",
      "output_type": "stream",
      "text": [
-      "Variance of bill length: 34.77 (from opendp query).\n"
+      "Variance of bill length: 34.29 (from opendp query).\n"
      ]
     }
    ],
@@ -1134,7 +1134,7 @@
    "id": "ddb058f9-42fa-4896-a38f-3bf67fc0b2fb",
    "metadata": {},
    "source": [
-    "#### Postprocessing: no additional provacy risk with DP"
+    "#### Postprocessing: no additional privacy risk with DP"
    ]
   },
   {
@@ -1175,7 +1175,7 @@
      "name": "stdout",
      "output_type": "stream",
      "text": [
-      "The 95% confidence interval of the bill length of all penguins is [43.21, 44.45].\n"
+      "The 95% confidence interval of the bill length of all penguins is [44.4, 45.64].\n"
      ]
     }
    ],
@@ -2339,7 +2339,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 50,
+   "execution_count": 48,
    "id": "ed292cec-8497-4e5b-b3bf-cca9227abf7d",
    "metadata": {},
    "outputs": [],
@@ -2355,7 +2355,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 51,
+   "execution_count": 49,
    "id": "6fab49d9-f5c5-437e-b70f-234b84c2f7e5",
    "metadata": {},
    "outputs": [
@@ -2389,32 +2389,32 @@
        "    <tr>\n",
        "      <th>0</th>\n",
        "      <td>Dream</td>\n",
-       "      <td>48.511585</td>\n",
-       "      <td>14.830853</td>\n",
+       "      <td>41.394843</td>\n",
+       "      <td>15.025560</td>\n",
        "    </tr>\n",
        "    <tr>\n",
        "      <th>1</th>\n",
        "      <td>Dream</td>\n",
-       "      <td>46.058849</td>\n",
-       "      <td>21.126229</td>\n",
+       "      <td>40.640350</td>\n",
+       "      <td>18.723828</td>\n",
        "    </tr>\n",
        "    <tr>\n",
        "      <th>2</th>\n",
        "      <td>Dream</td>\n",
-       "      <td>43.442156</td>\n",
-       "      <td>19.632390</td>\n",
+       "      <td>52.355482</td>\n",
+       "      <td>14.951575</td>\n",
        "    </tr>\n",
        "    <tr>\n",
        "      <th>3</th>\n",
        "      <td>Dream</td>\n",
-       "      <td>37.877035</td>\n",
-       "      <td>21.214375</td>\n",
+       "      <td>44.657039</td>\n",
+       "      <td>18.139905</td>\n",
        "    </tr>\n",
        "    <tr>\n",
        "      <th>4</th>\n",
        "      <td>Dream</td>\n",
-       "      <td>52.552031</td>\n",
-       "      <td>19.885068</td>\n",
+       "      <td>43.308159</td>\n",
+       "      <td>15.976547</td>\n",
        "    </tr>\n",
        "  </tbody>\n",
        "</table>\n",
@@ -2422,14 +2422,14 @@
       ],
       "text/plain": [
        "  island  bill_length_mm  bill_depth_mm\n",
-       "0  Dream       48.511585      14.830853\n",
-       "1  Dream       46.058849      21.126229\n",
-       "2  Dream       43.442156      19.632390\n",
-       "3  Dream       37.877035      21.214375\n",
-       "4  Dream       52.552031      19.885068"
+       "0  Dream       41.394843      15.025560\n",
+       "1  Dream       40.640350      18.723828\n",
+       "2  Dream       52.355482      14.951575\n",
+       "3  Dream       44.657039      18.139905\n",
+       "4  Dream       43.308159      15.976547"
       ]
      },
-     "execution_count": 51,
+     "execution_count": 49,
      "metadata": {},
      "output_type": "execute_result"
     }
@@ -2448,7 +2448,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 52,
+   "execution_count": 50,
    "id": "82ab4aee-b4af-4fbd-93cc-b8171f7f9a52",
    "metadata": {},
    "outputs": [
@@ -2458,7 +2458,7 @@
        "{'epsilon_cost': 1.0, 'delta_cost': 0.00015673368198174188}"
       ]
      },
-     "execution_count": 52,
+     "execution_count": 50,
      "metadata": {},
      "output_type": "execute_result"
     }
@@ -2482,7 +2482,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 53,
+   "execution_count": 51,
    "id": "fbc8b354-e4db-4472-957b-468e768eddc4",
    "metadata": {},
    "outputs": [],
@@ -2498,7 +2498,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 54,
+   "execution_count": 52,
    "id": "550fa89d-9537-4daf-9f96-42fa71f242b9",
    "metadata": {},
    "outputs": [
@@ -2532,32 +2532,32 @@
        "    <tr>\n",
        "      <th>0</th>\n",
        "      <td>Dream</td>\n",
-       "      <td>52.040945</td>\n",
-       "      <td>15.677968</td>\n",
+       "      <td>52.040896</td>\n",
+       "      <td>15.677953</td>\n",
        "    </tr>\n",
        "    <tr>\n",
        "      <th>1</th>\n",
        "      <td>Dream</td>\n",
-       "      <td>47.597868</td>\n",
-       "      <td>17.673375</td>\n",
+       "      <td>47.597791</td>\n",
+       "      <td>17.673353</td>\n",
        "    </tr>\n",
        "    <tr>\n",
        "      <th>2</th>\n",
        "      <td>Torgersen</td>\n",
-       "      <td>57.004345</td>\n",
-       "      <td>15.573335</td>\n",
+       "      <td>57.004282</td>\n",
+       "      <td>15.573296</td>\n",
        "    </tr>\n",
        "    <tr>\n",
        "      <th>3</th>\n",
        "      <td>Biscoe</td>\n",
-       "      <td>37.559502</td>\n",
-       "      <td>19.513633</td>\n",
+       "      <td>37.559513</td>\n",
+       "      <td>19.513630</td>\n",
        "    </tr>\n",
        "    <tr>\n",
        "      <th>4</th>\n",
        "      <td>Dream</td>\n",
-       "      <td>46.813858</td>\n",
-       "      <td>15.565315</td>\n",
+       "      <td>46.813786</td>\n",
+       "      <td>15.565273</td>\n",
        "    </tr>\n",
        "  </tbody>\n",
        "</table>\n",
@@ -2565,14 +2565,14 @@
       ],
       "text/plain": [
        "      island  bill_length_mm  bill_depth_mm\n",
-       "0      Dream       52.040945      15.677968\n",
-       "1      Dream       47.597868      17.673375\n",
-       "2  Torgersen       57.004345      15.573335\n",
-       "3     Biscoe       37.559502      19.513633\n",
-       "4      Dream       46.813858      15.565315"
+       "0      Dream       52.040896      15.677953\n",
+       "1      Dream       47.597791      17.673353\n",
+       "2  Torgersen       57.004282      15.573296\n",
+       "3     Biscoe       37.559513      19.513630\n",
+       "4      Dream       46.813786      15.565273"
       ]
      },
-     "execution_count": 54,
+     "execution_count": 52,
      "metadata": {},
      "output_type": "execute_result"
     }
@@ -2591,7 +2591,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 55,
+   "execution_count": 53,
    "id": "7d5a336e-80f0-48fa-84a3-33d0e51a2d3b",
    "metadata": {},
    "outputs": [],
@@ -2602,7 +2602,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 56,
+   "execution_count": 54,
    "id": "e890a8d9-0c7b-4805-be8a-81e66d5fa7ca",
    "metadata": {},
    "outputs": [
@@ -2610,7 +2610,7 @@
      "name": "stdout",
      "output_type": "stream",
      "text": [
-      "The average with Smartnoise-SQL on private data was 43.83.\n",
+      "The average with Smartnoise-SQL on private data was 45.02.\n",
       "The average with Smartnoise-Synth on synthetic data is 44.65.\n"
      ]
     }
@@ -2624,7 +2624,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 57,
+   "execution_count": 55,
    "id": "11a14d9f-0fe3-4a5e-b425-d2192acd1e84",
    "metadata": {},
    "outputs": [
@@ -2632,7 +2632,7 @@
      "name": "stdout",
      "output_type": "stream",
      "text": [
-      "The variance with opendp on private data was 34.77.\n",
+      "The variance with opendp on private data was 34.29.\n",
       "The variance with Smartnoise-Synth on synthetic data is 36.15.\n"
      ]
     }
@@ -2662,7 +2662,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 58,
+   "execution_count": 56,
    "id": "008fd230-cdfd-4e03-91ce-5a60b06c106d",
    "metadata": {},
    "outputs": [
@@ -2672,7 +2672,7 @@
        "4"
       ]
      },
-     "execution_count": 58,
+     "execution_count": 56,
      "metadata": {},
      "output_type": "execute_result"
     }
@@ -2684,7 +2684,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 59,
+   "execution_count": 57,
    "id": "b712b269-64f2-4c7e-b8bf-d1a608933eff",
    "metadata": {},
    "outputs": [
@@ -2703,15 +2703,15 @@
        " 'response': {'requested_by': 'Dr. Antartica',\n",
        "  'query_response': {'index': [0],\n",
        "   'columns': ['avg_bill_length_mm'],\n",
-       "   'data': [[43.825255002419716]],\n",
+       "   'data': [[45.01787281553949]],\n",
        "   'index_names': [None],\n",
        "   'column_names': [None]},\n",
        "  'spent_epsilon': 1.0,\n",
        "  'spent_delta': 4.999999999999449e-05},\n",
-       " 'timestamp': 1725364882.6916816}"
+       " 'timestamp': 1726136962.7276385}"
       ]
      },
-     "execution_count": 59,
+     "execution_count": 57,
      "metadata": {},
      "output_type": "execute_result"
     }
@@ -2724,7 +2724,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 60,
+   "execution_count": 58,
    "id": "8dfaf2b6-2b6c-480b-bcd7-250b0b2806a6",
    "metadata": {},
    "outputs": [
@@ -2741,13 +2741,13 @@
        "      output_measure = MaxDivergence(f64)),\n",
        "  'fixed_delta': None},\n",
        " 'response': {'requested_by': 'Dr. Antartica',\n",
-       "  'query_response': 34.769447201112094,\n",
+       "  'query_response': 34.28606138801999,\n",
        "  'spent_epsilon': 0.7122093023265228,\n",
        "  'spent_delta': 0},\n",
-       " 'timestamp': 1725364890.340155}"
+       " 'timestamp': 1726136975.3640535}"
       ]
      },
-     "execution_count": 60,
+     "execution_count": 58,
      "metadata": {},
      "output_type": "execute_result"
     }
@@ -2760,7 +2760,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 61,
+   "execution_count": 59,
    "id": "376315ec-6f38-4919-959e-d6bf244a4952",
    "metadata": {},
    "outputs": [
@@ -2791,10 +2791,10 @@
        "                                           epsilon=2.0))])},\n",
        "  'spent_epsilon': 2.0,\n",
        "  'spent_delta': 0.0},\n",
-       " 'timestamp': 1725364897.0776813}"
+       " 'timestamp': 1726136985.1675484}"
       ]
      },
-     "execution_count": 61,
+     "execution_count": 59,
      "metadata": {},
      "output_type": "execute_result"
     }
@@ -2807,7 +2807,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 62,
+   "execution_count": 60,
    "id": "638817e0-d88f-407a-8136-210309651cf2",
    "metadata": {},
    "outputs": [
@@ -2830,25 +2830,25 @@
        "  'nb_samples': 200},\n",
        " 'response': {'requested_by': 'Dr. Antartica',\n",
        "  'query_response':         island  bill_length_mm  bill_depth_mm\n",
-       "  0        Dream       52.040945      15.677968\n",
-       "  1        Dream       47.597868      17.673375\n",
-       "  2    Torgersen       57.004345      15.573335\n",
-       "  3       Biscoe       37.559502      19.513633\n",
-       "  4        Dream       46.813858      15.565315\n",
+       "  0        Dream       52.040896      15.677953\n",
+       "  1        Dream       47.597791      17.673353\n",
+       "  2    Torgersen       57.004282      15.573296\n",
+       "  3       Biscoe       37.559513      19.513630\n",
+       "  4        Dream       46.813786      15.565273\n",
        "  ..         ...             ...            ...\n",
-       "  195      Dream       41.009962      18.210038\n",
-       "  196      Dream       50.174224      18.534705\n",
-       "  197     Biscoe       42.890809      15.596290\n",
-       "  198  Torgersen       47.442859      15.934119\n",
-       "  199     Biscoe       45.300011      18.687155\n",
+       "  195      Dream       41.010001      18.210018\n",
+       "  196      Dream       50.174267      18.534693\n",
+       "  197     Biscoe       42.890791      15.596277\n",
+       "  198  Torgersen       47.442753      15.934110\n",
+       "  199     Biscoe       45.299875      18.687149\n",
        "  \n",
        "  [200 rows x 3 columns],\n",
        "  'spent_epsilon': 1.0,\n",
        "  'spent_delta': 0.00015673368198174188},\n",
-       " 'timestamp': 1725364925.3299136}"
+       " 'timestamp': 1726137013.8748171}"
       ]
      },
-     "execution_count": 62,
+     "execution_count": 60,
      "metadata": {},
      "output_type": "execute_result"
     }
diff --git a/client/requirements.txt b/client/requirements.txt
index 437c00e2..2b334b30 100644
--- a/client/requirements.txt
+++ b/client/requirements.txt
@@ -2,7 +2,7 @@ diffprivlib==0.6.4
 diffprivlib_logger==0.0.3
 git+https://github.com/ryan112358/private-pgm.git@01f02f17eba440f4e76c1d06fa5ee9eed0bd2bca
 jupyter==1.0.0
-numpy==1.24.0
+numpy==1.26.2
 opendp==0.10.0
 opendp_logger==0.3.0
 pandas==2.2.2
diff --git a/client/setup.py b/client/setup.py
index bdf7fdc3..54412120 100644
--- a/client/setup.py
+++ b/client/setup.py
@@ -10,7 +10,7 @@
 setup(
     name="lomas_client",
     packages=find_packages(),
-    version="0.3.2",
+    version="0.3.3",
     description="A client to interact with the Lomas server.",
     long_description=long_description,
     long_description_content_type="text/markdown",
@@ -25,6 +25,7 @@
         "License :: OSI Approved :: MIT License",
         "Natural Language :: English",
         "Programming Language :: Python :: 3.11",
+        "Programming Language :: Python :: 3.12",
         "Topic :: Software Development :: Libraries",
         "Topic :: Software Development :: Libraries :: Python Modules",
         "Topic :: Scientific/Engineering",
@@ -40,11 +41,11 @@
         "smartnoise-sql",
         "smartnoise-synth",
     ],
-    python_requires=">=3.10, <3.11",
+    python_requires=">=3.11, <3.13",
     install_requires=[
         "diffprivlib>=0.6.4",
         "diffprivlib_logger>=0.0.3",
-        "numpy>=1.24.0",
+        "numpy>=1.26.2",
         "opendp==0.10.0",
         "opendp_logger==0.3.0",
         "pandas>=2.2.2",
diff --git a/docs/versions.yaml b/docs/versions.yaml
index efe80cd2..7d30729e 100644
--- a/docs/versions.yaml
+++ b/docs/versions.yaml
@@ -26,7 +26,7 @@
    tag: "v0.3.0"
    languages:
     - "en"
-"v0.3.2":
-   tag: "v0.3.2"
+"v0.3.3":
+   tag: "v0.3.3"
    languages:
     - "en"
\ No newline at end of file
diff --git a/server/Dockerfile b/server/Dockerfile
index c8b5968f..2103f86b 100644
--- a/server/Dockerfile
+++ b/server/Dockerfile
@@ -1,5 +1,5 @@
 # To install packages including opendp_polars and opendp_logger 
-FROM python:3.11 AS lomas_server_base
+FROM python:3.12 AS lomas_server_base
 
 # To install missing packages
 COPY ./requirements.txt /requirements.txt
@@ -8,25 +8,25 @@ RUN pip install --no-cache-dir --upgrade -r requirements.txt
 WORKDIR /code
 
 FROM lomas_server_base AS lomas_server
-COPY ./lomas_server/ /code/
+COPY ./lomas_server/ /code/lomas_server/
 COPY ./LICENSE /code/LICENSE
 COPY ./data/ /data/
-CMD ["python", "uvicorn_serve.py"]
+CMD ["python", "-m", "lomas_server.uvicorn_serve"]
 
 FROM lomas_server_base AS lomas_server_dev
 ENV PYTHONDONTWRITEBYTECODE 1
-CMD ["python", "uvicorn_serve.py"]
+CMD ["python", "-m", "lomas_server.uvicorn_serve"]
 
 FROM lomas_server_base AS lomas_admin_dashboard_base
 COPY ./requirements_streamlit.txt /requirements_streamlit.txt
 RUN pip install --no-cache-dir --upgrade -r /requirements_streamlit.txt
 
 FROM lomas_admin_dashboard_base AS lomas_admin_dashboard
-COPY ./lomas_server/ /code/
+COPY ./lomas_server/ /code/lomas_server/
 COPY ./LICENSE /code/LICENSE
 COPY ./data/ /data/
-CMD ["streamlit", "run", "administration/dashboard/about.py"]
+CMD ["streamlit", "run", "lomas_server/administration/dashboard/about.py"]
 
 FROM lomas_admin_dashboard_base AS lomas_admin_dashboard_dev
 ENV PYTHONDONTWRITEBYTECODE 1
-CMD ["streamlit", "run", "administration/dashboard/about.py"]
+CMD ["streamlit", "run", "lomas_server/administration/dashboard/about.py"]
diff --git a/server/configs/example_config.yaml b/server/configs/example_config.yaml
index 6683d752..46ca98ed 100644
--- a/server/configs/example_config.yaml
+++ b/server/configs/example_config.yaml
@@ -16,9 +16,6 @@ runtime_args:
       address: "mongodb"
       port: 27017
       db_name: "defaultdb"
-    dataset_store:
-      ds_store_type: "LRU_cache"
-      max_memory_usage: 512
     dp_libraries:
       opendp:
         contrib: True
diff --git a/server/data/collections/metadata/birthday_metadata.yaml b/server/data/collections/metadata/birthday_metadata.yaml
index e37126e6..4962c6c4 100644
--- a/server/data/collections/metadata/birthday_metadata.yaml
+++ b/server/data/collections/metadata/birthday_metadata.yaml
@@ -1,6 +1,7 @@
 max_ids: 1
 row_privacy: true
 censor_dims: False
+rows: 5
 columns:
   birthday:
     type: datetime
diff --git a/server/data/collections/metadata/fso_income_synthetic_metadata.yaml b/server/data/collections/metadata/fso_income_synthetic_metadata.yaml
index 8084d940..8365bf98 100644
--- a/server/data/collections/metadata/fso_income_synthetic_metadata.yaml
+++ b/server/data/collections/metadata/fso_income_synthetic_metadata.yaml
@@ -1,19 +1,44 @@
 max_ids: 1
 row_privacy: true
+rows: 2032543
 columns:
   region:
     type: int
+    precision: 32
+    cardinality: 7
+    categories: [1, 2, 3, 4, 5, 6, 7]
   eco_branch:
     type: int
+    precision: 32
+    cardinality: 72
+    categories: [ 8, 10, 11, 13, 14, 15, 16, 17, 18, 20, 21, 22, 23, 24, 25, 26, 27,
+       28, 29, 30, 31, 32, 33, 35, 37, 38, 41, 42, 43, 45, 46, 47, 49, 50,
+       52, 53, 55, 56, 58, 59, 60, 61, 62, 63, 64, 65, 66, 68, 69, 70, 71,
+       72, 73, 74, 75, 77, 78, 79, 80, 81, 82, 85, 86, 87, 88, 90, 91, 92,
+       93, 94, 95, 96]
   profession:
     type: int
+    precision: 32
+    cardinality: 28
+    categories: [10, 21, 22, 23, 24, 25, 31, 32, 33, 34, 41, 42, 43, 51, 52, 53, 61,
+       62, 71, 72, 73, 74, 81, 83, 91, 92, 93, 94]
   education:
     type: int
+    precision: 32
+    cardinality: 8
+    categories: [1, 2, 3, 4, 5, 6, 7, 8]
   age:
     type: int
+    precision: 32
+    lower: 0
+    upper: 120
   sex:
     type: int
+    precision: 32
+    cardinality: 2
+    categories: [0, 1]
   income:
     type: float
+    precision: 64
     lower: 1000
     upper: 100000
diff --git a/server/data/collections/metadata/iris_metadata.yaml b/server/data/collections/metadata/iris_metadata.yaml
index 91644d08..a1e74cc3 100644
--- a/server/data/collections/metadata/iris_metadata.yaml
+++ b/server/data/collections/metadata/iris_metadata.yaml
@@ -3,18 +3,22 @@ row_privacy: true
 columns:
   petal_length:
     type: float
+    precision: 64
     lower: 0.5
     upper: 10.0
   petal_width:
     type: float
+    precision: 64
     lower: 0.05
     upper: 5.0
   sepal_length:
     type: float
+    precision: 64
     lower: 2.0
     upper: 10.0
   sepal_width:
     type: float
+    precision: 64
     lower: 1.0
     upper: 6.0
   species:
diff --git a/server/data/collections/metadata/penguin_metadata.yaml b/server/data/collections/metadata/penguin_metadata.yaml
index c8fdbd6e..b9770ddb 100644
--- a/server/data/collections/metadata/penguin_metadata.yaml
+++ b/server/data/collections/metadata/penguin_metadata.yaml
@@ -18,18 +18,22 @@ columns:
     - Dream
   bill_length_mm:
     type: float
+    precision: 64
     lower: 30.0
     upper: 65.0
   bill_depth_mm:
     type: float
+    precision: 64
     lower: 13.0
     upper: 23.0
   flipper_length_mm:
     type: float
+    precision: 64
     lower: 150.0
     upper: 250.0
   body_mass_g:
     type: float
+    precision: 64
     lower: 2000.0
     upper: 7000.0
   sex:
diff --git a/server/data/collections/metadata/pums_metadata.yaml b/server/data/collections/metadata/pums_metadata.yaml
index e89d55b8..a9a45064 100644
--- a/server/data/collections/metadata/pums_metadata.yaml
+++ b/server/data/collections/metadata/pums_metadata.yaml
@@ -5,16 +5,24 @@ columns:
   age:
     private_id: true # for testing purposes
     type: int
+    precision: 32
+    lower: 0
+    upper: 120
   sex:
     type: boolean
   educ:
     type: int
-    cardinality: 14
+    precision: 32
+    cardinality: 16
+    categories: [ 9,  1, 11,  5,  8, 13, 10,  7, 14,  4, 15, 16,  3, 12,  2,  6]
   race:
     type: int
+    precision: 32
     cardinality: 6
+    categories: [1, 3, 4, 2, 6, 5]
   income:
     type: float
+    precision: 64
     lower: 0.0
     upper: 500_000.0
   married:
diff --git a/server/data/collections/metadata/titanic_metadata.yaml b/server/data/collections/metadata/titanic_metadata.yaml
index 4b42ab10..6e741b7e 100644
--- a/server/data/collections/metadata/titanic_metadata.yaml
+++ b/server/data/collections/metadata/titanic_metadata.yaml
@@ -1,11 +1,10 @@
 max_ids: 1
 row_privacy: true
+rows: 887
 columns:
-  PassengerId:
-    private_id: true
-    type: int
   Pclass:
     type: int
+    precision: 32
     lower: 1
     upper: 3
   Name:
@@ -18,20 +17,24 @@ columns:
     - female
   Age:
     type: float
+    precision: 64
     lower: 0.1
     upper: 100.0
   SibSp:
     type: int
+    precision: 32
     lower: 0
     upper: 10
   Parch:
     type: int
+    precision: 32
     lower: 0
     upper: 10
   Ticket:
     type: string
   Fare:
     type: float
+    precision: 64
     lower: 0.0
     upper: 1000
   Cabin:
diff --git a/server/deploy/helm/charts/lomas_server/Chart.yaml b/server/deploy/helm/charts/lomas_server/Chart.yaml
index 71aa5ffa..958ec212 100644
--- a/server/deploy/helm/charts/lomas_server/Chart.yaml
+++ b/server/deploy/helm/charts/lomas_server/Chart.yaml
@@ -4,9 +4,9 @@ description: Lomas deployment chart
 
 type: application
 
-version: 0.3.2
+version: 0.3.3
 
-appVersion: "0.3.2"
+appVersion: "0.3.3"
 
 dependencies:
   - name: mongodb
diff --git a/server/deploy/helm/charts/lomas_server/templates/server/configmap.yaml b/server/deploy/helm/charts/lomas_server/templates/server/configmap.yaml
index eeac0d64..53a81428 100644
--- a/server/deploy/helm/charts/lomas_server/templates/server/configmap.yaml
+++ b/server/deploy/helm/charts/lomas_server/templates/server/configmap.yaml
@@ -19,9 +19,6 @@ data:
           password: {{ index .Values.mongodb.auth.passwords 0 }}
           username: {{ index .Values.mongodb.auth.usernames 0 }}
           db_name: {{ index .Values.mongodb.auth.databases 0}}
-        dataset_store:
-          ds_store_type: {{ .Values.server.runtime_args.settings.dataset_store.ds_store_type }}
-          max_memory_usage : {{ .Values.server.runtime_args.settings.dataset_store.max_memory_usage }}
         dp_libraries:
           opendp:
             contrib: {{ .Values.server.runtime_args.settings.dp_libraries.opendp.contrib }}
diff --git a/server/deploy/helm/charts/lomas_server/values.schema.json b/server/deploy/helm/charts/lomas_server/values.schema.json
index 2c43891a..9c96a810 100644
--- a/server/deploy/helm/charts/lomas_server/values.schema.json
+++ b/server/deploy/helm/charts/lomas_server/values.schema.json
@@ -277,23 +277,6 @@
                         "settings": {
                             "type": "object",
                             "properties": {
-                                "dataset_store": {
-                                    "type": "object",
-                                    "properties": {
-                                        "ds_store_type": {
-                                            "description": "The type of dataset store.",
-                                            "enum": [
-                                                "basic",
-                                                "LRU_cache"
-                                            ],
-                                            "type": "string"
-                                        },
-                                        "max_memory_usage": {
-                                            "description": "Max memory usage of dataset store. Only useful for LRU_cache type.",
-                                            "type": "integer"
-                                        }
-                                    }
-                                },
                                 "develop_mode": {
                                     "description": "When True, service database is reset to default user and dataset collections.",
                                     "type": "boolean"
diff --git a/server/deploy/helm/charts/lomas_server/values.yaml b/server/deploy/helm/charts/lomas_server/values.yaml
index 26b6fbb8..a30ef4e7 100644
--- a/server/deploy/helm/charts/lomas_server/values.yaml
+++ b/server/deploy/helm/charts/lomas_server/values.yaml
@@ -57,9 +57,6 @@ server:
         time_attack:
           method: "jitter" # or "stall"
           magnitude: 1
-      dataset_store:
-        ds_store_type: "LRU_cache"
-        max_memory_usage : 512
       dp_libraries:
         opendp:
           contrib: True
diff --git a/server/docker-compose.yml b/server/docker-compose.yml
index 6125c114..f67c22e7 100644
--- a/server/docker-compose.yml
+++ b/server/docker-compose.yml
@@ -9,7 +9,7 @@ services:
       - 80:80
     restart: always
     volumes:
-      - ./lomas_server:/code/
+      - ./lomas_server:/code/lomas_server
       - ./configs/example_config.yaml:/usr/lomas_server/runtime.yaml
       - ./configs/example_secrets.yaml:/usr/lomas_server/secrets.yaml
       - ./data/collections/metadata:/data/collections/metadata/
@@ -87,7 +87,7 @@ services:
       - 8501:8501
     restart: always
     volumes:
-      - ./lomas_server:/code/
+      - ./lomas_server:/code/lomas_server
       - ./configs/example_config.yaml:/usr/lomas_server/runtime.yaml
       - ./configs/example_secrets.yaml:/usr/lomas_server/secrets.yaml
       - ./configs/example_dashboard_config.yaml:/usr/lomas_dashboard/dashboard.yaml
diff --git a/server/lomas_server/admin_database/admin_database.py b/server/lomas_server/admin_database/admin_database.py
index 0c983fe2..5d7a8f6b 100644
--- a/server/lomas_server/admin_database/admin_database.py
+++ b/server/lomas_server/admin_database/admin_database.py
@@ -4,11 +4,12 @@
 from functools import wraps
 from typing import Callable, Dict, List
 
-from constants import MODEL_INPUT_TO_LIB
-from utils.error_handler import (
+from lomas_server.utils.collection_models import Metadata
+from lomas_server.utils.error_handler import (
     InvalidQueryException,
     UnauthorizedAccessException,
 )
+from lomas_server.utils.query_models import RequestModel, model_input_to_lib
 
 
 def user_must_exist(func: Callable) -> Callable:  # type: ignore
@@ -155,7 +156,7 @@ def does_dataset_exist(self, dataset_name: str) -> bool:
     @abstractmethod
     @dataset_must_exist
     @user_must_have_access_to_dataset
-    def get_dataset_metadata(self, dataset_name: str) -> dict:
+    def get_dataset_metadata(self, dataset_name: str) -> Metadata:
         """
         Returns the metadata dictionnary of the dataset.
 
@@ -165,7 +166,7 @@ def get_dataset_metadata(self, dataset_name: str) -> dict:
             dataset_name (str): name of the dataset to get the metadata
 
         Returns:
-            dict: The metadata dict.
+            Metadata: The metadata object.
         """
 
     @abstractmethod
@@ -425,14 +426,14 @@ def get_user_previous_queries(
         """
 
     def prepare_save_query(
-        self, user_name: str, query_json: dict, response: dict
+        self, user_name: str, query_json: RequestModel, response: dict
     ) -> dict:
         """
         Prepare the query to save in archives
 
         Args:
             user_name (str): name of the user
-            query_json (dict): json received from client
+            query_json (RequestModel): request received from client
             response (dict): response sent to the client
 
         Raises:
@@ -441,11 +442,10 @@ def prepare_save_query(
         Returns:
             dict: The query archive dictionary.
         """
-        model_input = query_json.__class__.__name__
         to_archive = {
             "user_name": user_name,
             "dataset_name": query_json.dataset_name,
-            "dp_librairy": MODEL_INPUT_TO_LIB[model_input],
+            "dp_librairy": model_input_to_lib(query_json),
             "client_input": query_json.model_dump(),
             "response": response,
             "timestamp": time.time(),
@@ -455,7 +455,7 @@ def prepare_save_query(
 
     @abstractmethod
     def save_query(
-        self, user_name: str, query_json: dict, response: dict
+        self, user_name: str, query_json: RequestModel, response: dict
     ) -> None:
         """
         Save queries of user on datasets in a separate collection (table)
diff --git a/server/lomas_server/admin_database/factory.py b/server/lomas_server/admin_database/factory.py
index f8ad7b02..814f4db1 100644
--- a/server/lomas_server/admin_database/factory.py
+++ b/server/lomas_server/admin_database/factory.py
@@ -1,10 +1,9 @@
-from admin_database.admin_database import AdminDatabase
-from admin_database.mongodb_database import AdminMongoDatabase
-from admin_database.utils import get_mongodb_url
-from admin_database.yaml_database import AdminYamlDatabase
-from constants import AdminDBType
-from utils.config import DBConfig
-from utils.error_handler import InternalServerException
+from lomas_server.admin_database.admin_database import AdminDatabase
+from lomas_server.admin_database.mongodb_database import AdminMongoDatabase
+from lomas_server.admin_database.utils import get_mongodb_url
+from lomas_server.admin_database.yaml_database import AdminYamlDatabase
+from lomas_server.utils.config import DBConfig, MongoDBConfig, YamlDBConfig
+from lomas_server.utils.error_handler import InternalServerException
 
 
 def admin_database_factory(config: DBConfig) -> AdminDatabase:
@@ -21,19 +20,15 @@ def admin_database_factory(config: DBConfig) -> AdminDatabase:
     Returns:
         AdminDatabase: A instance of the correct type of AdminDatabase.
     """
-    db_type = config.db_type
 
-    match db_type:
-        case AdminDBType.MONGODB:
+    match config:
+        case MongoDBConfig():
             db_url = get_mongodb_url(config)
             db_name = config.db_name
             return AdminMongoDatabase(db_url, db_name)
 
-        case AdminDBType.YAML:
+        case YamlDBConfig():
             yaml_database_file = config.db_file
             return AdminYamlDatabase(yaml_database_file)
-
         case _:
-            raise InternalServerException(
-                f"Database type {db_type} not supported."
-            )
+            raise InternalServerException("Database type not supported.")
diff --git a/server/lomas_server/admin_database/mongodb_database.py b/server/lomas_server/admin_database/mongodb_database.py
index 0f8509f3..e0149352 100644
--- a/server/lomas_server/admin_database/mongodb_database.py
+++ b/server/lomas_server/admin_database/mongodb_database.py
@@ -5,13 +5,15 @@
 from pymongo.errors import WriteConcernError
 from pymongo.results import _WriteResult
 
-from admin_database.admin_database import (
+from lomas_server.admin_database.admin_database import (
     AdminDatabase,
     dataset_must_exist,
     user_must_exist,
     user_must_have_access_to_dataset,
 )
-from utils.error_handler import InvalidQueryException
+from lomas_server.utils.collection_models import Metadata
+from lomas_server.utils.error_handler import InvalidQueryException
+from lomas_server.utils.query_models import RequestModel
 
 
 class AdminMongoDatabase(AdminDatabase):
@@ -59,7 +61,7 @@ def does_dataset_exist(self, dataset_name: str) -> bool:
         return False
 
     @dataset_must_exist
-    def get_dataset_metadata(self, dataset_name: str) -> dict:
+    def get_dataset_metadata(self, dataset_name: str) -> Metadata:
         """Returns the metadata dictionnary of the dataset.
 
         Wrapped by :py:func:`dataset_must_exist`.
@@ -68,12 +70,12 @@ def get_dataset_metadata(self, dataset_name: str) -> dict:
             dataset_name (str): name of the dataset to get the metadata
 
         Returns:
-            dict: The metadata dict.
+            Metadata: The metadata model.
         """
         metadatas = self.db.metadata.find_one(
             {dataset_name: {"$exists": True}}
         )
-        return metadatas[dataset_name]  # type: ignore
+        return Metadata.model_validate(metadatas[dataset_name])  # type: ignore
 
     @user_must_exist
     def set_may_user_query(self, user_name: str, may_query: bool) -> None:
@@ -248,14 +250,14 @@ def get_user_previous_queries(
         return list(queries)
 
     def save_query(
-        self, user_name: str, query_json: dict, response: dict
+        self, user_name: str, query_json: RequestModel, response: dict
     ) -> None:
         """Save queries of user on datasets in a separate collection (table)
         named "queries_archives" in the DB
 
         Args:
             user_name (str): name of the user
-            query_json (dict): json received from client
+            query_json (RequestModel): json received from client
             response (dict): response sent to the client
 
         Raises:
diff --git a/server/lomas_server/admin_database/utils.py b/server/lomas_server/admin_database/utils.py
index 2d37c2e0..03f3c1d7 100644
--- a/server/lomas_server/admin_database/utils.py
+++ b/server/lomas_server/admin_database/utils.py
@@ -1,22 +1,21 @@
-from types import SimpleNamespace
-
 from pymongo import MongoClient
 from pymongo.database import Database
 
-from mongodb_admin import (
+from lomas_server.mongodb_admin import (
     add_datasets_via_yaml,
     add_users_via_yaml,
     drop_collection,
 )
-from utils.config import DBConfig, get_config
-from utils.logger import LOG
+from lomas_server.utils.config import MongoDBConfig, get_config
+from lomas_server.utils.error_handler import InternalServerException
+from lomas_server.utils.logger import LOG
 
 
-def get_mongodb_url(config: DBConfig) -> str:
+def get_mongodb_url(config: MongoDBConfig) -> str:
     """Get URL of the administration MongoDB.
 
     Args:
-        config (DBConfig): An instance of DBConfig.
+        config (MongoDBConfig): An instance of DBConfig.
 
     Returns:
         str: A correctly formatted url for connecting to the
@@ -46,9 +45,15 @@ def get_mongodb() -> Database:
         str: A correctly formatted url for connecting to the
             MongoDB database.
     """
-    db_args = SimpleNamespace(**vars(get_config().admin_database))
-    db_url = get_mongodb_url(db_args)
-    return MongoClient(db_url)[db_args.db_name]
+    admin_config = get_config().admin_database
+    if isinstance(admin_config, MongoDBConfig):
+        db_url = get_mongodb_url(admin_config)
+    else:
+        raise InternalServerException(
+            "Expected MongoDBConfig, found {type(admin_config)}."
+        )
+
+    return MongoClient(db_url)[admin_config.db_name]
 
 
 def add_demo_data_to_mongodb_admin(
diff --git a/server/lomas_server/admin_database/yaml_database.py b/server/lomas_server/admin_database/yaml_database.py
index 069a93ec..7eddc56c 100644
--- a/server/lomas_server/admin_database/yaml_database.py
+++ b/server/lomas_server/admin_database/yaml_database.py
@@ -3,13 +3,18 @@
 
 import yaml
 
-from admin_database.admin_database import (
+from lomas_server.admin_database.admin_database import (
     AdminDatabase,
     dataset_must_exist,
     user_must_exist,
     user_must_have_access_to_dataset,
 )
-from utils.error_handler import InternalServerException, InvalidQueryException
+from lomas_server.utils.collection_models import Metadata
+from lomas_server.utils.error_handler import (
+    InternalServerException,
+    InvalidQueryException,
+)
+from lomas_server.utils.query_models import RequestModel
 
 
 class AdminYamlDatabase(AdminDatabase):
@@ -58,7 +63,7 @@ def does_dataset_exist(self, dataset_name: str) -> bool:
         return False
 
     @dataset_must_exist
-    def get_dataset_metadata(self, dataset_name: str) -> dict:
+    def get_dataset_metadata(self, dataset_name: str) -> Metadata:
         """Returns the metadata dictionnary of the dataset.
 
         Wrapped by :py:func:`dataset_must_exist`.
@@ -67,7 +72,7 @@ def get_dataset_metadata(self, dataset_name: str) -> dict:
             dataset_name (str): name of the dataset to get the metadata
 
         Returns:
-            dict: The metadata dict.
+            Metadata: The metadata model.
         """
         for dt in self.database["datasets"]:
             if dt["dataset_name"] == dataset_name:
@@ -76,7 +81,7 @@ def get_dataset_metadata(self, dataset_name: str) -> dict:
                 with open(metadata_path, mode="r", encoding="utf-8") as f:
                     metadata = yaml.safe_load(f)
 
-        return metadata
+        return Metadata.model_validate(metadata)
 
     @user_must_exist
     def set_may_user_query(self, user_name: str, may_query: bool) -> None:
@@ -248,14 +253,14 @@ def get_user_previous_queries(
         return previous_queries
 
     def save_query(
-        self, user_name: str, query_json: dict, response: dict
+        self, user_name: str, query_json: RequestModel, response: dict
     ) -> None:
         """Save queries of user on datasets in a separate collection (table)
         named "queries_archives" in the DB
 
         Args:
             user_name (str): name of the user
-            query_json (dict): json received from client
+            query_json (RequestModel): request received from client
             response (dict): response sent to the client
         """
         to_archive = super().prepare_save_query(
diff --git a/server/lomas_server/administration/dashboard/about.py b/server/lomas_server/administration/dashboard/about.py
index 8bf0e936..31cb9999 100644
--- a/server/lomas_server/administration/dashboard/about.py
+++ b/server/lomas_server/administration/dashboard/about.py
@@ -17,7 +17,7 @@
     sys.path.append(src_dir)
 
     st.set_page_config(page_title="Lomas Dashboard")
-    FOLDER = "./administration/dashboard"  # TODO move
+    FOLDER = "./lomas_server/administration/dashboard"  # TODO move
     show_pages(
         [
             Page(f"{FOLDER}/about.py", "Home Page", "🏠"),
diff --git a/server/lomas_server/administration/dashboard/pages/01_server_overview.py b/server/lomas_server/administration/dashboard/pages/01_server_overview.py
index 0480d1b6..d77493cf 100644
--- a/server/lomas_server/administration/dashboard/pages/01_server_overview.py
+++ b/server/lomas_server/administration/dashboard/pages/01_server_overview.py
@@ -2,7 +2,7 @@
 import streamlit as st
 from config import get_config
 
-from constants import AdminDBType, DatasetStoreType
+from constants import AdminDBType
 from utils.config import get_config as get_server_config
 from utils.error_handler import InternalServerException
 
@@ -62,7 +62,7 @@ def get_server_data(fastapi_address, endpoint):
     else:
         st.write(":red[The server is in PRODUCTION mode.]")
 
-    tab_1, tab_2, tab_3 = st.columns(3)
+    tab_1, tab_2 = st.columns(2)
     with tab_1:
         st.subheader("Server configurations")
 
@@ -108,23 +108,3 @@ def get_server_data(fastapi_address, endpoint):
             raise InternalServerException(
                 f"Admin database type {db_type} not supported."
             )
-
-    with tab_3:
-        st.subheader("Dataset Store")
-        ds_store_type = st.session_state.config.dataset_store.ds_store_type
-        st.write("The dataset store type is: ", ds_store_type)
-        if ds_store_type == DatasetStoreType.LRU:
-            st.write(
-                "The maximum memory usage is: ",
-                st.session_state.config.dataset_store.max_memory_usage,
-                "MB.",
-            )
-            memory = get_server_data(
-                st.session_state.dashboard_config.server_service,
-                "get_memory_usage",
-            )
-            st.write(
-                "Current memory usage with loaded datasets: ",
-                memory["memory_usage"],
-                "MB.",
-            )
diff --git a/server/lomas_server/app.py b/server/lomas_server/app.py
index 74cde1dc..085947ef 100644
--- a/server/lomas_server/app.py
+++ b/server/lomas_server/app.py
@@ -4,23 +4,25 @@
 
 from fastapi import FastAPI, Request, Response
 
-from admin_database.factory import admin_database_factory
-from admin_database.utils import add_demo_data_to_mongodb_admin
-from constants import (
+from lomas_server.admin_database.factory import admin_database_factory
+from lomas_server.admin_database.utils import add_demo_data_to_mongodb_admin
+from lomas_server.constants import (
     CONFIG_NOT_LOADED,
     DB_NOT_LOADED,
-    QUERY_HANDLER_NOT_LOADED,
     SERVER_LIVE,
     AdminDBType,
 )
-from dataset_store.factory import dataset_store_factory
-from dp_queries.dp_libraries.opendp import set_opendp_features_config
-from dp_queries.dp_logic import QueryHandler
-from routes import routes_admin, routes_dp
-from utils.anti_timing_att import anti_timing_att
-from utils.config import get_config
-from utils.error_handler import InternalServerException, add_exception_handlers
-from utils.logger import LOG
+from lomas_server.dp_queries.dp_libraries.opendp import (
+    set_opendp_features_config,
+)
+from lomas_server.routes import routes_admin, routes_dp
+from lomas_server.utils.anti_timing_att import anti_timing_att
+from lomas_server.utils.config import get_config
+from lomas_server.utils.error_handler import (
+    InternalServerException,
+    add_exception_handlers,
+)
+from lomas_server.utils.logger import LOG
 
 
 @asynccontextmanager
@@ -45,8 +47,6 @@ async def lifespan(
 
     # Set some app state
     app.state.admin_database = None
-    app.state.query_handler = None
-    app.state.dataset_store = None
 
     # General server state, can add fields if need be.
     app.state.server_state = {
@@ -62,6 +62,7 @@ async def lifespan(
         LOG.info("Loading config")
         app.state.server_state["message"].append("Loading config")
         config = get_config()
+        app.state.private_credentials = config.private_db_credentials
     except InternalServerException:
         LOG.info("Config could not loaded")
         app.state.server_state["state"].append(CONFIG_NOT_LOADED)
@@ -91,7 +92,7 @@ async def lifespan(
                 config.admin_database
             )
         except InternalServerException as e:
-            LOG.exception("Failed at startup:" + str(e))
+            LOG.exception(f"Failed at startup: {str(e)}")
             app.state.server_state["state"].append(DB_NOT_LOADED)
             app.state.server_state["message"].append(
                 f"Admin database could not be loaded: {str(e)}"
@@ -99,33 +100,9 @@ async def lifespan(
             app.state.server_state["LIVE"] = False
             status_ok = False
 
-    # Load query handler
-    if status_ok:
-        LOG.info("Loading query handler")
-        app.state.server_state["message"].append("Loading dataset store")
-        app.state.dataset_store = dataset_store_factory(
-            config.dataset_store,
-            app.state.admin_database,
-            config.private_db_credentials,
-        )
-
-        app.state.server_state["message"].append("Loading query handler")
-        app.state.query_handler = QueryHandler(
-            app.state.admin_database, app.state.dataset_store
-        )
-
         app.state.server_state["state"].append("Startup completed")
         app.state.server_state["message"].append("Startup completed")
 
-        if app.state.query_handler is None:
-            LOG.info("QueryHandler not loaded")
-            app.state.server_state["state"].append(QUERY_HANDLER_NOT_LOADED)
-            app.state.server_state["message"].append(
-                "Server could not be started!"
-            )
-            app.state.server_state["LIVE"] = False
-            status_ok = False
-
     # Set DP Libraries config
     set_opendp_features_config(config.dp_libraries.opendp)
 
diff --git a/server/lomas_server/constants.py b/server/lomas_server/constants.py
index d67b098f..2e06156b 100644
--- a/server/lomas_server/constants.py
+++ b/server/lomas_server/constants.py
@@ -2,6 +2,8 @@
 import string
 from enum import StrEnum
 
+import pkg_resources
+
 # Get config and secrets from correct location
 if "LOMAS_CONFIG_PATH" in os.environ:
     CONFIG_PATH = f"""{os.environ.get("LOMAS_CONFIG_PATH")}"""
@@ -29,9 +31,6 @@ class ConfigKeys(StrEnum):
     DB_TYPE_MONGODB: str = "mongodb"
     MONGODB_ADDR: str = "address"
     MONGODB_PORT: str = "port"
-    DATASET_STORE: str = "dataset_store"
-    DATASET_STORE_TYPE: str = "ds_store_type"
-    LRU_DATASET_STORE_MAX_SIZE: str = "max_memory_usage"
     DP_LIBRARY: str = "dp_libraries"
 
 
@@ -42,13 +41,6 @@ class AdminDBType(StrEnum):
     MONGODB: str = "mongodb"
 
 
-class DatasetStoreType(StrEnum):
-    """Types of classes to handle datasets in memory"""
-
-    BASIC: str = "basic"
-    LRU: str = "LRU_cache"
-
-
 class TimeAttackMethod(StrEnum):
     """Possible methods against timing attacks"""
 
@@ -57,7 +49,6 @@ class TimeAttackMethod(StrEnum):
 
 
 # Server states
-QUERY_HANDLER_NOT_LOADED = "QueryHander not loaded"
 DB_NOT_LOADED = "User database not loaded"
 CONFIG_NOT_LOADED = "Config not loaded"
 SERVER_LIVE = "LIVE"
@@ -68,13 +59,11 @@ class TimeAttackMethod(StrEnum):
 )
 
 # General values
-DEFAULT_DATE_FORMAT = "%Y-%m-%d"
 SECONDS_IN_A_DAY = 60 * 60 * 24
 
-
-# DP constants
-EPSILON_LIMIT: float = 5.0
-DELTA_LIMIT: float = 0.0004
+# DP constants (max budget per user per dataset)
+EPSILON_LIMIT: float = 10.0
+DELTA_LIMIT: float = 0.01
 
 
 # Supported DP libraries
@@ -87,15 +76,6 @@ class DPLibraries(StrEnum):
     DIFFPRIVLIB = "diffprivlib"
 
 
-# Query model input to DP librairy
-MODEL_INPUT_TO_LIB = {
-    "SmartnoiseSQLModel": DPLibraries.SMARTNOISE_SQL,
-    "SmartnoiseSynthQueryModel": DPLibraries.SMARTNOISE_SYNTH,
-    "OpenDPModel": DPLibraries.OPENDP,
-    "DiffPrivLibModel": DPLibraries.DIFFPRIVLIB,
-}
-
-
 # Private Databases
 class PrivateDatabaseType(StrEnum):
     """Type of Private Database for the private data"""
@@ -178,13 +158,9 @@ class OpenDPDatasetInputMetric(StrEnum):
 # Dummy dataset generation
 DUMMY_NB_ROWS = 100
 DUMMY_SEED = 42
-DEFAULT_NUMERICAL_MIN = -10000
-DEFAULT_NUMERICAL_MAX = 10000
 RANDOM_STRINGS = list(
     string.ascii_lowercase + string.ascii_uppercase + string.digits
 )
-RANDOM_DATE_START = "01/01/2000"
-RANDOM_DATE_RANGE = 50 * 365 * 24 * 60 * 60  # 50 years
 NB_RANDOM_NONE = 5  # if nullable, how many random none to add
 
 
@@ -192,8 +168,9 @@ class OpenDPDatasetInputMetric(StrEnum):
 NUMERICAL_DTYPES = ["int16", "int32", "int64", "float16", "float32", "float64"]
 
 # Example pipeline inputs
+OPENDP_VERSION = pkg_resources.get_distribution("opendp").version
 OPENDP_PIPELINE = (
-    '{"version": "0.10.0", '
+    f'{{"version": "{OPENDP_VERSION}", '
     '"ast": {'
     '"_type": "partial_chain", "lhs": {'
     '"_type": "partial_chain", "lhs": {'
diff --git a/server/lomas_server/data_connector/data_connector.py b/server/lomas_server/data_connector/data_connector.py
index 3aa7957f..e7356e59 100644
--- a/server/lomas_server/data_connector/data_connector.py
+++ b/server/lomas_server/data_connector/data_connector.py
@@ -3,8 +3,7 @@
 
 import pandas as pd
 
-from dataset_store.data_connector_observer import DataConnectorObserver
-from utils.collection_models import Metadata
+from lomas_server.utils.collection_models import DatetimeMetadata, Metadata
 
 
 class DataConnector(ABC):
@@ -20,73 +19,47 @@ def __init__(self, metadata: Metadata) -> None:
         Args:
             metadata (Metadata): The metadata for this dataset
         """
-        self.metadata: dict = metadata
-        self.dataset_observers: List[DataConnectorObserver] = []
+        self.metadata: Metadata = metadata
 
-        dtypes, datetime_columns = get_column_dtypes(metadata)
+        dtypes, datetime_columns = get_column_dtypes(self.metadata)
         self.dtypes: Dict[str, str] = dtypes
         self.datetime_columns: List[str] = datetime_columns
 
     @abstractmethod
-    def get_pandas_df(self, dataset_name: str) -> pd.DataFrame:
+    def get_pandas_df(self) -> pd.DataFrame:
         """Get the data in pandas dataframe format
 
-        Args:
-            dataset_name (str): name of the private dataset
-
         Returns:
             pd.DataFrame: The pandas dataframe for this dataset.
         """
 
-    def get_metadata(self) -> dict:
+    def get_metadata(self) -> Metadata:
         """Get the metadata for this dataset
 
         Returns:
-            dict: The metadata dictionary.
+            Metadata: The metadata object.
         """
         return self.metadata
 
-    def get_memory_usage(self) -> int:
-        """Returns the memory usage of this dataset, in MiB.
-
-        The number returned only takes into account the memory usage
-        of the pandas DataFrame "cached" in the instance.
-
-        Returns:
-            int: The memory usage, in MiB.
-        """
-        if self.df is None:
-            return 0
-        return self.df.memory_usage().sum() / (1024**2)
-
-    def subscribe_for_memory_usage_updates(
-        self, dataset_observer: DataConnectorObserver
-    ) -> None:
-        """Add the DataConnectorObserver to the list of dataset_observers.
-
-        Args:
-            dataset_observer (DataConnectorObserver):
-                The observer of this dataset.
-        """
-        self.dataset_observers.append(dataset_observer)
-
 
-def get_column_dtypes(metadata: dict) -> Tuple[Dict[str, str], List[str]]:
-    """Extract and return the column types from the metadata.
+def get_column_dtypes(metadata: Metadata) -> Tuple[Dict[str, str], List[str]]:
+    """Extracts and returns the column types from the metadata.
 
     Args:
-        metadata (dict): The metadata dictionary.
+        metadata (Metadata): The metadata.
 
     Returns:
-        dict: The dictionary of the column type.
-        list: The list of columns of datetime type
+        Tuple[Dict[str, str], List[str]]:
+           dict: The dictionary of the column type.
+            list: The list of columns of datetime type
     """
+
     dtypes = {}
     datetime_columns = []
-    for col_name, data in metadata["columns"].items():
-        if data["type"] == "datetime":
+    for col_name, data in metadata.columns.items():
+        if isinstance(data, DatetimeMetadata):
             dtypes[col_name] = "string"
             datetime_columns.append(col_name)
         else:
-            dtypes[col_name] = data["type"]
+            dtypes[col_name] = data.type
     return dtypes, datetime_columns
diff --git a/server/lomas_server/data_connector/factory.py b/server/lomas_server/data_connector/factory.py
index bea7f19d..7fabc5d2 100644
--- a/server/lomas_server/data_connector/factory.py
+++ b/server/lomas_server/data_connector/factory.py
@@ -1,12 +1,12 @@
 from typing import List
 
-from admin_database.admin_database import AdminDatabase
-from constants import PrivateDatabaseType
-from data_connector.data_connector import DataConnector
-from data_connector.path_connector import PathConnector
-from data_connector.s3_connector import S3Connector
-from utils.config import PrivateDBCredentials, S3CredentialsConfig
-from utils.error_handler import InternalServerException
+from lomas_server.admin_database.admin_database import AdminDatabase
+from lomas_server.constants import PrivateDatabaseType
+from lomas_server.data_connector.data_connector import DataConnector
+from lomas_server.data_connector.path_connector import PathConnector
+from lomas_server.data_connector.s3_connector import S3Connector
+from lomas_server.utils.config import PrivateDBCredentials, S3CredentialsConfig
+from lomas_server.utils.error_handler import InternalServerException
 
 
 def data_connector_factory(
@@ -39,7 +39,7 @@ def data_connector_factory(
             dataset_path = admin_database.get_dataset_field(
                 dataset_name, "dataset_path"
             )
-            private_db = PathConnector(ds_metadata, dataset_path)
+            return PathConnector(ds_metadata, dataset_path)
         case PrivateDatabaseType.S3:
 
             credentials_name = admin_database.get_dataset_field(
@@ -50,6 +50,11 @@ def data_connector_factory(
                 private_db_credentials, database_type, credentials_name
             )
 
+            if not isinstance(credentials, S3CredentialsConfig):
+                raise InternalServerException(
+                    "Could not get correct credentials"
+                )
+
             credentials.endpoint_url = admin_database.get_dataset_field(
                 dataset_name, "endpoint_url"
             )
@@ -60,14 +65,12 @@ def data_connector_factory(
                 dataset_name, "key"
             )
 
-            private_db = S3Connector(ds_metadata, credentials)
+            return S3Connector(ds_metadata, credentials)
         case _:
             raise InternalServerException(
                 f"Unknown database type: {database_type}"
             )
 
-    return private_db
-
 
 def get_dataset_credentials(
     private_db_credentials: List[PrivateDBCredentials],
@@ -80,10 +83,9 @@ def get_dataset_credentials(
     credentials name.
 
     Args:
-        private_db_credentials (List[PrivateDBCredentials]):\
+        private_db_credentials (Sequence[PrivateDBCredentials]):\
             The list of private database credentials.
         db_type (PrivateDatabaseType): The type of the database.
-        ds_infos (Dict[str, str]): _description_
 
     Raises:
         InternalServerException: If the credentials are not found.
diff --git a/server/lomas_server/data_connector/in_memory_connector.py b/server/lomas_server/data_connector/in_memory_connector.py
index d292a3bf..21987d87 100644
--- a/server/lomas_server/data_connector/in_memory_connector.py
+++ b/server/lomas_server/data_connector/in_memory_connector.py
@@ -1,7 +1,7 @@
 import pandas as pd
 
-from data_connector.data_connector import DataConnector
-from utils.collection_models import Metadata
+from lomas_server.data_connector.data_connector import DataConnector
+from lomas_server.utils.collection_models import Metadata
 
 
 class InMemoryConnector(DataConnector):
@@ -21,7 +21,7 @@ def __init__(
             dataset_df (pd.DataFrame): Dataframe of the dataset
         """
         super().__init__(metadata)
-        self.df = dataset_df.copy()
+        self.df: pd.DataFrame = dataset_df.copy()
 
     def get_pandas_df(self) -> pd.DataFrame:
         """Get the data in pandas dataframe format
diff --git a/server/lomas_server/data_connector/path_connector.py b/server/lomas_server/data_connector/path_connector.py
index fff6a195..b3c1e706 100644
--- a/server/lomas_server/data_connector/path_connector.py
+++ b/server/lomas_server/data_connector/path_connector.py
@@ -2,9 +2,12 @@
 
 import pandas as pd
 
-from data_connector.data_connector import DataConnector
-from utils.collection_models import Metadata
-from utils.error_handler import InternalServerException, InvalidQueryException
+from lomas_server.data_connector.data_connector import DataConnector
+from lomas_server.utils.collection_models import Metadata
+from lomas_server.utils.error_handler import (
+    InternalServerException,
+    InvalidQueryException,
+)
 
 
 class PathConnector(DataConnector):
@@ -57,8 +60,4 @@ def get_pandas_df(self) -> pd.DataFrame:
                     "loading into pandas DataFrame."
                 )
 
-            # Notify observer since memory usage has changed
-            for observer in self.dataset_observers:
-                observer.update_memory_usage()
-
         return self.df
diff --git a/server/lomas_server/data_connector/s3_connector.py b/server/lomas_server/data_connector/s3_connector.py
index da60d230..14f57770 100644
--- a/server/lomas_server/data_connector/s3_connector.py
+++ b/server/lomas_server/data_connector/s3_connector.py
@@ -3,10 +3,10 @@
 import boto3
 import pandas as pd
 
-from data_connector.data_connector import DataConnector
-from utils.collection_models import Metadata
-from utils.config import S3CredentialsConfig
-from utils.error_handler import InternalServerException
+from lomas_server.data_connector.data_connector import DataConnector
+from lomas_server.utils.collection_models import Metadata
+from lomas_server.utils.config import S3CredentialsConfig
+from lomas_server.utils.error_handler import InternalServerException
 
 
 class S3Connector(DataConnector):
@@ -56,8 +56,4 @@ def get_pandas_df(self) -> pd.DataFrame:
                     + f"{self.bucket}/{self.key}: {err}"
                 ) from err
 
-            # Notify observer since memory usage has changed
-            for observer in self.dataset_observers:
-                observer.update_memory_usage()
-
         return self.df
diff --git a/server/lomas_server/dataset_store/__init__.py b/server/lomas_server/dataset_store/__init__.py
deleted file mode 100644
index e69de29b..00000000
diff --git a/server/lomas_server/dataset_store/basic_dataset_store.py b/server/lomas_server/dataset_store/basic_dataset_store.py
deleted file mode 100644
index 28de7da0..00000000
--- a/server/lomas_server/dataset_store/basic_dataset_store.py
+++ /dev/null
@@ -1,77 +0,0 @@
-from typing import Dict, List
-
-from admin_database.admin_database import AdminDatabase
-from constants import DPLibraries
-from data_connector.factory import data_connector_factory
-from dataset_store.dataset_store import DatasetStore
-from dp_queries.dp_libraries.factory import querier_factory
-from dp_queries.dp_querier import DPQuerier
-from utils.config import PrivateDBCredentials
-
-
-class BasicDatasetStore(DatasetStore):
-    """
-    Basic implementation of the QuerierManager interface.
-
-    The queriers are initialized lazily and put into a dict.
-    There is no memory management => The manager will fail if the datasets are
-    too large to fit in memory.
-    """
-
-    dp_queriers: Dict[str, Dict[str, DPQuerier]] = {}
-
-    def __init__(
-        self,
-        admin_database: AdminDatabase,
-        private_db_credentials: List[PrivateDBCredentials],
-    ) -> None:
-        """Initializer.
-
-        Args:
-            admin_database (AdminDatabase): An initialized AdminDatabase.
-            private_db_credentials (List[PrivateDBCredentials]):\
-                The private database credentials from the server config.
-        """
-        super().__init__(admin_database, private_db_credentials)
-        self.dp_queriers = {}
-
-    def _add_dataset(self, dataset_name: str) -> None:
-        """Adds a dataset to the manager.
-
-        Args:
-            dataset_name (str): The name of the dataset.
-        """
-        # Should not call this function if dataset already present.
-        assert (
-            dataset_name not in self.dp_queriers
-        ), "BasicQuerierManager: \
-        Trying to add a dataset already in self.dp_queriers"
-
-        # Metadata and data getter
-        data_connector = data_connector_factory(
-            dataset_name, self.admin_database, self.private_db_credentials
-        )
-
-        # Initialize dict
-        self.dp_queriers[dataset_name] = {}
-
-        for lib in DPLibraries:
-            querier = querier_factory(lib.value, data_connector)
-            self.dp_queriers[dataset_name][lib.value] = querier
-
-    def get_querier(self, dataset_name: str, query_type: str) -> DPQuerier:
-        """
-        Returns the querier for the given dataset and library
-
-        Args:
-            dataset_name (str): The dataset name.
-            query_type (str): The type of DP library.
-                One of :py:class:`constants.DPLibraries`
-
-        Returns:
-            DPQuerier: The DPQuerier for the specified dataset and library.
-        """
-        if dataset_name not in self.dp_queriers:
-            self._add_dataset(dataset_name)
-
-        return self.dp_queriers[dataset_name][query_type]
diff --git a/server/lomas_server/dataset_store/data_connector_observer.py b/server/lomas_server/dataset_store/data_connector_observer.py
deleted file mode 100644
index 2323a103..00000000
--- a/server/lomas_server/dataset_store/data_connector_observer.py
+++ /dev/null
@@ -1,14 +0,0 @@
-from abc import ABC, abstractmethod
-
-
-class DataConnectorObserver(ABC):
-    """
-    We use this abstract class to "subscribe" to object instances
-    (DataConnector) so that they can notify instances of this
-    abstract class (LRUDatasetStore or other DatasetStore implementing
-    caching) when their memory usage changes.
-    """
-
-    @abstractmethod
-    def update_memory_usage(self) -> None:
-        """Abstract method to update total memory used by datasets"""
diff --git a/server/lomas_server/dataset_store/dataset_store.py b/server/lomas_server/dataset_store/dataset_store.py
deleted file mode 100644
index 14b5eb13..00000000
--- a/server/lomas_server/dataset_store/dataset_store.py
+++ /dev/null
@@ -1,56 +0,0 @@
-from abc import ABC, abstractmethod
-from typing import List
-
-from admin_database.admin_database import AdminDatabase
-from dp_queries.dp_querier import DPQuerier
-from utils.config import PrivateDBCredentials
-
-
-class DatasetStore(ABC):
-    """
-    Manages the DPQueriers for the different datasets and libraries
-
-    Holds a reference to the user database in order to get information
-    about users.
-
-    We make the _add_dataset function private to enforce lazy loading
-    of queriers.
-    """
-
-    admin_database: AdminDatabase
-
-    def __init__(
-        self,
-        admin_database: AdminDatabase,
-        private_db_credentials: List[PrivateDBCredentials],
-    ) -> None:
-        """Initializer
-
-        Args:
-            admin_database (AdminDatabase): The admin database for the server.
-            private_db_credentials (List[PrivateDBCredentials]):\
-                The private database credentials from the server config.
-        """
-        self.admin_database = admin_database
-        self.private_db_credentials = private_db_credentials
-
-    @abstractmethod
-    def _add_dataset(self, dataset_name: str) -> None:
-        """Adds a dataset to the manager
-
-        Args:
-            dataset_name (str): The dataset name.
-        """
-
-    @abstractmethod
-    def get_querier(self, dataset_name: str, library: str) -> DPQuerier:
-        """Returns the querier for the given dataset and library
-
-        Args:
-            dataset_name (str): The dataset name.
-            library (str): The type of DP library.
-                One of :py:class:`constants.DPLibraries`
-
-        Returns:
-            DPQuerier: The DPQuerier for the specified dataset and library.
-        """
diff --git a/server/lomas_server/dataset_store/factory.py b/server/lomas_server/dataset_store/factory.py
deleted file mode 100644
index ec3e6e71..00000000
--- a/server/lomas_server/dataset_store/factory.py
+++ /dev/null
@@ -1,45 +0,0 @@
-from typing import List
-
-from admin_database.admin_database import AdminDatabase
-from constants import DatasetStoreType
-from dataset_store.basic_dataset_store import BasicDatasetStore
-from dataset_store.dataset_store import DatasetStore
-from dataset_store.lru_dataset_store import LRUDatasetStore
-from utils.config import DatasetStoreConfig, PrivateDBCredentials
-from utils.error_handler import InternalServerException
-
-
-def dataset_store_factory(
-    config: DatasetStoreConfig,
-    admin_database: AdminDatabase,
-    private_db_credentials: List[PrivateDBCredentials],
-) -> DatasetStore:
-    """
-    Instantiates and returns the correct DatasetStore based on the config.
-
-    Args:
-        config (DatasetStoreConfig): A valid DatasetStoreConfig.
-        admin_database (AdminDatabase): An initialized AdminDatabase instance.
-        private_db_credentials (List[PrivateDBCredentials]):\
-             A list of private database credentials.
-
-    Raises:
-        InternalServerException: If the dataset store type from the config
-            does not exist.
-
-    Returns:
-        DatasetStore: The correct DatasetStore instance.
-    """
-    ds_store_type = config.ds_store_type
-
-    match config.ds_store_type:
-        case DatasetStoreType.BASIC:
-            return BasicDatasetStore(admin_database, private_db_credentials)
-        case DatasetStoreType.LRU:
-            return LRUDatasetStore(
-                admin_database, private_db_credentials, config.max_memory_usage
-            )
-        case _:
-            raise InternalServerException(
-                f"Dataset Store type {ds_store_type} does not exist."
-            )
diff --git a/server/lomas_server/dataset_store/lru_dataset_store.py b/server/lomas_server/dataset_store/lru_dataset_store.py
deleted file mode 100644
index 74f94e58..00000000
--- a/server/lomas_server/dataset_store/lru_dataset_store.py
+++ /dev/null
@@ -1,120 +0,0 @@
-from collections import OrderedDict
-from typing import List
-
-from admin_database.admin_database import AdminDatabase
-from data_connector.data_connector import DataConnector
-from data_connector.factory import data_connector_factory
-from dataset_store.data_connector_observer import DataConnectorObserver
-from dataset_store.dataset_store import DatasetStore
-from dp_queries.dp_libraries.factory import querier_factory
-from dp_queries.dp_querier import DPQuerier
-from utils.config import PrivateDBCredentials
-from utils.error_handler import InternalServerException
-from utils.logger import LOG
-
-
-class LRUDatasetStore(DatasetStore, DataConnectorObserver):
-    """
-    Implementation of the DatasetStore interface, with an LRU cache.
-
-    Subscribes to the DataConnectors to get notified if their memory usage
-    changes and then clears the cache accordingly in order stay below
-    the maximum memory usage.
-    """
-
-    dataset_cache: OrderedDict[str, DataConnector]
-
-    def __init__(
-        self,
-        admin_database: AdminDatabase,
-        private_db_credentials: List[PrivateDBCredentials],
-        max_memory_usage: int = 1024,
-    ) -> None:
-        """Initializer.
-
-        Args:
-            admin_database (AdminDatabase): An initialized AdminDatabase.
-            max_memory_usage (int, optional): Maximum memory usage limit\
-                for the manager.. Defaults to 1024.
-            private_db_credentials (List[PrivateDBCredentials]):\
-                The private database credentials from the server config.
-        """
-        super().__init__(admin_database, private_db_credentials)
-        self.max_memory_usage = max_memory_usage
-
-        self.dataset_cache = OrderedDict()
-        self.memory_usage = 0
-
-    def _add_dataset(self, dataset_name: str) -> None:
-        """Adds a dataset to the manager.
-
-        Makes sure the memory usage limit is not exceeded.
-
-        Args:
-            dataset_name (str): The name of the dataset.
-        """
-        # Should not call this function if dataset already present.
-        assert (
-            dataset_name not in self.dataset_cache.keys()
-        ), "BasicQuerierManager: \
-        Trying to add a dataset already in self.dp_queriers"
-
-        # Make private dataset
-        data_connector = data_connector_factory(
-            dataset_name, self.admin_database, self.private_db_credentials
-        )
-        data_connector.subscribe_for_memory_usage_updates(self)
-
-        # Remove least recently used dataset from cache if not enough space
-        data_connector_mem_usage = data_connector.get_memory_usage()
-
-        if data_connector_mem_usage > self.max_memory_usage:
-            raise InternalServerException(
-                f"Dataset {dataset_name} too large"
-                "to fit in dataset manager memory."
-            )
-
-        self.dataset_cache[dataset_name] = data_connector
-        self.memory_usage += data_connector_mem_usage
-
-        LOG.info(f"New dataset cache size: {self.memory_usage} MiB")
-        self.update_memory_usage()
-
-    def update_memory_usage(self) -> None:
-        """Remove least recently used datasets until the cache
-        is back to below or equal to its maximum size.
-        """
-        self.memory_usage = sum(
-            private_ds.get_memory_usage()
-            for private_ds in self.dataset_cache.values()
-        )
-
-        while self.memory_usage > self.max_memory_usage:
-            evicted_ds_name, evicted_ds = self.dataset_cache.popitem(
-                last=False
-            )
-            self.memory_usage -= evicted_ds.get_memory_usage()
-            LOG.info(f"Dataset {evicted_ds_name} was evicted from cache.")
-
-        LOG.info(f"New dataset cache size: {self.memory_usage} MiB")
-
-    def get_querier(self, dataset_name: str, library: str) -> DPQuerier:
-        """Returns the querier for the given dataset and library
-
-        Args:
-            dataset_name (str): The dataset name.
-            library (str): The type of DP library.
-                One of :py:class:`constants.DPLibraries`
-
-        Returns:
-            DPQuerier: The DPQuerier for the specified dataset and library.
-        """
-        # Add dataset to cache if not present and get it.
-        if dataset_name not in self.dataset_cache:
-            self._add_dataset(dataset_name)
-        else:
-            self.dataset_cache.move_to_end(dataset_name)
-        assert dataset_name in self.dataset_cache.keys()
-
-        data_connector = self.dataset_cache[dataset_name]
-        return querier_factory(library, data_connector)
diff --git a/server/lomas_server/dp_queries/dp_libraries/diffprivlib.py b/server/lomas_server/dp_queries/dp_libraries/diffprivlib.py
index a240f14f..a2c1ef8f 100644
--- a/server/lomas_server/dp_queries/dp_libraries/diffprivlib.py
+++ b/server/lomas_server/dp_queries/dp_libraries/diffprivlib.py
@@ -7,30 +7,43 @@
 from sklearn.model_selection import train_test_split
 from sklearn.pipeline import Pipeline
 
-from constants import DPLibraries
-from data_connector.data_connector import DataConnector
-from dp_queries.dp_libraries.utils import handle_missing_data, serialise_model
-from dp_queries.dp_querier import DPQuerier
-from utils.error_handler import (
+from lomas_server.admin_database.admin_database import AdminDatabase
+from lomas_server.constants import DPLibraries
+from lomas_server.data_connector.data_connector import DataConnector
+from lomas_server.dp_queries.dp_libraries.utils import (
+    handle_missing_data,
+    serialise_model,
+)
+from lomas_server.dp_queries.dp_querier import DPQuerier
+from lomas_server.utils.error_handler import (
     ExternalLibraryException,
     InternalServerException,
 )
-from utils.query_models import DiffPrivLibModel
+from lomas_server.utils.query_models import (
+    DiffPrivLibQueryModel,
+    DiffPrivLibRequestModel,
+)
 
 
-class DiffPrivLibQuerier(DPQuerier):
+class DiffPrivLibQuerier(
+    DPQuerier[DiffPrivLibRequestModel, DiffPrivLibQueryModel]
+):
     """
     Concrete implementation of the DPQuerier ABC for the DiffPrivLib library.
     """
 
-    def __init__(self, data_connector: DataConnector) -> None:
-        super().__init__(data_connector)
+    def __init__(
+        self,
+        data_connector: DataConnector,
+        admin_database: AdminDatabase,
+    ) -> None:
+        super().__init__(data_connector, admin_database)
         self.dpl_pipeline: Optional[Pipeline] = None
         self.x_test: Optional[pd.DataFrame] = None
         self.y_test: Optional[pd.DataFrame] = None
 
     def fit_model_on_data(
-        self, query_json: DiffPrivLibModel
+        self, query_json: DiffPrivLibRequestModel
     ) -> tuple[Pipeline, pd.DataFrame, pd.DataFrame]:
         """Perform necessary steps to fit the model on the data
 
@@ -77,11 +90,11 @@ def fit_model_on_data(
 
         return dpl_pipeline, x_test, y_test
 
-    def cost(self, query_json: DiffPrivLibModel) -> tuple[float, float]:
+    def cost(self, query_json: DiffPrivLibRequestModel) -> tuple[float, float]:
         """Estimate cost of query
 
         Args:
-            query_json (BaseModel): The JSON request object for the query.
+            query_json (DiffPrivLibRequestModel): The request model object.
 
         Raises:
             ExternalLibraryException: For exceptions from libraries
@@ -104,12 +117,13 @@ def cost(self, query_json: DiffPrivLibModel) -> tuple[float, float]:
         return spent_epsilon, spent_delta
 
     def query(
-        self, query_json: DiffPrivLibModel  # pylint: disable=unused-argument
+        self,
+        query_json: DiffPrivLibQueryModel,  # pylint: disable=unused-argument
     ) -> Dict:
         """Perform the query and return the response.
 
         Args:
-            query_json (BaseModel): The JSON request object for the query.
+            query_json (DiffPrivLibQueryModel): The request model object.
 
         Raises:
             ExternalLibraryException: For exceptions from libraries
@@ -137,12 +151,12 @@ def query(
 
 
 def split_train_test_data(
-    df: pd.DataFrame, query_json: DiffPrivLibModel
+    df: pd.DataFrame, query_json: DiffPrivLibRequestModel
 ) -> tuple[pd.DataFrame, pd.DataFrame, pd.DataFrame, pd.DataFrame]:
     """Split the data between train and test set
     Args:
         df (pd.DataFrame): dataframe with the data
-        query_json (DiffPrivLibModel): user input query indication
+        query_json (DiffPrivLibRequestModel): user input query indication
             feature_columns (list[str]): columns from data to use as features
             target_columns (list[str]): columns from data to use as target (to predict)
             test_size (float): proportion of data in the test set
diff --git a/server/lomas_server/dp_queries/dp_libraries/factory.py b/server/lomas_server/dp_queries/dp_libraries/factory.py
index 4eb1c2bb..776dd280 100644
--- a/server/lomas_server/dp_queries/dp_libraries/factory.py
+++ b/server/lomas_server/dp_queries/dp_libraries/factory.py
@@ -1,20 +1,31 @@
-from constants import DPLibraries
-from data_connector.data_connector import DataConnector
-from dp_queries.dp_libraries.diffprivlib import DiffPrivLibQuerier
-from dp_queries.dp_libraries.opendp import OpenDPQuerier
-from dp_queries.dp_libraries.smartnoise_sql import SmartnoiseSQLQuerier
-from dp_queries.dp_libraries.smartnoise_synth import SmartnoiseSynthQuerier
-from dp_queries.dp_querier import DPQuerier
-from utils.error_handler import InternalServerException
-
-
-def querier_factory(lib: str, data_connector: DataConnector) -> DPQuerier:
+from lomas_server.admin_database.admin_database import AdminDatabase
+from lomas_server.constants import DPLibraries
+from lomas_server.data_connector.data_connector import DataConnector
+from lomas_server.dp_queries.dp_libraries.diffprivlib import DiffPrivLibQuerier
+from lomas_server.dp_queries.dp_libraries.opendp import OpenDPQuerier
+from lomas_server.dp_queries.dp_libraries.smartnoise_sql import (
+    SmartnoiseSQLQuerier,
+)
+from lomas_server.dp_queries.dp_libraries.smartnoise_synth import (
+    SmartnoiseSynthQuerier,
+)
+from lomas_server.dp_queries.dp_querier import DPQuerier
+from lomas_server.utils.error_handler import InternalServerException
+
+
+def querier_factory(
+    lib: str,
+    data_connector: DataConnector,
+    admin_database: AdminDatabase,
+) -> DPQuerier:
     """Builds the correct DPQuerier instance.
 
     Args:
         lib (str): The library to build the querier for.
             One of :py:class:`DPLibraries`.
         data_connector (DataConnector): The dataset to query.
+        admin_database (AdminDatabase): An initialized instance of
+                an AdminDatabase.
 
     Raises:
         InternalServerException: If the library is unknown.
@@ -22,18 +33,19 @@ def querier_factory(lib: str, data_connector: DataConnector) -> DPQuerier:
     Returns:
         DPQuerier: The built DPQuerier.
     """
+    querier: DPQuerier
     match lib:
         case DPLibraries.SMARTNOISE_SQL:
-            querier = SmartnoiseSQLQuerier(data_connector)
+            querier = SmartnoiseSQLQuerier(data_connector, admin_database)
 
         case DPLibraries.SMARTNOISE_SYNTH:
-            querier = SmartnoiseSynthQuerier(data_connector)
+            querier = SmartnoiseSynthQuerier(data_connector, admin_database)
 
         case DPLibraries.OPENDP:
-            querier = OpenDPQuerier(data_connector)
+            querier = OpenDPQuerier(data_connector, admin_database)
 
         case DPLibraries.DIFFPRIVLIB:
-            querier = DiffPrivLibQuerier(data_connector)
+            querier = DiffPrivLibQuerier(data_connector, admin_database)
 
         case _:
             raise InternalServerException(f"Unknown library: {lib}")
diff --git a/server/lomas_server/dp_queries/dp_libraries/opendp.py b/server/lomas_server/dp_queries/dp_libraries/opendp.py
index 209949cf..6690f311 100644
--- a/server/lomas_server/dp_queries/dp_libraries/opendp.py
+++ b/server/lomas_server/dp_queries/dp_libraries/opendp.py
@@ -5,29 +5,36 @@
 from opendp.mod import enable_features
 from opendp_logger import make_load_json
 
-from constants import DPLibraries, OpenDPDatasetInputMetric, OpenDPMeasurement
-from dp_queries.dp_querier import DPQuerier
-from utils.config import OpenDPConfig
-from utils.error_handler import (
+from lomas_server.constants import (
+    DPLibraries,
+    OpenDPDatasetInputMetric,
+    OpenDPMeasurement,
+)
+from lomas_server.dp_queries.dp_querier import DPQuerier
+from lomas_server.utils.config import OpenDPConfig
+from lomas_server.utils.error_handler import (
     ExternalLibraryException,
     InternalServerException,
     InvalidQueryException,
 )
-from utils.logger import LOG
-from utils.query_models import OpenDPModel
+from lomas_server.utils.logger import LOG
+from lomas_server.utils.query_models import (
+    OpenDPQueryModel,
+    OpenDPRequestModel,
+)
 
 
-class OpenDPQuerier(DPQuerier):
+class OpenDPQuerier(DPQuerier[OpenDPRequestModel, OpenDPQueryModel]):
     """
     Concrete implementation of the DPQuerier ABC for the OpenDP library.
     """
 
-    def cost(self, query_json: OpenDPModel) -> tuple[float, float]:
+    def cost(self, query_json: OpenDPRequestModel) -> tuple[float, float]:
         """
         Estimate cost of query
 
         Args:
-            query_json (BaseModel): The JSON request object for the query.
+            query_json (OpenDPRequestModel): The request model object.
 
         Raises:
             ExternalLibraryException: For exceptions from libraries
@@ -49,7 +56,7 @@ def cost(self, query_json: OpenDPModel) -> tuple[float, float]:
             opendp_pipe = dp.combinators.make_zCDP_to_approxDP(opendp_pipe)
             measurement_type = OpenDPMeasurement.SMOOTHED_MAX_DIVERGENCE
 
-        max_ids = self.data_connector.get_metadata()["max_ids"]
+        max_ids = self.data_connector.get_metadata().max_ids
         try:
             # d_in is int as input metric is a dataset metric
             cost = opendp_pipe.map(d_in=int(max_ids))
@@ -82,11 +89,11 @@ def cost(self, query_json: OpenDPModel) -> tuple[float, float]:
 
         return epsilon, delta
 
-    def query(self, query_json: OpenDPModel) -> Union[List, int, float]:
+    def query(self, query_json: OpenDPQueryModel) -> Union[List, int, float]:
         """Perform the query and return the response.
 
         Args:
-            query_json (BaseModel): The JSON request object for the query.
+            query_json (OpenDPQueryModel): The input model for the query.
 
         Raises:
             ExternalLibraryException: For exceptions from libraries
diff --git a/server/lomas_server/dp_queries/dp_libraries/smartnoise_sql.py b/server/lomas_server/dp_queries/dp_libraries/smartnoise_sql.py
index 41114c4f..d393ef4e 100644
--- a/server/lomas_server/dp_queries/dp_libraries/smartnoise_sql.py
+++ b/server/lomas_server/dp_queries/dp_libraries/smartnoise_sql.py
@@ -4,28 +4,40 @@
 from snsql import Mechanism, Privacy, Stat, from_connection
 from snsql.reader.base import Reader
 
-from constants import SSQL_MAX_ITERATION, SSQL_STATS, DPLibraries
-from data_connector.data_connector import DataConnector
-from dp_queries.dp_querier import DPQuerier
-from utils.collection_models import Metadata
-from utils.error_handler import (
+from lomas_server.admin_database.admin_database import AdminDatabase
+from lomas_server.constants import SSQL_MAX_ITERATION, SSQL_STATS, DPLibraries
+from lomas_server.data_connector.data_connector import DataConnector
+from lomas_server.dp_queries.dp_querier import DPQuerier
+from lomas_server.utils.collection_models import Metadata
+from lomas_server.utils.error_handler import (
     ExternalLibraryException,
     InternalServerException,
     InvalidQueryException,
 )
-from utils.query_models import SmartnoiseSQLCostModel, SmartnoiseSQLModel
+from lomas_server.utils.query_models import (
+    SmartnoiseSQLQueryModel,
+    SmartnoiseSQLRequestModel,
+)
 
 
-class SmartnoiseSQLQuerier(DPQuerier):
+class SmartnoiseSQLQuerier(
+    DPQuerier[SmartnoiseSQLRequestModel, SmartnoiseSQLQueryModel]
+):
     """
     Concrete implementation of the DPQuerier ABC for the SmartNoiseSQL library.
     """
 
-    def __init__(self, data_connector: DataConnector) -> None:
-        super().__init__(data_connector)
+    def __init__(
+        self,
+        data_connector: DataConnector,
+        admin_database: AdminDatabase,
+    ) -> None:
+        super().__init__(data_connector, admin_database)
         self.reader: Optional[Reader] = None
 
-    def cost(self, query_json: SmartnoiseSQLCostModel) -> tuple[float, float]:
+    def cost(
+        self, query_json: SmartnoiseSQLRequestModel
+    ) -> tuple[float, float]:
         """Estimate cost of query
 
         Args:
@@ -61,11 +73,24 @@ def cost(self, query_json: SmartnoiseSQLCostModel) -> tuple[float, float]:
 
         return epsilon, delta
 
-    def query(self, query_json: SmartnoiseSQLModel, nb_iter: int = 0) -> dict:
+    def query(self, query_json: SmartnoiseSQLQueryModel) -> dict:
+        """Performs the query and returns the response.
+
+        Args:
+            query_json (SmartnoiseSQLQueryModel): The request model object.
+
+        Returns:
+            dict: The dictionary encoding of the result pd.DataFrame.
+        """
+        return self.query_with_iter(query_json)
+
+    def query_with_iter(
+        self, query_json: SmartnoiseSQLQueryModel, nb_iter: int = 0
+    ) -> dict:
         """Perform the query and return the response.
 
         Args:
-            query_json (SmartnoiseSQLModel): JSON request object for the query.
+            query_json (SmartnoiseSQLQueryModel): Request object for the query.
             nb_iter (int, optional): Number of trials if output is Nan.
                 Defaults to 0.
 
@@ -114,7 +139,7 @@ def query(self, query_json: SmartnoiseSQLModel, nb_iter: int = 0) -> dict:
             # Try again up to SSQL_MAX_ITERATION
             if nb_iter < SSQL_MAX_ITERATION:
                 nb_iter += 1
-                return self.query(query_json, nb_iter)
+                return self.query_with_iter(query_json, nb_iter)
 
             raise InvalidQueryException(
                 f"SQL Reader generated NAN results. "
@@ -151,7 +176,7 @@ def convert_to_smartnoise_metadata(metadata: Metadata) -> dict:
     Returns:
         dict: metadata of the dataset in smartnoise-sql format
     """
-    metadata = dict(metadata)
-    metadata.update(metadata["columns"])
-    del metadata["columns"]
-    return {"": {"": {"df": metadata}}}
+    metadata_dict = metadata.model_dump()
+    metadata_dict.update(metadata_dict["columns"])
+    del metadata_dict["columns"]
+    return {"": {"": {"df": metadata_dict}}}
diff --git a/server/lomas_server/dp_queries/dp_libraries/smartnoise_synth.py b/server/lomas_server/dp_queries/dp_libraries/smartnoise_synth.py
index 4ac6e2a7..7bc238a5 100644
--- a/server/lomas_server/dp_queries/dp_libraries/smartnoise_synth.py
+++ b/server/lomas_server/dp_queries/dp_libraries/smartnoise_synth.py
@@ -1,6 +1,6 @@
 import re
 from datetime import datetime
-from typing import Dict, List, Optional, Union
+from typing import Dict, List, Optional, TypeAlias, TypeGuard, Union
 
 import pandas as pd
 from smartnoise_synth_logger import deserialise_constraints
@@ -16,122 +16,178 @@
 from snsynth.transform.datetime import DateTimeTransformer
 from snsynth.transform.table import TableTransformer
 
-from constants import (
-    DEFAULT_DATE_FORMAT,
+from lomas_server.admin_database.admin_database import AdminDatabase
+from lomas_server.constants import (
     SECONDS_IN_A_DAY,
     SSYNTH_DEFAULT_BINS,
     SSYNTH_MIN_ROWS_PATE_GAN,
     SSYNTH_PRIVATE_COLUMN,
     DPLibraries,
-    SSynthColumnType,
     SSynthGanSynthesizer,
     SSynthMarginalSynthesizer,
     SSynthTableTransStyle,
 )
-from data_connector.data_connector import DataConnector
-from dp_queries.dp_libraries.utils import serialise_model
-from dp_queries.dp_querier import DPQuerier
-from utils.collection_models import Metadata
-from utils.error_handler import (
+from lomas_server.data_connector.data_connector import DataConnector
+from lomas_server.dp_queries.dp_libraries.utils import serialise_model
+from lomas_server.dp_queries.dp_querier import DPQuerier
+from lomas_server.utils.collection_models import (
+    BooleanMetadata,
+    ColumnMetadata,
+    DatetimeMetadata,
+    FloatMetadata,
+    IntCategoricalMetadata,
+    IntMetadata,
+    Metadata,
+    StrCategoricalMetadata,
+    StrMetadata,
+)
+from lomas_server.utils.error_handler import (
     ExternalLibraryException,
     InternalServerException,
     InvalidQueryException,
 )
-from utils.query_models import (
-    SmartnoiseSynthCostModel,
+from lomas_server.utils.query_models import (
     SmartnoiseSynthQueryModel,
+    SmartnoiseSynthRequestModel,
 )
 
 
-def datetime_to_float(upper, lower):
+def datetime_to_float(upper: datetime, lower: datetime) -> float:
     """Convert the upper date as the distance between the upper date and
     lower date as float
 
     Args:
-            upper (str): date to convert
-            lower (str): start date to convert from
+            upper (datetime): date to convert
+            lower (datetime): start date to convert from
 
         Returns:
             float: number of days between upper and lower
     """
-    distance = datetime.strptime(
-        upper, DEFAULT_DATE_FORMAT
-    ) - datetime.strptime(lower, DEFAULT_DATE_FORMAT)
+    distance = upper - lower
     return float(distance.total_seconds() / SECONDS_IN_A_DAY)
 
 
-class SmartnoiseSynthQuerier(DPQuerier):
+# TODO maybe a better place to put this? See issue #336
+SSynthColumnType: TypeAlias = Union[
+    StrMetadata,
+    StrCategoricalMetadata,
+    BooleanMetadata,
+    IntCategoricalMetadata,
+    IntMetadata,
+    FloatMetadata,
+    DatetimeMetadata,
+]
+
+
+class SmartnoiseSynthQuerier(
+    DPQuerier[SmartnoiseSynthRequestModel, SmartnoiseSynthQueryModel]
+):
     """
     Concrete implementation of the DPQuerier ABC for the SmartNoiseSynth library.
     """
 
-    def __init__(self, data_connector: DataConnector) -> None:
-        super().__init__(data_connector)
+    def __init__(
+        self,
+        data_connector: DataConnector,
+        admin_database: AdminDatabase,
+    ) -> None:
+        super().__init__(data_connector, admin_database)
         self.model: Optional[Synthesizer] = None
 
-    def _categorize_column(self, data: dict) -> str:
+    def _is_categorical(
+        self, col_metadata: ColumnMetadata
+    ) -> TypeGuard[
+        StrMetadata
+        | StrCategoricalMetadata
+        | BooleanMetadata
+        | IntCategoricalMetadata
+    ]:
         """
-        Categorize the column based on its metadata.
+        Checks if the column type is categorical
 
         Args:
-            data (dict): Metadata of the column.
+            col_metadata (ColumnMetadata): The column metadata
 
         Returns:
-            str: Category of the column.
+            TypeGuard[StrMetadata | StrCategoricalMetadata| BooleanMetadata |
+                      IntCategoricalMetadata]:
+                TypeGuard for categorical columns metadata
         """
-        match data["type"]:
-            case "string" | "boolean":
-                return SSynthColumnType.CATEGORICAL
-            case "int" | "float":
-                if "lower" in data.keys():
-                    return SSynthColumnType.CONTINUOUS
-                return SSynthColumnType.CATEGORICAL
-            case "datetime":
-                return SSynthColumnType.DATETIME
-            case _:
-                raise InternalServerException(
-                    f"Unknown column type in metadata: {data['type']}"
-                )
+        return isinstance(
+            col_metadata,
+            (
+                StrMetadata,
+                StrCategoricalMetadata,
+                BooleanMetadata,
+                IntCategoricalMetadata,
+            ),
+        )
 
-    def _get_column_by_types(
-        self,
-        metadata: Metadata,
-        select_cols: List[str],
-    ) -> Dict[str, List[str]]:
+    def _is_continuous(
+        self, col_metadata: ColumnMetadata
+    ) -> TypeGuard[IntMetadata | FloatMetadata]:
+        """Checks if the column type is continuous
+
+        Args:
+            col_metadata (ColumnMetadata): The column metadata
+
+        Returns:
+            TypeGuard[IntMetadata | FloatMetadata]:
+                TypeGuard for continuous columns metadata
         """
-        Sort the column in categories based on their types and metadata
+        return isinstance(col_metadata, (IntMetadata, FloatMetadata))
+
+    def _is_datetime(
+        self, col_metadata: ColumnMetadata
+    ) -> TypeGuard[DatetimeMetadata]:
+        """Checks if the column type is datetime
 
         Args:
-            metadata (Metadata): Metadata of the dataset
-            select_cols (List[str]): List of columns to select
+            col_metadata (ColumnMetadata): The column metadata
 
         Returns:
-            Dict[str, List[str]]: Dictionnary of list of columns by categories
+            TypeGuard[DatetimeMetadata]: TypeGuard for datetime metadata.
+        """
+        return isinstance(col_metadata, DatetimeMetadata)
+
+    def _get_and_check_valid_column_types(
+        self, metadata: Metadata, select_cols: List[str]
+    ) -> Dict[str, SSynthColumnType]:
         """
-        col_categories: Dict[str, List[str]] = {
-            SSynthColumnType.CATEGORICAL: [],
-            SSynthColumnType.CONTINUOUS: [],
-            SSynthColumnType.DATETIME: [],
-            SSynthColumnType.PRIVATE_ID: [],
-        }
-        for col_name, data in metadata["columns"].items():
+        Ensures the type of the selected columns can be handled with
+        SmartnoiseSynth and returns the dict of column metadata
+        for the selected columns.
+
+        Args:
+            metadata (Metadata): Dataset metadata
+            select_cols (List[str]): List of selected columns
+
+        Raises:
+            InternalServerException: If one of the column types
+                cannot be handled with SmartnoiseSynth.
+
+        Returns:
+            Dict[str, SSynthColumnType]: The filtered dict of selected columns.
+        """
+        columns: Dict[str, SSynthColumnType] = {}
+
+        for col_name, data in metadata.columns.items():
             if select_cols and col_name not in select_cols:
                 continue
 
-            if "private_id" in data.keys():
-                col_categories[SSynthColumnType.PRIVATE_ID].append(col_name)
-                continue
+            if not isinstance(data, SSynthColumnType):  # type: ignore[misc, arg-type]
+                raise InternalServerException(
+                    f"Column type {data.type} not supported for SmartnoiseSynth"
+                )
 
-            # Sort the column in categories based on their types and metadata
-            category = self._categorize_column(data)
-            col_categories[category].append(col_name)
+            columns[col_name] = data
 
-        return col_categories
+        return columns
 
     def _get_default_constraints(
         self,
         metadata: Metadata,
-        query_json: dict,
+        query_json: SmartnoiseSynthRequestModel,
         table_transformer_style: str,
     ) -> TableTransformer:
         """
@@ -142,7 +198,7 @@ def _get_default_constraints(
 
         Args:
             metadata (Metadata): Metadata of the dataset
-            query_json (SmartnoiseSynthModelCost): JSON request object for the query
+            query_json (SmartnoiseSynthRequestModel): JSON request object for the query
                 select_cols (List[str]): List of columns to select
                 nullable (bool): True is the data can have Null values, False otherwise
             table_transformer_style (str): 'gan' or 'cube'
@@ -150,69 +206,72 @@ def _get_default_constraints(
         Returns:
             table_tranformer (TableTransformer) to pre and post-process the data
         """
-        col_categories = self._get_column_by_types(
+        columns = self._get_and_check_valid_column_types(
             metadata, query_json.select_cols
         )
 
         constraints = {}
-        for col in col_categories[SSynthColumnType.PRIVATE_ID]:
-            constraints[col] = AnonymizationTransformer(SSYNTH_PRIVATE_COLUMN)
-
         nullable = query_json.nullable
-        if table_transformer_style == SSynthTableTransStyle.GAN:  # gan
-            for col in col_categories[SSynthColumnType.CATEGORICAL]:
-                constraints[col] = ChainTransformer(
-                    [LabelTransformer(nullable=nullable), OneHotEncoder()]
+        for col, col_metadata in columns.items():
+            if col_metadata.private_id:
+                constraints[col] = AnonymizationTransformer(
+                    SSYNTH_PRIVATE_COLUMN
                 )
-            for col in col_categories[SSynthColumnType.CONTINUOUS]:
-                constraints[col] = MinMaxTransformer(
-                    lower=metadata["columns"][col]["lower"],
-                    upper=metadata["columns"][col]["upper"],
-                    nullable=nullable,
-                )
-            for col in col_categories[SSynthColumnType.DATETIME]:
-                constraints[col] = ChainTransformer(
-                    [
-                        DateTimeTransformer(
-                            epoch=metadata["columns"][col]["lower"]
-                        ),
-                        MinMaxTransformer(
-                            lower=0.0,  # because start epoch at lower bound
-                            upper=datetime_to_float(
-                                metadata["columns"][col]["upper"],
-                                metadata["columns"][col]["lower"],
+
+            if table_transformer_style == SSynthTableTransStyle.GAN:  # gan
+                if self._is_categorical(
+                    col_metadata
+                ):  # TODO any way of specifying cardinality? See issue #337
+                    constraints[col] = ChainTransformer(
+                        [LabelTransformer(nullable=nullable), OneHotEncoder()]
+                    )
+                elif self._is_continuous(col_metadata):
+                    constraints[col] = MinMaxTransformer(
+                        lower=col_metadata.lower,
+                        upper=col_metadata.upper,
+                        nullable=nullable,
+                    )
+                elif self._is_datetime(col_metadata):
+                    constraints[col] = ChainTransformer(
+                        [
+                            DateTimeTransformer(epoch=col_metadata.lower),
+                            MinMaxTransformer(
+                                lower=0.0,  # because start epoch at lower bound
+                                upper=datetime_to_float(
+                                    col_metadata.upper,
+                                    col_metadata.lower,
+                                ),
+                                nullable=nullable,
                             ),
-                            nullable=nullable,
-                        ),
-                    ]
-                )
-        else:  # cube
-            for col in col_categories[SSynthColumnType.CATEGORICAL]:
-                constraints[col] = LabelTransformer(nullable=nullable)
-            for col in col_categories[SSynthColumnType.CONTINUOUS]:
-                constraints[col] = BinTransformer(
-                    lower=metadata["columns"][col]["lower"],
-                    upper=metadata["columns"][col]["upper"],
-                    bins=SSYNTH_DEFAULT_BINS,
-                    nullable=nullable,
-                )
-            for col in col_categories[SSynthColumnType.DATETIME]:
-                constraints[col] = ChainTransformer(
-                    [
-                        DateTimeTransformer(
-                            epoch=metadata["columns"][col]["lower"]
-                        ),
-                        BinTransformer(
-                            lower=0.0,  # because start epoch at lower bound
-                            upper=datetime_to_float(
-                                metadata["columns"][col]["upper"],
-                                metadata["columns"][col]["lower"],
+                        ]
+                    )
+            else:  # Cube
+                if self._is_categorical(
+                    col_metadata
+                ):  # TODO any way of specifying cardinality? See issue #337
+                    constraints[col] = LabelTransformer(nullable=nullable)
+                elif self._is_continuous(col_metadata):
+                    constraints[col] = BinTransformer(
+                        lower=col_metadata.lower,
+                        upper=col_metadata.upper,
+                        bins=SSYNTH_DEFAULT_BINS,
+                        nullable=nullable,
+                    )
+                elif self._is_datetime(col_metadata):
+                    constraints[col] = ChainTransformer(
+                        [
+                            DateTimeTransformer(epoch=col_metadata.lower),
+                            BinTransformer(
+                                lower=0.0,  # because start epoch at lower bound
+                                upper=datetime_to_float(
+                                    col_metadata.upper,
+                                    col_metadata.lower,
+                                ),
+                                bins=SSYNTH_DEFAULT_BINS,
+                                nullable=nullable,
                             ),
-                            bins=SSYNTH_DEFAULT_BINS,
-                            nullable=nullable,
-                        ),
-                    ]
-                )
+                        ]
+                    )
 
         return constraints
 
@@ -220,7 +279,7 @@ def _get_fit_model(
         self,
         private_data: pd.DataFrame,
         transformer: TableTransformer,
-        query_json: dict,
+        query_json: SmartnoiseSynthRequestModel,
     ) -> Synthesizer:
         """
         Create and fit the synthesizer model.
@@ -228,7 +287,7 @@ def _get_fit_model(
         Args:
             private_data (pd.DataFrame): Private data for fitting the model
             transformer (TableTransformer): Transformer to pre/postprocess data
-            query_json (SmartnoiseSynthModelCost): JSON request object for the query
+            query_json (SmartnoiseSynthRequestModel): JSON request object for the query
                 synth_name (str): name of the Yanthesizer model to use
                 epsilon (float): epsilon budget value
                 nullable (bool): True if some data cells may be null
@@ -285,12 +344,12 @@ def _get_fit_model(
         return model
 
     def _model_pipeline(
-        self, query_json: SmartnoiseSynthCostModel
+        self, query_json: SmartnoiseSynthRequestModel
     ) -> Synthesizer:
         """Return a trained Synthesizer model based on query_json
 
         Args:
-            query_json (SmartnoiseSynthCostModel): JSON request object for the query.
+            query_json (SmartnoiseSynthRequestModel): JSON request object for the query.
 
         Returns:
             model: Smartnoise Synthesizer
@@ -320,7 +379,7 @@ def _model_pipeline(
         # Preprocessing information from metadata
         metadata = self.data_connector.get_metadata()
         if query_json.synth_name == SSynthGanSynthesizer.PATE_GAN:
-            if metadata["rows"] < SSYNTH_MIN_ROWS_PATE_GAN:
+            if metadata.rows < SSYNTH_MIN_ROWS_PATE_GAN:
                 raise ExternalLibraryException(
                     DPLibraries.SMARTNOISE_SYNTH,
                     f"{SSynthGanSynthesizer.PATE_GAN} not reliable "
@@ -332,9 +391,9 @@ def _model_pipeline(
         )
 
         # Overwrite default constraint with custom constraint (if any)
-        custom_constraints = query_json.constraints
-        if custom_constraints:
-            custom_constraints = deserialise_constraints(custom_constraints)
+        constraints_json = query_json.constraints
+        if constraints_json:
+            custom_constraints = deserialise_constraints(constraints_json)
             custom_constraints = {
                 key: custom_constraints[key]
                 for key in query_json.select_cols
@@ -364,12 +423,12 @@ def _model_pipeline(
         return model
 
     def cost(
-        self, query_json: SmartnoiseSynthCostModel
+        self, query_json: SmartnoiseSynthRequestModel
     ) -> tuple[float, float]:
         """Return cost of query_json
 
         Args:
-            query_json (SmartnoiseSynthModelCost): JSON request object for the query.
+            query_json (SmartnoiseSynthRequestModel): JSON request object for the query.
 
         Returns:
             tuple[float, float]: The tuple of costs, the first value
@@ -387,12 +446,14 @@ def cost(
         return epsilon, delta
 
     def query(
-        self, query_json: SmartnoiseSynthQueryModel
+        self,
+        query_json: SmartnoiseSynthQueryModel,
     ) -> Union[pd.DataFrame, str]:
         """Perform the query and return the response.
 
         Args:
-            query_json (SmartnoiseSynthModel): The JSON request object for the query.
+            query_json (SmartnoiseSynthQueryModel):
+                The request object for the query.
 
         Raises:
             ExternalLibraryException: For exceptions from libraries
diff --git a/server/lomas_server/dp_queries/dp_libraries/utils.py b/server/lomas_server/dp_queries/dp_libraries/utils.py
index 84e6d75f..9332fa9b 100644
--- a/server/lomas_server/dp_queries/dp_libraries/utils.py
+++ b/server/lomas_server/dp_queries/dp_libraries/utils.py
@@ -6,8 +6,8 @@
 import pandas as pd
 from sklearn.impute import SimpleImputer
 
-from constants import NUMERICAL_DTYPES
-from utils.error_handler import InvalidQueryException
+from lomas_server.constants import NUMERICAL_DTYPES
+from lomas_server.utils.error_handler import InvalidQueryException
 
 
 def handle_missing_data(
diff --git a/server/lomas_server/dp_queries/dp_logic.py b/server/lomas_server/dp_queries/dp_logic.py
deleted file mode 100644
index 0a32e104..00000000
--- a/server/lomas_server/dp_queries/dp_logic.py
+++ /dev/null
@@ -1,204 +0,0 @@
-from pydantic import BaseModel
-
-from admin_database.admin_database import AdminDatabase
-from constants import DPLibraries
-from dataset_store.dataset_store import DatasetStore
-from dp_queries.dp_querier import DPQuerier
-from utils.error_handler import (
-    KNOWN_EXCEPTIONS,
-    InternalServerException,
-    InvalidQueryException,
-    UnauthorizedAccessException,
-)
-
-
-class QueryHandler:
-    """
-    Query handler for the server.
-
-    Holds a reference to the admin database and the DatasetStore.
-    """
-
-    admin_database: AdminDatabase
-    dataset_store: DatasetStore
-
-    def __init__(
-        self, admin_database: AdminDatabase, dataset_store: DatasetStore
-    ) -> None:
-        """Initializer.
-
-        Args:
-            admin_database (AdminDatabase): An initialized instance of
-                an AdminDatabase.
-            dataset_store (DatasetStore): A dataset store.
-        """
-        self.admin_database = admin_database
-        self.dataset_store = dataset_store
-
-    def _get_querier(
-        self,
-        query_type: str,
-        query_json: BaseModel,
-    ) -> DPQuerier:
-        """Internal function to get a correct querier.
-
-        Get the querier for the query_type and dataset_name in
-        the query_json.
-
-        Args:
-            query_type (str): The type of DP library, one of
-            :py:class:DPLibraries`
-            query_json (BasicModel): The JSON request object for the query.
-
-        Raises:
-            InternalServerException: If the query type does not exist.
-            InternalServerException: If the querier cannot be received from
-                the querier manager.
-
-        Returns:
-            DPQuerier: The correct querier instance for the given
-                request and libray.
-        """
-        # Check query type
-        supported_lib = [lib.value for lib in DPLibraries]
-        if query_type not in supported_lib:
-            raise InternalServerException(
-                f"Query type {query_type} not supported in QueryHandler"
-            )
-
-        # Get querier
-        try:
-            dp_querier = self.dataset_store.get_querier(
-                query_json.dataset_name, query_type
-            )
-        except InvalidQueryException as e:
-            raise e
-        except Exception as e:
-            raise InternalServerException(
-                "Failed to get querier for dataset "
-                + f"{query_json.dataset_name}: {str(e)}"
-            ) from e
-        return dp_querier
-
-    def estimate_cost(
-        self,
-        query_type: str,
-        query_json: BaseModel,
-    ) -> dict[str, float]:
-        """Estimate query cost.
-
-        Args:
-            query_type (str): The type of DP library,
-                one of :py:class:DPLibraries`
-            query_json (BasicModel): The JSON request object for the query.
-
-        Returns:
-            dict[str, float]: Dictionary containing:
-                - epsilon_cost (float): The estimated epsilon cost.
-                - delta_cost (float): The estimated delta cost.
-        """
-        # Get querier
-        dp_querier = self._get_querier(query_type, query_json)
-
-        # Get cost of the query
-        eps_cost, delta_cost = dp_querier.cost(query_json)
-
-        return {"epsilon_cost": eps_cost, "delta_cost": delta_cost}
-
-    def handle_query(
-        self,
-        query_type: str,
-        query_json: BaseModel,
-        user_name: str,
-    ) -> dict:
-        """
-        Handle DP query.
-
-        Args:
-            query_type (str): The type of DP library,
-                one of :py:class:DPLibraries`
-            query_json (BasicModel): The JSON request object for the query.
-            user_name (str, optional): User name.
-
-        Raises:
-            UnauthorizedAccessException: A query is already
-                ongoing for this user,
-            the user does not exist or does not have access to the dataset.
-            InvalidQueryException: If the query is not valid.
-            InternalServerException: For any other unforseen exceptions.
-
-        Returns:
-            dict: A dictionary containing:
-                - requested_by (str): The user name.
-                - query_response (pd.DataFrame): A DataFrame containing
-                  the query response.
-                - spent_epsilon (float): The amount of epsilon budget spent
-                for the query.
-                - spent_delta (float): The amount of delta budget spent
-                  for the query.
-
-        """
-        # Block access to other queries to user
-        if not self.admin_database.get_and_set_may_user_query(
-            user_name, False
-        ):
-            raise UnauthorizedAccessException(
-                f"User {user_name} is trying to query"
-                + " before end of previous query."
-            )
-
-        try:
-            # Get querier
-            dp_querier = self._get_querier(query_type, query_json)
-
-            # Get cost of the query
-            eps_cost, delta_cost = dp_querier.cost(query_json)
-
-            # Check that enough budget to do the query
-            try:
-                (
-                    eps_remain,
-                    delta_remain,
-                ) = self.admin_database.get_remaining_budget(
-                    user_name, query_json.dataset_name
-                )
-            except UnauthorizedAccessException as e:
-                raise e
-
-            if (eps_remain < eps_cost) or (delta_remain < delta_cost):
-                raise InvalidQueryException(
-                    "Not enough budget for this query epsilon remaining "
-                    f"{eps_remain}, delta remaining {delta_remain}."
-                )
-
-            # Query
-            try:
-                query_response = dp_querier.query(query_json)
-            except KNOWN_EXCEPTIONS as e:
-                raise e
-            except Exception as e:
-                raise InternalServerException(e) from e
-
-            # Deduce budget from user
-            self.admin_database.update_budget(
-                user_name, query_json.dataset_name, eps_cost, delta_cost
-            )
-            response = {
-                "requested_by": user_name,
-                "query_response": query_response,
-                "spent_epsilon": eps_cost,
-                "spent_delta": delta_cost,
-            }
-
-            # Add query to db (for archive)
-            self.admin_database.save_query(user_name, query_json, response)
-
-        except Exception as e:
-            self.admin_database.set_may_user_query(user_name, True)
-            raise e
-
-        # Re-enable user to query
-        self.admin_database.set_may_user_query(user_name, True)
-
-        # Return response
-        return response
diff --git a/server/lomas_server/dp_queries/dp_querier.py b/server/lomas_server/dp_queries/dp_querier.py
index 5828b789..4a7c9750 100644
--- a/server/lomas_server/dp_queries/dp_querier.py
+++ b/server/lomas_server/dp_queries/dp_querier.py
@@ -1,11 +1,24 @@
 from abc import ABC, abstractmethod
+from typing import Any, Generic, List, TypeVar
 
-from pydantic import BaseModel
+from lomas_server.admin_database.admin_database import AdminDatabase
+from lomas_server.data_connector.data_connector import DataConnector
+from lomas_server.utils.error_handler import (
+    KNOWN_EXCEPTIONS,
+    InternalServerException,
+    InvalidQueryException,
+    UnauthorizedAccessException,
+)
+from lomas_server.utils.query_models import (  # pylint: disable=W0611
+    QueryModel,
+    RequestModel,
+)
 
-from data_connector.data_connector import DataConnector
+RequestModelGeneric = TypeVar("RequestModelGeneric", bound="RequestModel")
+QueryModelGeneric = TypeVar("QueryModelGeneric", bound="QueryModel")
 
 
-class DPQuerier(ABC):
+class DPQuerier(ABC, Generic[RequestModelGeneric, QueryModelGeneric]):
     """
     Abstract Base Class for Queriers to external DP library.
 
@@ -13,36 +26,139 @@ class DPQuerier(ABC):
     a querier instance is specific to a DataConnector instance.
     """
 
-    def __init__(self, data_connector: DataConnector) -> None:
+    def __init__(
+        self,
+        data_connector: DataConnector,
+        admin_database: AdminDatabase,
+    ) -> None:
         """Initialise with specific dataset
 
         Args:
             data_connector (DataConnector): The private dataset to query.
+            admin_database (AdminDatabase): An initialized instance of
+                an AdminDatabase.
         """
         self.data_connector = data_connector
+        self.admin_database = admin_database
 
     @abstractmethod
-    def cost(self, query_json: BaseModel) -> tuple[float, float]:
+    def cost(self, query_json: RequestModelGeneric) -> tuple[float, float]:
         """
         Estimate cost of query.
 
         Args:
-            query_json (BaseModel): The JSON request object for the query.
-
+            query_json (RequestModelGeneric): The input object of the request.
+                Must be a subclass of RequestModel.
         Returns:
             tuple[float, float]: The tuple of costs, the first value is
                 the epsilon cost, the second value is the delta value.
         """
 
     @abstractmethod
-    def query(self, query_json: BaseModel) -> str:
+    def query(
+        self, query_json: QueryModelGeneric
+    ) -> dict | int | float | List[Any] | Any | str:
         """
         Perform the query and return the response.
 
         Args:
-            query_json (BaseModel): The JSON request object for the query.
+            query_json (QueryModelGeneric): The input object of the query.
+              Must be a subclass of QueryModel.
+
+        Returns:
+            dict | int | float | List[Any] | Any | str:
+                The query result, to be added to the response dict.
+        """
+
+    def handle_query(
+        self,
+        query_json: QueryModel,
+        user_name: str,
+    ) -> dict:
+        """
+        Handle DP query.
+
+        Args:
+            query_json (RequestModel): The input object of the query.
+              Must be a subclass of QueryModel.
+            user_name (str, optional): User name.
+
+        Raises:
+            UnauthorizedAccessException: A query is already
+                ongoing for this user,
+            the user does not exist or does not have access to the dataset.
+            InvalidQueryException: If the query is not valid.
+            InternalServerException: For any other unforseen exceptions.
 
         Returns:
-            TODO check this.
-            str: The JSON encoded string representation of the query result.
+            dict: A dictionary containing:
+                - requested_by (str): The user name.
+                - query_response (pd.DataFrame): A DataFrame containing
+                  the query response.
+                - spent_epsilon (float): The amount of epsilon budget spent
+                for the query.
+                - spent_delta (float): The amount of delta budget spent
+                  for the query.
+
         """
+        # Block access to other queries to user
+        if not self.admin_database.get_and_set_may_user_query(
+            user_name, False
+        ):
+            raise UnauthorizedAccessException(
+                f"User {user_name} is trying to query"
+                + " before end of previous query."
+            )
+
+        try:
+            # Get cost of the query
+            eps_cost, delta_cost = self.cost(query_json)  # type: ignore [arg-type]
+
+            # Check that enough budget to do the query
+            try:
+                (
+                    eps_remain,
+                    delta_remain,
+                ) = self.admin_database.get_remaining_budget(
+                    user_name, query_json.dataset_name
+                )
+            except UnauthorizedAccessException as e:
+                raise e
+
+            if (eps_remain < eps_cost) or (delta_remain < delta_cost):
+                raise InvalidQueryException(
+                    "Not enough budget for this query epsilon remaining "
+                    f"{eps_remain}, delta remaining {delta_remain}."
+                )
+
+            # Query
+            try:
+                query_response = self.query(query_json)  # type: ignore [arg-type]
+            except KNOWN_EXCEPTIONS as e:
+                raise e
+            except Exception as e:
+                raise InternalServerException(str(e)) from e
+
+            # Deduce budget from user
+            self.admin_database.update_budget(
+                user_name, query_json.dataset_name, eps_cost, delta_cost
+            )
+            response = {
+                "requested_by": user_name,
+                "query_response": query_response,
+                "spent_epsilon": eps_cost,
+                "spent_delta": delta_cost,
+            }
+
+            # Add query to db (for archive)
+            self.admin_database.save_query(user_name, query_json, response)
+
+        except Exception as e:
+            self.admin_database.set_may_user_query(user_name, True)
+            raise e
+
+        # Re-enable user to query
+        self.admin_database.set_may_user_query(user_name, True)
+
+        # Return response
+        return response
diff --git a/server/lomas_server/dp_queries/dummy_dataset.py b/server/lomas_server/dp_queries/dummy_dataset.py
index 0c365e70..92b97d44 100644
--- a/server/lomas_server/dp_queries/dummy_dataset.py
+++ b/server/lomas_server/dp_queries/dummy_dataset.py
@@ -1,33 +1,35 @@
-import datetime
-import random
-
 import numpy as np
 import pandas as pd
 
-from admin_database.admin_database import AdminDatabase
-from constants import (
-    DEFAULT_NUMERICAL_MAX,
-    DEFAULT_NUMERICAL_MIN,
+from lomas_server.admin_database.admin_database import AdminDatabase
+from lomas_server.constants import (
     DUMMY_NB_ROWS,
     DUMMY_SEED,
     NB_RANDOM_NONE,
-    RANDOM_DATE_RANGE,
-    RANDOM_DATE_START,
     RANDOM_STRINGS,
 )
-from data_connector.in_memory_connector import InMemoryConnector
-from utils.error_handler import InternalServerException
-from utils.query_models import GetDummyDataset
+from lomas_server.data_connector.in_memory_connector import InMemoryConnector
+from lomas_server.utils.collection_models import (
+    BooleanMetadata,
+    CategoricalColumnMetadata,
+    DatetimeMetadata,
+    FloatMetadata,
+    IntMetadata,
+    Metadata,
+    StrMetadata,
+)
+from lomas_server.utils.error_handler import InternalServerException
+from lomas_server.utils.query_models import RequestModel
 
 
 def make_dummy_dataset(  # pylint: disable=too-many-locals
-    metadata: dict, nb_rows: int = DUMMY_NB_ROWS, seed: int = DUMMY_SEED
+    metadata: Metadata, nb_rows: int = DUMMY_NB_ROWS, seed: int = DUMMY_SEED
 ) -> pd.DataFrame:
     """
     Create a dummy dataset based on a metadata dictionnary
 
     Args:
-        metadata (dict): dictionnary of the metadata of the real dataset
+        metadata (Metadata): The metadata model for the real dataset.
         nb_rows (int, optional): _description_. Defaults to DUMMY_NB_ROWS.
         seed (int, optional): _description_. Defaults to DUMMY_SEED.
 
@@ -37,100 +39,62 @@ def make_dummy_dataset(  # pylint: disable=too-many-locals
     Returns:
         pd.DataFrame: dummy dataframe based on metadata
     """
-    # Setting seed
-    random.seed(seed)
-    np.random.seed(seed)
+    # Creating new random generator with fixed seed
+    rng = np.random.default_rng(seed)
 
     # Create dataframe
     df = pd.DataFrame()
-    for col_name, data in metadata["columns"].items():
+    for col_name, data in metadata.columns.items():
         # Create a random serie based on the data type
-        match data["type"]:
-            case "string":
-                if "cardinality" in data.keys():
-                    cardinality = data["cardinality"]
-                    if "categories" in data.keys():
-                        categories = data["categories"]
-                        serie = pd.Series(
-                            random.choices(categories, k=nb_rows)
-                        )
-                    else:
-                        serie = pd.Series(
-                            random.choices(
-                                RANDOM_STRINGS[:cardinality], k=nb_rows
-                            )
-                        )
-                else:
-                    serie = pd.Series(
-                        random.choices(RANDOM_STRINGS, k=nb_rows)
-                    )
-            case "boolean":
-                # type boolean instead of bool will allow null values
+        match data:
+            case CategoricalColumnMetadata():
+                categories = data.categories
+                serie = pd.Series(rng.choice(categories, size=nb_rows))
+            case StrMetadata():
+                serie = pd.Series(rng.choice(RANDOM_STRINGS, size=nb_rows))
+            case BooleanMetadata():
+                # type boolean instead of bool will allow null values below
                 serie = pd.Series(
-                    random.choices([True, False], k=nb_rows), dtype="boolean"
-                )
-            case "int" | "float":
-                column_min = (
-                    data["lower"]
-                    if "lower" in data.keys()
-                    else DEFAULT_NUMERICAL_MIN
+                    rng.choice([True, False], size=nb_rows), dtype="boolean"
                 )
-                column_max = (
-                    data["upper"]
-                    if "upper" in data.keys()
-                    else DEFAULT_NUMERICAL_MAX
-                )
-                if data["type"] == "int":
-                    # pd.Series to ensure consistency between different types
-                    serie = pd.Series(
-                        np.random.randint(column_min, column_max, size=nb_rows)
-                    )
-                else:
-                    serie = pd.Series(
-                        np.random.uniform(column_min, column_max, size=nb_rows)
-                    )
-            case "datetime":
-                # From start date and random on a range above
-                start_date = (
-                    data["lower"]
-                    if "lower" in data.keys()
-                    else datetime.datetime.strptime(
-                        RANDOM_DATE_START, "%m/%d/%Y"
-                    )
+            case IntMetadata():
+                # pd.Series to ensure consistency between different types
+                dtype = f"{data.type}{data.precision}"
+                serie = pd.Series(
+                    rng.integers(
+                        data.lower,
+                        high=data.upper,
+                        endpoint=True,
+                        size=nb_rows,
+                    ),
+                    dtype=np.dtype(dtype),
                 )
-                end_date = (
-                    data["upper"]
-                    if "upper" in data.keys()
-                    else datetime.datetime.strptime(
-                        RANDOM_DATE_START, "%m/%d/%Y"
-                    )
-                    + datetime.timedelta(
-                        seconds=random.randrange(RANDOM_DATE_RANGE)
-                    )
+            case FloatMetadata():
+                dtype = f"{data.type}{data.precision}"
+                serie = pd.Series(
+                    data.lower
+                    + (data.upper - data.lower)
+                    * rng.random(size=nb_rows, dtype=np.dtype(dtype))
                 )
+            case DatetimeMetadata():
                 serie = pd.Series(
-                    np.random.choice(
-                        pd.date_range(start=start_date, end=end_date),
+                    rng.choice(
+                        pd.date_range(start=data.lower, end=data.upper),
                         size=nb_rows,
                     )
                 )
-            case "unknown":
-                # Unknown column are ignored by smartnoise sql
-                continue
             case _:
                 raise InternalServerException(
                     f"unknown column type in metadata: \
-                    {data['type']} in column {col_name}"
+                    {type(data)} in column {col_name}"
                 )
 
         # Add None value if the column is nullable
-        nullable = data["nullable"] if "nullable" in data.keys() else False
-
-        if nullable:
+        if data.nullable:
             # Get the indexes of 'serie'
             indexes = serie.index.tolist()
             for _ in range(0, NB_RANDOM_NONE):
-                index_to_insert = random.choice(indexes)
+                index_to_insert = rng.choice(indexes)
                 serie.at[index_to_insert] = None
 
         # Add randomly generated data as new column of dataframe
@@ -140,14 +104,14 @@ def make_dummy_dataset(  # pylint: disable=too-many-locals
 
 
 def get_dummy_dataset_for_query(
-    admin_database: AdminDatabase, query_json: GetDummyDataset
+    admin_database: AdminDatabase, query_json: RequestModel
 ) -> InMemoryConnector:
     """Get a dummy dataset for a given query.
 
     Args:
         admin_database (AdminDatabase): An initialized instance
             of AdminDatabase.
-        query_json (GetDummyDataset): The JSON request object for the query.
+        query_json (RequestModel): The request object for the query.
 
 
     Returns:
@@ -156,7 +120,9 @@ def get_dummy_dataset_for_query(
     # Create dummy dataset based on seed and number of rows
     ds_metadata = admin_database.get_dataset_metadata(query_json.dataset_name)
     ds_df = make_dummy_dataset(
-        ds_metadata, query_json.dummy_nb_rows, query_json.dummy_seed
+        ds_metadata,
+        query_json.dummy_nb_rows,
+        query_json.dummy_seed,
     )
     ds_data_connector = InMemoryConnector(ds_metadata, ds_df)
 
diff --git a/server/lomas_server/mongodb_admin.py b/server/lomas_server/mongodb_admin.py
index dfb17812..53867f2f 100644
--- a/server/lomas_server/mongodb_admin.py
+++ b/server/lomas_server/mongodb_admin.py
@@ -8,11 +8,18 @@
 from pymongo.database import Database
 from pymongo.results import _WriteResult
 
-from admin_database.mongodb_database import check_result_acknowledged
-from constants import PrivateDatabaseType
-from utils.collection_models import DatasetsCollection, UserCollection
-from utils.error_handler import InternalServerException
-from utils.logger import LOG
+from lomas_server.admin_database.mongodb_database import (
+    check_result_acknowledged,
+)
+from lomas_server.constants import PrivateDatabaseType
+from lomas_server.utils.collection_models import (
+    DatasetsCollection,
+    MetadataOfPathDB,
+    MetadataOfS3DB,
+    UserCollection,
+)
+from lomas_server.utils.error_handler import InternalServerException
+from lomas_server.utils.logger import LOG
 
 
 def check_user_exists(enforce_true: bool) -> Callable:
@@ -427,8 +434,10 @@ def add_users_via_yaml(
     # Load yaml data and insert it
     if isinstance(yaml_file, str):
         with open(yaml_file, encoding="utf-8") as f:
-            yaml_file = yaml.safe_load(f)
-    user_dict = UserCollection(**yaml_file)
+            yaml_dict: dict = yaml.safe_load(f)
+    else:
+        yaml_dict = yaml_file
+    user_dict = UserCollection(**yaml_dict)
 
     # Filter out duplicates
     new_users = []
@@ -662,8 +671,10 @@ def add_datasets_via_yaml(  # pylint: disable=R0912, R0914, R0915
 
     if isinstance(yaml_file, str):
         with open(yaml_file, encoding="utf-8") as f:
-            yaml_file = yaml.safe_load(f)
-    dataset_dict = DatasetsCollection(**yaml_file)
+            yaml_dict: dict = yaml.safe_load(f)
+    else:
+        yaml_dict = yaml_file
+    dataset_dict = DatasetsCollection(**yaml_dict)
 
     # Step 1: add datasets
     new_datasets = []
@@ -704,12 +715,12 @@ def add_datasets_via_yaml(  # pylint: disable=R0912, R0914, R0915
         dataset_name = d.dataset_name
         metadata_db_type = d.metadata.database_type
 
-        match metadata_db_type:
-            case PrivateDatabaseType.PATH:
+        match d.metadata:
+            case MetadataOfPathDB():
                 with open(d.metadata.metadata_path, encoding="utf-8") as f:
                     metadata_dict = yaml.safe_load(f)
 
-            case PrivateDatabaseType.S3:
+            case MetadataOfS3DB():
                 client = boto3.client(
                     "s3",
                     endpoint_url=d.metadata.endpoint_url,
diff --git a/server/lomas_server/mongodb_admin_cli.py b/server/lomas_server/mongodb_admin_cli.py
index a6ac9fd8..6057e54b 100644
--- a/server/lomas_server/mongodb_admin_cli.py
+++ b/server/lomas_server/mongodb_admin_cli.py
@@ -3,8 +3,9 @@
 from pymongo import MongoClient
 from pymongo.database import Database
 
-from admin_database.utils import get_mongodb_url
-from mongodb_admin import (
+from lomas_server.admin_database.utils import get_mongodb_url
+from lomas_server.constants import AdminDBType
+from lomas_server.mongodb_admin import (
     add_dataset,
     add_dataset_to_user,
     add_datasets_via_yaml,
@@ -26,6 +27,7 @@
     set_budget_field,
     set_may_query,
 )
+from lomas_server.utils.config import MongoDBConfig
 
 if __name__ == "__main__":
     ########################################################################
@@ -386,7 +388,15 @@
 
     #######################  FUNCTION CALL  ###################### # noqa: E266
     # Get MongoDB
-    db_url = get_mongodb_url(args)
+    mongo_config = MongoDBConfig(
+        db_type=AdminDBType.MONGODB,
+        address=args.address,
+        port=args.port,
+        username=args.username,
+        password=args.password,
+        db_name=args.db_name,
+    )
+    db_url = get_mongodb_url(mongo_config)
     mongo_db: Database = MongoClient(db_url)[args.db_name]
 
     function_map = {
diff --git a/server/lomas_server/routes/routes_admin.py b/server/lomas_server/routes/routes_admin.py
index 41de6d38..db8e569b 100644
--- a/server/lomas_server/routes/routes_admin.py
+++ b/server/lomas_server/routes/routes_admin.py
@@ -1,19 +1,19 @@
 from fastapi import APIRouter, Body, Depends, Header, Request
 from fastapi.responses import JSONResponse, RedirectResponse
 
-from data_connector.data_connector import get_column_dtypes
-from dp_queries.dummy_dataset import make_dummy_dataset
-from routes.utils import server_live
-from utils.error_handler import (
+from lomas_server.data_connector.data_connector import get_column_dtypes
+from lomas_server.dp_queries.dummy_dataset import make_dummy_dataset
+from lomas_server.routes.utils import server_live
+from lomas_server.utils.error_handler import (
     KNOWN_EXCEPTIONS,
     InternalServerException,
     UnauthorizedAccessException,
 )
-from utils.query_examples import (
+from lomas_server.utils.query_examples import (
     example_get_admin_db_data,
     example_get_dummy_dataset,
 )
-from utils.query_models import GetDbData, GetDummyDataset
+from lomas_server.utils.query_models import GetDbData, GetDummyDataset
 
 router = APIRouter()
 
@@ -52,29 +52,6 @@ async def get_state(
     )
 
 
-@router.get(
-    "/get_memory_usage",
-    dependencies=[Depends(server_live)],
-    tags=["ADMIN_USER"],
-)
-async def get_memory_usage(request: Request) -> JSONResponse:
-    """Return the dataset store object memory usage
-    Args:
-        request (Request): Raw request object
-        user_name (str, optional): The user name. Defaults to Header(None).
-
-    Returns:
-        JSONResponse: with DatasetStore object memory usage
-    """
-    app = request.app
-
-    return JSONResponse(
-        content={
-            "memory_usage": app.state.dataset_store.memory_usage,
-        }
-    )
-
-
 # Metadata query
 @router.post(
     "/get_dataset_metadata",
@@ -122,7 +99,7 @@ def get_dataset_metadata(
     except KNOWN_EXCEPTIONS as e:
         raise e
     except Exception as e:
-        raise InternalServerException(e) from e
+        raise InternalServerException(str(e)) from e
 
     return ds_metadata
 
@@ -177,7 +154,9 @@ def get_dummy_dataset(
         dtypes, datetime_columns = get_column_dtypes(ds_metadata)
 
         dummy_df = make_dummy_dataset(
-            ds_metadata, query_json.dummy_nb_rows, query_json.dummy_seed
+            ds_metadata,
+            query_json.dummy_nb_rows,
+            query_json.dummy_seed,
         )
 
         for col in datetime_columns:
@@ -186,7 +165,7 @@ def get_dummy_dataset(
     except KNOWN_EXCEPTIONS as e:
         raise e
     except Exception as e:
-        raise InternalServerException(e) from e
+        raise InternalServerException(str(e)) from e
 
     return JSONResponse(
         content={
@@ -245,7 +224,7 @@ def get_initial_budget(
     except KNOWN_EXCEPTIONS as e:
         raise e
     except Exception as e:
-        raise InternalServerException(e) from e
+        raise InternalServerException(str(e)) from e
 
     return JSONResponse(
         content={
@@ -303,7 +282,7 @@ def get_total_spent_budget(
     except KNOWN_EXCEPTIONS as e:
         raise e
     except Exception as e:
-        raise InternalServerException(e) from e
+        raise InternalServerException(str(e)) from e
 
     return JSONResponse(
         content={
@@ -358,7 +337,7 @@ def get_remaining_budget(
     except KNOWN_EXCEPTIONS as e:
         raise e
     except Exception as e:
-        raise InternalServerException(e) from e
+        raise InternalServerException(str(e)) from e
 
     return JSONResponse(
         content={
@@ -414,6 +393,6 @@ def get_user_previous_queries(
     except KNOWN_EXCEPTIONS as e:
         raise e
     except Exception as e:
-        raise InternalServerException(e) from e
+        raise InternalServerException(str(e)) from e
 
     return JSONResponse(content={"previous_queries": previous_queries})
diff --git a/server/lomas_server/routes/routes_dp.py b/server/lomas_server/routes/routes_dp.py
index f1eb7897..7f7b972d 100644
--- a/server/lomas_server/routes/routes_dp.py
+++ b/server/lomas_server/routes/routes_dp.py
@@ -1,14 +1,14 @@
 from fastapi import APIRouter, Body, Depends, Header, Request
 from fastapi.responses import JSONResponse
 
-from constants import DPLibraries
-from routes.utils import (
+from lomas_server.constants import DPLibraries
+from lomas_server.routes.utils import (
     handle_cost_query,
-    handle_query_on_data_connector,
     handle_query_on_dummy_dataset,
+    handle_query_on_private_dataset,
     server_live,
 )
-from utils.query_examples import (
+from lomas_server.utils.query_examples import (
     example_diffprivlib,
     example_dummy_diffprivlib,
     example_dummy_opendp,
@@ -20,17 +20,19 @@
     example_smartnoise_synth_cost,
     example_smartnoise_synth_query,
 )
-from utils.query_models import (
-    DiffPrivLibModel,
-    DummyDiffPrivLibModel,
-    DummyOpenDPModel,
-    DummySmartnoiseSQLModel,
-    DummySmartnoiseSynthQueryModel,
-    OpenDPModel,
-    SmartnoiseSQLCostModel,
-    SmartnoiseSQLModel,
-    SmartnoiseSynthCostModel,
+from lomas_server.utils.query_models import (
+    DiffPrivLibDummyQueryModel,
+    DiffPrivLibQueryModel,
+    DiffPrivLibRequestModel,
+    OpenDPDummyQueryModel,
+    OpenDPQueryModel,
+    OpenDPRequestModel,
+    SmartnoiseSQLDummyQueryModel,
+    SmartnoiseSQLQueryModel,
+    SmartnoiseSQLRequestModel,
+    SmartnoiseSynthDummyQueryModel,
     SmartnoiseSynthQueryModel,
+    SmartnoiseSynthRequestModel,
 )
 
 router = APIRouter()
@@ -43,7 +45,7 @@
 )
 def smartnoise_sql_handler(
     request: Request,
-    query_json: SmartnoiseSQLModel = Body(example_smartnoise_sql),
+    query_json: SmartnoiseSQLQueryModel = Body(example_smartnoise_sql),
     user_name: str = Header(None),
 ) -> JSONResponse:
     """
@@ -88,7 +90,7 @@ def smartnoise_sql_handler(
             - spent_delta (float): The amount of delta budget spent
               for the query.
     """
-    return handle_query_on_data_connector(
+    return handle_query_on_private_dataset(
         request, query_json, user_name, DPLibraries.SMARTNOISE_SQL
     )
 
@@ -101,7 +103,9 @@ def smartnoise_sql_handler(
 )
 def dummy_smartnoise_sql_handler(
     request: Request,
-    query_json: DummySmartnoiseSQLModel = Body(example_dummy_smartnoise_sql),
+    query_json: SmartnoiseSQLDummyQueryModel = Body(
+        example_dummy_smartnoise_sql
+    ),
     user_name: str = Header(None),
 ) -> JSONResponse:
     """
@@ -152,7 +156,7 @@ def dummy_smartnoise_sql_handler(
 )
 def estimate_smartnoise_sql_cost(
     request: Request,
-    query_json: SmartnoiseSQLCostModel = Body(example_smartnoise_sql_cost),
+    query_json: SmartnoiseSQLRequestModel = Body(example_smartnoise_sql_cost),
     user_name: str = Header(None),
 ) -> JSONResponse:
     """
@@ -160,7 +164,7 @@ def estimate_smartnoise_sql_cost(
 
     Args:
         request (Request): Raw request object
-        query_json (SmartnoiseSQLCostModel, optional):
+        query_json (SmartnoiseSQLRequestModel, optional):
             A JSON object containing the following:
             - query: The SQL query to estimate the cost for.
               NOTE: the table name is "df", the query must end with "FROM df".
@@ -205,7 +209,7 @@ def smartnoise_synth_handler(
     Handles queries for the SmartNoise Synth library.
     Args:
         request (Request): Raw request object
-        query_json (SNSQLInp): A JSON object containing:
+        query_json (SmartnoiseSynthQueryModel): A JSON object containing:
             - synth_name (str): name of the Synthesizer model to use.
             - epsilon (float): Privacy parameter (e.g., 0.1).
             - delta (float): Privacy parameter (e.g., 1e-5).
@@ -245,7 +249,7 @@ def smartnoise_synth_handler(
             - spent_delta (float): The amount of delta budget spent
               for the query.
     """
-    return handle_query_on_data_connector(
+    return handle_query_on_private_dataset(
         request, query_json, user_name, DPLibraries.SMARTNOISE_SYNTH
     )
 
@@ -257,7 +261,7 @@ def smartnoise_synth_handler(
 )
 def dummy_smartnoise_synth_handler(
     request: Request,
-    query_json: DummySmartnoiseSynthQueryModel = Body(
+    query_json: SmartnoiseSynthDummyQueryModel = Body(
         example_dummy_smartnoise_synth_query
     ),
     user_name: str = Header(None),
@@ -266,7 +270,7 @@ def dummy_smartnoise_synth_handler(
     Handles queries for the SmartNoise Synth library.
     Args:
         request (Request): Raw request object
-        query_json (SNSQLInp): A JSON object containing:
+        query_json (SmartnoiseSynthDummyQueryModel): A JSON object containing:
             - synth_name (str): name of the Synthesizer model to use.
             - epsilon (float): Privacy parameter (e.g., 0.1).
             - delta (float): Privacy parameter (e.g., 1e-5).
@@ -323,14 +327,16 @@ def dummy_smartnoise_synth_handler(
 )
 def estimate_smartnoise_synth_cost(
     request: Request,
-    query_json: SmartnoiseSynthCostModel = Body(example_smartnoise_synth_cost),
+    query_json: SmartnoiseSynthRequestModel = Body(
+        example_smartnoise_synth_cost
+    ),
     user_name: str = Header(None),
 ) -> JSONResponse:
     """
     Handles queries for the SmartNoise Synth library.
     Args:
         request (Request): Raw request object
-        query_json (SNSQLInp): A JSON object containing:
+        query_json (SmartnoiseSynthRequestModel): A JSON object containing:
             - synth_name (str): name of the Synthesizer model to use.
             - epsilon (float): Privacy parameter (e.g., 0.1).
             - delta (float): Privacy parameter (e.g., 1e-5).
@@ -373,7 +379,7 @@ def estimate_smartnoise_synth_cost(
 )
 def opendp_query_handler(
     request: Request,
-    query_json: OpenDPModel = Body(example_opendp),
+    query_json: OpenDPQueryModel = Body(example_opendp),
     user_name: str = Header(None),
 ) -> JSONResponse:
     """
@@ -381,7 +387,7 @@ def opendp_query_handler(
 
     Args:
         request (Request): Raw request object.
-        query_json (OpenDPModel, optional): A JSON object containing the following:
+        query_json (OpenDPQueryModel, optional): A JSON object containing the following:
             - opendp_pipeline: The OpenDP pipeline for the query.
             - fixed_delta: If the pipeline measurement is of type
                 "ZeroConcentratedDivergence" (e.g. with "make_gaussian") then it is
@@ -414,7 +420,7 @@ def opendp_query_handler(
             - spent_delta (float): The amount of delta budget spent
               for the query.
     """
-    response = handle_query_on_data_connector(
+    response = handle_query_on_private_dataset(
         request, query_json, user_name, DPLibraries.OPENDP
     )
     return JSONResponse(content=response)
@@ -427,7 +433,7 @@ def opendp_query_handler(
 )
 def dummy_opendp_query_handler(
     request: Request,
-    query_json: DummyOpenDPModel = Body(example_dummy_opendp),
+    query_json: OpenDPDummyQueryModel = Body(example_dummy_opendp),
     user_name: str = Header(None),
 ) -> JSONResponse:
     """
@@ -435,7 +441,7 @@ def dummy_opendp_query_handler(
 
     Args:
         request (Request): Raw request object.
-        query_json (DummyOpenDPModel, optional):
+        query_json (OpenDPDummyQueryModel, optional):
             A JSON object containing the following:
             - opendp_pipeline: The OpenDP pipeline for the query.
             - fixed_delta: If the pipeline measurement is of type\
@@ -475,7 +481,7 @@ def dummy_opendp_query_handler(
 )
 def estimate_opendp_cost(
     request: Request,
-    query_json: OpenDPModel = Body(example_opendp),
+    query_json: OpenDPRequestModel = Body(example_opendp),
     user_name: str = Header(None),
 ) -> JSONResponse:
     """
@@ -483,7 +489,7 @@ def estimate_opendp_cost(
 
     Args:
         request (Request): Raw request object
-        query_json (OpenDPModel, optional):
+        query_json (OpenDPRequestModel, optional):
             A JSON object containing the following:
             - "opendp_pipeline": The OpenDP pipeline for the query.
 
@@ -513,7 +519,7 @@ def estimate_opendp_cost(
 )
 def diffprivlib_query_handler(
     request: Request,
-    query_json: DiffPrivLibModel = Body(example_diffprivlib),
+    query_json: DiffPrivLibQueryModel = Body(example_diffprivlib),
     user_name: str = Header(None),
 ):
     """
@@ -521,13 +527,14 @@ def diffprivlib_query_handler(
 
     Args:
         request (Request): Raw request object.
-        query_json (OpenDPModel, optional): A JSON object containing the following:
-            - pipeline: The DiffPrivLib pipeline for the query.
-            - feature_columns: the list of feature column to train
-            - target_columns: the list of target column to predict
-            - test_size: proportion of the test set
-            - test_train_split_seed: seed for the random train test split,
-            - imputer_strategy: imputation strategy
+        query_json (DiffPrivLibQueryModel, optional):
+            A JSON object containing the following:
+                - pipeline: The DiffPrivLib pipeline for the query.
+                - feature_columns: the list of feature column to train
+                - target_columns: the list of target column to predict
+                - test_size: proportion of the test set
+                - test_train_split_seed: seed for the random train test split,
+                - imputer_strategy: imputation strategy
 
             Defaults to Body(example_diffprivlib).
 
@@ -553,7 +560,7 @@ def diffprivlib_query_handler(
             - spent_delta (float): The amount of delta budget spent
               for the query.
     """
-    return handle_query_on_data_connector(
+    return handle_query_on_private_dataset(
         request, query_json, user_name, DPLibraries.DIFFPRIVLIB
     )
 
@@ -565,7 +572,7 @@ def diffprivlib_query_handler(
 )
 def dummy_diffprivlib_query_handler(
     request: Request,
-    query_json: DummyDiffPrivLibModel = Body(example_dummy_diffprivlib),
+    query_json: DiffPrivLibDummyQueryModel = Body(example_dummy_diffprivlib),
     user_name: str = Header(None),
 ):
     """
@@ -573,17 +580,18 @@ def dummy_diffprivlib_query_handler(
 
     Args:
         request (Request): Raw request object.
-        query_json (DiffPrivLibModel, optional): A JSON object containing the following:
-            - pipeline: The DiffPrivLib pipeline for the query.
-            - feature_columns: the list of feature column to train
-            - target_columns: the list of target column to predict
-            - test_size: proportion of the test set
-            - test_train_split_seed: seed for the random train test split,
-            - imputer_strategy: imputation strategy
-            - nb_rows (int, optional): The number of rows in the dummy dataset
-              (default: 100).
-            - seed (int, optional): The random seed for generating
-              the dummy dataset (default: 42).
+        query_json (DiffPrivLibDummyQueryModel, optional):
+            A JSON object containing the following:
+                - pipeline: The DiffPrivLib pipeline for the query.
+                - feature_columns: the list of feature column to train
+                - target_columns: the list of target column to predict
+                - test_size: proportion of the test set
+                - test_train_split_seed: seed for the random train test split,
+                - imputer_strategy: imputation strategy
+                - nb_rows (int, optional):
+                  The number of rows in the dummy dataset (default: 100).
+                - seed (int, optional): The random seed for generating
+                the dummy dataset (default: 42).
               Defaults to Body(example_dummy_diffprivlib)
 
     Raises:
@@ -610,7 +618,7 @@ def dummy_diffprivlib_query_handler(
 )
 def estimate_diffprivlib_cost(
     request: Request,
-    query_json: DiffPrivLibModel = Body(example_diffprivlib),
+    query_json: DiffPrivLibRequestModel = Body(example_diffprivlib),
     user_name: str = Header(None),
 ):
     """
@@ -618,7 +626,7 @@ def estimate_diffprivlib_cost(
 
     Args:
         request (Request): Raw request object
-        query_json (DiffPrivLibModel, optional):
+        query_json (DiffPrivLibRequestModel, optional):
         A JSON object containing the following:
             - pipeline: The DiffPrivLib pipeline for the query.
             - feature_columns: the list of feature column to train
diff --git a/server/lomas_server/routes/utils.py b/server/lomas_server/routes/utils.py
index 069e5af7..b98d54aa 100644
--- a/server/lomas_server/routes/utils.py
+++ b/server/lomas_server/routes/utils.py
@@ -2,16 +2,21 @@
 
 from fastapi import Request
 from fastapi.responses import JSONResponse
-from pydantic import BaseModel
 
-from constants import DPLibraries
-from dp_queries.dp_libraries.factory import querier_factory
-from dp_queries.dummy_dataset import get_dummy_dataset_for_query
-from utils.error_handler import (
+from lomas_server.constants import DPLibraries
+from lomas_server.data_connector.factory import data_connector_factory
+from lomas_server.dp_queries.dp_libraries.factory import querier_factory
+from lomas_server.dp_queries.dummy_dataset import get_dummy_dataset_for_query
+from lomas_server.utils.error_handler import (
     KNOWN_EXCEPTIONS,
     InternalServerException,
     UnauthorizedAccessException,
 )
+from lomas_server.utils.query_models import (
+    DummyQueryModel,
+    QueryModel,
+    RequestModel,
+)
 
 
 async def server_live(request: Request) -> AsyncGenerator:
@@ -35,9 +40,9 @@ async def server_live(request: Request) -> AsyncGenerator:
     yield
 
 
-def handle_query_on_data_connector(
+def handle_query_on_private_dataset(
     request: Request,
-    query_json: BaseModel,
+    query_json: QueryModel,
     user_name: str,
     dp_library: DPLibraries,
 ):
@@ -71,21 +76,29 @@ def handle_query_on_data_connector(
     """
     app = request.app
 
+    data_connector = data_connector_factory(
+        query_json.dataset_name,
+        app.state.admin_database,
+        app.state.private_credentials,
+    )
+    dp_querier = querier_factory(
+        dp_library,
+        data_connector=data_connector,
+        admin_database=app.state.admin_database,
+    )
     try:
-        response = app.state.query_handler.handle_query(
-            dp_library, query_json, user_name
-        )
+        response = dp_querier.handle_query(query_json, user_name)
     except KNOWN_EXCEPTIONS as e:
         raise e
     except Exception as e:
-        raise InternalServerException(e) from e
+        raise InternalServerException(str(e)) from e
 
     return response
 
 
 def handle_query_on_dummy_dataset(
     request: Request,
-    query_json: BaseModel,
+    query_json: DummyQueryModel,
     user_name: str,
     dp_library: DPLibraries,
 ):
@@ -122,24 +135,32 @@ def handle_query_on_dummy_dataset(
         app.state.admin_database, query_json
     )
     dummy_querier = querier_factory(
-        dp_library, data_connector=ds_data_connector
+        dp_library,
+        data_connector=ds_data_connector,
+        admin_database=app.state.admin_database,
     )
 
     try:
-        _ = dummy_querier.cost(query_json)  # verify cost works
+        eps_cost, delta_cost = dummy_querier.cost(query_json)
         response_df = dummy_querier.query(query_json)
-        response = JSONResponse(content={"query_response": response_df})
+        response = JSONResponse(
+            content={
+                "query_response": response_df,
+                "epsilon": eps_cost,
+                "delta": delta_cost,
+            }
+        )
     except KNOWN_EXCEPTIONS as e:
         raise e
     except Exception as e:
-        raise InternalServerException(e) from e
+        raise InternalServerException(str(e)) from e
 
     return response
 
 
 def handle_cost_query(
     request: Request,
-    query_json: BaseModel,
+    query_json: RequestModel,
     user_name: str,
     dp_library: DPLibraries,
 ):
@@ -173,13 +194,23 @@ def handle_cost_query(
             f"{user_name} does not have access to {dataset_name}.",
         )
 
+    data_connector = data_connector_factory(
+        query_json.dataset_name,
+        app.state.admin_database,
+        app.state.private_credentials,
+    )
+    dp_querier = querier_factory(
+        dp_library,
+        data_connector=data_connector,
+        admin_database=app.state.admin_database,
+    )
     try:
-        response = app.state.query_handler.estimate_cost(
-            dp_library, query_json
-        )
+        eps_cost, delta_cost = dp_querier.cost(query_json)
     except KNOWN_EXCEPTIONS as e:
         raise e
     except Exception as e:
-        raise InternalServerException(e) from e
+        raise InternalServerException(str(e)) from e
 
-    return JSONResponse(content=response)
+    return JSONResponse(
+        content={"epsilon_cost": eps_cost, "delta_cost": delta_cost}
+    )
diff --git a/server/lomas_server/utils/anti_timing_att.py b/server/lomas_server/utils/anti_timing_att.py
index 1a95f8e4..8e8381b3 100644
--- a/server/lomas_server/utils/anti_timing_att.py
+++ b/server/lomas_server/utils/anti_timing_att.py
@@ -4,8 +4,8 @@
 
 from fastapi import Request, Response
 
-from utils.config import Config
-from utils.error_handler import InternalServerException
+from lomas_server.utils.config import Config
+from lomas_server.utils.error_handler import InternalServerException
 
 
 async def anti_timing_att(
diff --git a/server/lomas_server/utils/collection_models.py b/server/lomas_server/utils/collection_models.py
index 8dc31bea..c706d1a7 100644
--- a/server/lomas_server/utils/collection_models.py
+++ b/server/lomas_server/utils/collection_models.py
@@ -1,8 +1,10 @@
-from typing import Dict, List, Literal, Optional, Union
+from datetime import datetime
+from enum import IntEnum
+from typing import Annotated, Any, Dict, List, Literal, Optional, Union
 
-from pydantic import BaseModel, Field
+from pydantic import BaseModel, Discriminator, Field, Tag, model_validator
 
-from constants import PrivateDatabaseType
+from lomas_server.constants import PrivateDatabaseType
 
 
 class DatasetOfUser(BaseModel):
@@ -81,15 +83,162 @@ class DatasetOfS3DB(Dataset):
 class DatasetsCollection(BaseModel):
     """BaseModel for datasets collection"""
 
-    datasets: List[Union[DatasetOfPathDB, DatasetOfS3DB]] = Field(
-        ..., discriminator="database_type"
-    )
+    datasets: Annotated[
+        List[Union[DatasetOfPathDB, DatasetOfS3DB]],
+        Field(discriminator="database_type"),
+    ]
+
+
+class ColumnMetadata(BaseModel):
+    """Base model for column metadata"""
+
+    private_id: bool = False
+    nullable: bool = False
+    # TODO create validator and test these, see issue #323
+    max_partition_length: Optional[Annotated[int, Field(gt=0)]] = None
+    max_influenced_partitions: Optional[Annotated[int, Field(gt=0)]] = None
+    max_partition_contributions: Optional[Annotated[int, Field(gt=0)]] = None
+
+
+class StrMetadata(ColumnMetadata):
+    """Model for string metadata"""
+
+    type: Literal["string"]
+
+
+class CategoricalColumnMetadata(ColumnMetadata):
+    """Model for categorical column metadata"""
+
+    @model_validator(mode="after")
+    def validate_categories(self):
+        """Makes sure number of categories matches cardinality."""
+        if len(self.categories) != self.cardinality:
+            raise ValueError(
+                "Number of categories should be equal to cardinality."
+            )
+        return self
+
+
+class StrCategoricalMetadata(CategoricalColumnMetadata):
+    """Model for categorical string metadata"""
+
+    type: Literal["string"]
+    cardinality: int
+    categories: List[str]
+
+
+class Precision(IntEnum):
+    """Precision of integer and float data"""
+
+    SINGLE = 32
+    DOUBLE = 64
+
+
+class BoundedColumnMetadata(ColumnMetadata):
+    """Model for columns with bounded data"""
+
+    @model_validator(mode="after")
+    def validate_bounds(self):
+        """Validates column bounds."""
+        if (
+            self.lower is not None
+            and self.upper is not None
+            and self.lower > self.upper
+        ):
+            raise ValueError("Lower bound cannot be larger than upper bound.")
+
+        return self
+
+
+class IntMetadata(BoundedColumnMetadata):
+    """Model for integer column metadata"""
+
+    type: Literal["int"]
+    precision: Precision
+    lower: int
+    upper: int
+
+
+class IntCategoricalMetadata(CategoricalColumnMetadata):
+    """Model for integer categorical column metadata"""
+
+    type: Literal["int"]
+    precision: Precision
+    cardinality: int
+    categories: List[int]
+
+
+class FloatMetadata(BoundedColumnMetadata):
+    """Model for float column metadata"""
+
+    type: Literal["float"]
+    precision: Precision
+    lower: float
+    upper: float
+
+
+class BooleanMetadata(ColumnMetadata):
+    """Model for boolean column metadata"""
+
+    type: Literal["boolean"]
+
+
+class DatetimeMetadata(BoundedColumnMetadata):
+    """Model for datetime column metadata"""
+
+    type: Literal["datetime"]  # TODO make these constants, see issue #268
+    lower: datetime
+    upper: datetime
+
+
+def get_column_metadata_discriminator(v: Any) -> str:
+    """Discriminator function for determining the type of column metadata.
+
+    Args:
+        v (Any): The unparsed column metadata (either dict or class object)
+
+    Raises:
+        ValueError: If the column type cannot be found.
+
+    Returns:
+        str: The metadata string type.
+    """
+    if isinstance(v, dict):
+        col_type = v.get("type")
+    else:
+        col_type = getattr(v, "type")
+
+    if (col_type in ("string", "int")) and (
+        ((isinstance(v, dict)) and "cardinality" in v)
+        or (hasattr(v, "cardinality"))
+    ):
+        col_type = f"categorical_{col_type}"
+
+    if not isinstance(col_type, str):
+        raise ValueError("Could not find column type.")
+
+    return col_type
 
 
 class Metadata(BaseModel):
     """BaseModel for a metadata format"""
 
-    max_ids: int
+    max_ids: Annotated[int, Field(gt=0)]
+    rows: Annotated[int, Field(gt=0)]
     row_privacy: bool
     censor_dims: Optional[bool] = False
-    columns: Dict[str, Dict[str, Union[int, float, str, List[str]]]]
+    columns: Dict[
+        str,
+        Annotated[
+            Union[
+                Annotated[StrMetadata, Tag("string")],
+                Annotated[StrCategoricalMetadata, Tag("categorical_string")],
+                Annotated[IntMetadata, Tag("int")],
+                Annotated[IntCategoricalMetadata, Tag("categorical_int")],
+                Annotated[FloatMetadata, Tag("float")],
+                Annotated[BooleanMetadata, Tag("boolean")],
+                Annotated[DatetimeMetadata, Tag("datetime")],
+            ],
+            Discriminator(get_column_metadata_discriminator),
+        ],
+    ]
diff --git a/server/lomas_server/utils/config.py b/server/lomas_server/utils/config.py
index 5f934fe2..22f91592 100644
--- a/server/lomas_server/utils/config.py
+++ b/server/lomas_server/utils/config.py
@@ -3,16 +3,15 @@
 import yaml
 from pydantic import BaseModel, ConfigDict, Field
 
-from constants import (
+from lomas_server.constants import (
     CONFIG_PATH,
     SECRETS_PATH,
     AdminDBType,
     ConfigKeys,
-    DatasetStoreType,
     PrivateDatabaseType,
     TimeAttackMethod,
 )
-from utils.error_handler import InternalServerException
+from lomas_server.utils.error_handler import InternalServerException
 
 
 class TimeAttack(BaseModel):
@@ -33,23 +32,6 @@ class Server(BaseModel):
     workers: int
 
 
-class DatasetStoreConfig(BaseModel):
-    """BaseModel for dataset store configs"""
-
-
-class BasicDatasetStoreConfig(DatasetStoreConfig):
-    """BaseModel for basic dataset store configs"""
-
-    ds_store_type: Literal[DatasetStoreType.BASIC.value]  # type: ignore
-
-
-class LRUDatasetStoreConfig(DatasetStoreConfig):
-    """BaseModel for dataset store configs in case of a LRU dataset store"""
-
-    ds_store_type: Literal[DatasetStoreType.LRU]  # type: ignore
-    max_memory_usage: int
-
-
 class DBConfig(BaseModel):
     """BaseModel for database type config"""
 
@@ -123,10 +105,6 @@ class Config(BaseModel):
         ..., discriminator="db_type"
     )
 
-    dataset_store: Union[BasicDatasetStoreConfig, LRUDatasetStoreConfig] = (
-        Field(..., discriminator="ds_store_type")
-    )
-
     dp_libraries: DPLibraryConfig
 
 
@@ -141,7 +119,7 @@ class ConfigLoader:
     """
 
     _instance = None
-    _config = None
+    _config: Config | None = None
 
     def __new__(cls):
         if cls._instance is None:
@@ -219,7 +197,8 @@ def get_config(self) -> Config:
         """
         if self._config is None:
             self.load_config()
-        return self._config  # type: ignore
+        assert isinstance(self._config, Config)  # Helps mypy
+        return self._config
 
 
 CONFIG_LOADER = ConfigLoader()
diff --git a/server/lomas_server/utils/error_handler.py b/server/lomas_server/utils/error_handler.py
index a09a3fda..db66b946 100644
--- a/server/lomas_server/utils/error_handler.py
+++ b/server/lomas_server/utils/error_handler.py
@@ -1,9 +1,11 @@
+from typing import Type
+
 from fastapi import FastAPI, Request, status
 from fastapi.responses import JSONResponse
 from pymongo.errors import WriteConcernError
 
-from constants import INTERNAL_SERVER_ERROR
-from utils.logger import LOG
+from lomas_server.constants import INTERNAL_SERVER_ERROR
+from lomas_server.utils.logger import LOG
 
 
 class InvalidQueryException(Exception):
@@ -63,7 +65,7 @@ def __init__(self, error_message: str) -> None:
         self.error_message = error_message
 
 
-KNOWN_EXCEPTIONS: tuple[type, ...] = (
+KNOWN_EXCEPTIONS: tuple[Type[BaseException], ...] = (
     ExternalLibraryException,
     InternalServerException,
     InvalidQueryException,
diff --git a/server/lomas_server/utils/query_examples.py b/server/lomas_server/utils/query_examples.py
index e081d713..9825f9e7 100644
--- a/server/lomas_server/utils/query_examples.py
+++ b/server/lomas_server/utils/query_examples.py
@@ -1,4 +1,4 @@
-from constants import (
+from lomas_server.constants import (
     DIFFPRIVLIB_PIPELINE,
     OPENDP_PIPELINE,
     SSynthGanSynthesizer,
diff --git a/server/lomas_server/utils/query_models.py b/server/lomas_server/utils/query_models.py
index f5eccfbd..d91c510e 100644
--- a/server/lomas_server/utils/query_models.py
+++ b/server/lomas_server/utils/query_models.py
@@ -2,12 +2,12 @@
 
 from pydantic import BaseModel, Field
 
-from constants import (
-    DELTA_LIMIT,
-    EPSILON_LIMIT,
+from lomas_server.constants import (
+    DPLibraries,
     SSynthGanSynthesizer,
     SSynthMarginalSynthesizer,
 )
+from lomas_server.utils.error_handler import InternalServerException
 
 
 class GetDbData(BaseModel):
@@ -24,54 +24,69 @@ class GetDummyDataset(BaseModel):
     dummy_seed: int
 
 
-class SmartnoiseSQLModel(BaseModel):
-    """Model input for a smarnoise-sql query"""
+class RequestModel(BaseModel):
+    """
+    Base input model for any request on a dataset.
+
+    We differentiate between requests and queries:
+        - a request does not necessarily require an algorithm
+          to be executed on the private dataset (e.g. some cost requests).
+        - a query requires executing an algorithm on a private
+          dataset (or a potentially a dummy).
+    """
 
-    query_str: str
     dataset_name: str
-    epsilon: float = Field(
-        ...,
-        gt=0,
-        le=EPSILON_LIMIT,
-    )
-    delta: float = Field(
-        ...,
-        gt=0,
-        le=DELTA_LIMIT,
-    )
-    mechanisms: dict
-    postprocess: bool
 
 
-class DummySmartnoiseSQLModel(BaseModel):
-    """Model input for a smarnoise-sql dummy query"""
+class QueryModel(RequestModel):
+    """
+    Base input model for any query on a dataset.
+
+    We differentiate between requests and queries:
+        - a request does not necessarily require an algorithm
+          to be executed on the private dataset (e.g. some cost requests).
+        - a query requires executing an algorithm on a private
+          dataset (or a potentially a dummy).
+    """
+
+
+class DummyQueryModel(QueryModel):
+    """
+    Input model for a query on a dummy dataset.
+    """
 
-    query_str: str
-    dataset_name: str
     dummy_nb_rows: int = Field(..., gt=0)
     dummy_seed: int
-    epsilon: float = Field(..., gt=0)
-    delta: float = Field(..., gt=0)
-    mechanisms: dict
-    postprocess: bool
 
 
-class SmartnoiseSQLCostModel(BaseModel):
-    """Model input for a smarnoise-sql cost query"""
+# SmartnoiseSQL
+# ----------------------------------------------------------------------------
+class SmartnoiseSQLRequestModel(RequestModel):
+    """Base input model for a smarnoise-sql request"""
 
     query_str: str
-    dataset_name: str
     epsilon: float = Field(..., gt=0)
     delta: float = Field(..., gt=0)
     mechanisms: dict
 
 
-class SmartnoiseSynthCostModel(BaseModel):
-    """Model input for a smarnoise-synth cost"""
+class SmartnoiseSQLQueryModel(SmartnoiseSQLRequestModel, QueryModel):
+    """Base input model for a smartnoise-sql query"""
+
+    postprocess: bool
+
+
+class SmartnoiseSQLDummyQueryModel(SmartnoiseSQLQueryModel, DummyQueryModel):
+    """Input model for a smartnoise-sql query on a dummy dataset."""
+
+
+# SmartnoiseSynth
+# ----------------------------------------------------------------------------
+class SmartnoiseSynthRequestModel(RequestModel):
+    """Base input model for a SmartnoiseSynth request"""
 
-    dataset_name: str
     synth_name: Union[SSynthMarginalSynthesizer, SSynthGanSynthesizer]
-    epsilon: float = Field(..., gt=0, le=EPSILON_LIMIT)
+    epsilon: float = Field(..., gt=0)
     delta: Optional[float] = None
     select_cols: List
     synth_params: dict
@@ -79,43 +94,47 @@ class SmartnoiseSynthCostModel(BaseModel):
     constraints: str
 
 
-class SmartnoiseSynthQueryModel(SmartnoiseSynthCostModel):
-    """Model input for a smarnoise-synth query"""
+class SmartnoiseSynthQueryModel(SmartnoiseSynthRequestModel, QueryModel):
+    """Base input model for a smarnoise-synth query"""
 
     return_model: bool
     condition: str
     nb_samples: int
 
 
-class DummySmartnoiseSynthQueryModel(SmartnoiseSynthQueryModel):
-    """Dummy Model input for a smarnoise-synth query"""
+class SmartnoiseSynthDummyQueryModel(
+    SmartnoiseSynthQueryModel, DummyQueryModel
+):
+    """Input model for a smarnoise-synth query on a dummy dataset"""
 
-    dummy_nb_rows: int = Field(..., gt=0)
-    dummy_seed: int
+    # Same as normal query.
+    return_model: bool
+    condition: str
+    nb_samples: int
 
 
-class OpenDPModel(BaseModel):
-    """Model input for an opendp query"""
+# OpenDP
+# ----------------------------------------------------------------------------
+class OpenDPRequestModel(RequestModel):
+    """Base input model for an opendp request"""
 
-    dataset_name: str
     opendp_json: str
     fixed_delta: Optional[float] = None
 
 
-class DummyOpenDPModel(BaseModel):
-    """Model input for a dummy opendp query"""
+class OpenDPQueryModel(OpenDPRequestModel, QueryModel):
+    """Base input model for an opendp query"""
 
-    dataset_name: str
-    opendp_json: str
-    dummy_nb_rows: int = Field(..., gt=0)
-    dummy_seed: int
-    fixed_delta: Optional[float] = None
 
+class OpenDPDummyQueryModel(OpenDPRequestModel, DummyQueryModel):
+    """Input model for an opendp query on a dummy dataset"""
 
-class DiffPrivLibModel(BaseModel):
-    """Model input for a diffprivlib query"""
 
-    dataset_name: str
+# DiffPrivLib
+# ----------------------------------------------------------------------------
+class DiffPrivLibRequestModel(RequestModel):
+    """Base input model for a diffprivlib request"""
+
     diffprivlib_json: str
     feature_columns: list
     target_columns: Optional[list]
@@ -124,15 +143,40 @@ class DiffPrivLibModel(BaseModel):
     imputer_strategy: str
 
 
-class DummyDiffPrivLibModel(BaseModel):
-    """Model input for a dummy diffprivlib query"""
+class DiffPrivLibQueryModel(DiffPrivLibRequestModel, QueryModel):
+    """Base input model for a diffprivlib query"""
 
-    dataset_name: str
-    diffprivlib_json: str
-    feature_columns: list
-    target_columns: Optional[list]
-    test_size: float = Field(..., gt=0.0, lt=1.0)
-    test_train_split_seed: int
-    imputer_strategy: str
-    dummy_nb_rows: int = Field(..., gt=0)
-    dummy_seed: int
+
+class DiffPrivLibDummyQueryModel(DiffPrivLibQueryModel, DummyQueryModel):
+    """Input model for a DiffPrivLib query on a dummy dataset"""
+
+
+# Utils
+# ----------------------------------------------------------------------------
+
+
+def model_input_to_lib(request: RequestModel) -> DPLibraries:
+    """Return the type of DP library given a RequestModel.
+
+    Args:
+        request (RequestModel): The user request
+
+    Raises:
+        InternalServerException: If the library type cannot be determined.
+
+    Returns:
+        DPLibraries: The type of library for the request.
+    """
+    match request:
+        case SmartnoiseSQLRequestModel():
+            return DPLibraries.SMARTNOISE_SQL
+        case SmartnoiseSynthRequestModel():
+            return DPLibraries.SMARTNOISE_SYNTH
+        case OpenDPRequestModel():
+            return DPLibraries.OPENDP
+        case DiffPrivLibRequestModel():
+            return DPLibraries.DIFFPRIVLIB
+        case _:
+            raise InternalServerException(
+                "Cannot find library type for given model."
+            )
diff --git a/server/lomas_server/uvicorn_serve.py b/server/lomas_server/uvicorn_serve.py
index 2d63614a..8559530d 100644
--- a/server/lomas_server/uvicorn_serve.py
+++ b/server/lomas_server/uvicorn_serve.py
@@ -1,24 +1,21 @@
-import os
-
 import uvicorn
 
-from utils.config import get_config
-from utils.logger import LOG
+from lomas_server.utils.config import get_config
+from lomas_server.utils.logger import LOG
 
 if __name__ == "__main__":
-    os.chdir("/code/")
 
     config = get_config()
 
     if config.server.workers != 1:
-        LOG.WARNING(
+        LOG.WARN(
             "Only supports one server worker.",
             "Overwriting server.workers config",
             f" from {config.server.workers} to 1.",
         )
 
     uvicorn.run(
-        "app:app",
+        "lomas_server.app:app",
         host=config.server.host_ip,
         port=config.server.host_port,
         log_level=config.server.log_level,
diff --git a/server/pyproject.toml b/server/pyproject.toml
index a5ce0933..595d2014 100644
--- a/server/pyproject.toml
+++ b/server/pyproject.toml
@@ -26,5 +26,6 @@ disable = [
     "E0611",  # no-name-in-module
     "W0511",  # fixme
     "C0302",  # too-many-lines
+    "W1203",  # use lazy % formatting in logging functions
     "duplicate-code"  # duplicate-code
 ]
diff --git a/server/requirements.txt b/server/requirements.txt
index d055fa95..9decdfb3 100644
--- a/server/requirements.txt
+++ b/server/requirements.txt
@@ -6,7 +6,7 @@ git+https://github.com/ryan112358/private-pgm.git@01f02f17eba440f4e76c1d06fa5ee9
 httpx==0.27.0
 jax==0.4.31
 jaxlib==0.4.31
-numpy==1.24.0
+numpy==1.26.2
 opendp==0.10.0
 opendp-logger==0.3.0
 packaging==24.1
diff --git a/server/setup.py b/server/setup.py
index 10b83110..55a88004 100644
--- a/server/setup.py
+++ b/server/setup.py
@@ -1,10 +1,30 @@
 from setuptools import find_packages, setup
 
+
 setup(
-    name="lomas_server",
-    description="Lomas Server Package.",
-    version="0.0.1",
-    author="DSCC",
-    author_email="DSCC@bfs.admin.ch",
+    name="lomas-server",
     packages=find_packages(),
+    version="0.3.2",
+    description="Lomas server.",
+    # long_description=long_description,
+    # long_description_content_type="text/markdown",
+    url="https://github.com/dscc-admin-ch/lomas/",
+    author="Data Science Competence Center, Swiss Federal Statistical Office",
+    author_email="dscc@bfs.admin.ch",
+    license="MIT",
+    classifiers=[
+        "Development Status :: 4 - Beta",
+        "Intended Audience :: Developers",
+        "Intended Audience :: Science/Research",
+        "License :: OSI Approved :: MIT License",
+        "Natural Language :: English",
+        "Programming Language :: Python :: 3.11",
+        "Programming Language :: Python :: 3.12",
+        "Topic :: Software Development :: Libraries",
+        "Topic :: Software Development :: Libraries :: Python Modules",
+        "Topic :: Scientific/Engineering",
+        "Topic :: Security",
+    ],
+    python_requires=">=3.11, <3.13",
+    # install_requires=..
 )


Test patch:
diff --git a/.github/workflows/test_and_coverage_server.yml b/.github/workflows/test_and_coverage_server.yml
index 01083f6c..42834b57 100644
--- a/.github/workflows/test_and_coverage_server.yml
+++ b/.github/workflows/test_and_coverage_server.yml
@@ -23,7 +23,7 @@ jobs:
       uses: actions/checkout@v3
     - uses: actions/setup-python@v5
       with:
-        python-version: '3.11'
+        python-version: '3.12'
     - name: Install dependencies
       run: |
         cd server
diff --git a/.github/workflows/test_basic_server.yml b/.github/workflows/test_basic_server.yml
index 6d487e3b..e8b06bcd 100644
--- a/.github/workflows/test_basic_server.yml
+++ b/.github/workflows/test_basic_server.yml
@@ -6,13 +6,21 @@ on:
   
 jobs:
   tests:
+    name: Run tests on Python ${{ matrix.python-version }}
     runs-on: ubuntu-latest
+    strategy:
+      matrix:
+        python-version: ["3.11", "3.12"]
+
     steps:
     - name: Checkout
       uses: actions/checkout@v3
-    - uses: actions/setup-python@v5
+
+    - name: Set up Python ${{ matrix.python-version }}
+      uses: actions/setup-python@v5
       with:
-        python-version: '3.11'
+        python-version: ${{ matrix.python-version }}
+
     - name: Install dependencies
       run: |
         cd server
diff --git a/.github/workflows/test_linters.yml b/.github/workflows/test_linters.yml
index d58e6cb3..8c9c4b62 100644
--- a/.github/workflows/test_linters.yml
+++ b/.github/workflows/test_linters.yml
@@ -15,7 +15,7 @@ jobs:
       - name: Set up Python environment
         uses: actions/setup-python@v4
         with:
-          python-version: "3.11"
+          python-version: "3.12"
       - name: Install dependencies
         run: |
           pip install mypy==1.10.0
diff --git a/server/lomas_server/tests/test_api.py b/server/lomas_server/tests/test_api.py
index b784a1fc..114ee7af 100644
--- a/server/lomas_server/tests/test_api.py
+++ b/server/lomas_server/tests/test_api.py
@@ -12,24 +12,23 @@
 from opendp_logger import enable_logging
 from pymongo.database import Database
 
-from admin_database.factory import admin_database_factory
-from admin_database.utils import get_mongodb
-from app import app
-from constants import EPSILON_LIMIT, DatasetStoreType, DPLibraries
-from mongodb_admin import (
+from lomas_server.admin_database.factory import admin_database_factory
+from lomas_server.admin_database.utils import get_mongodb
+from lomas_server.app import app
+from lomas_server.constants import DPLibraries
+from lomas_server.mongodb_admin import (
     add_datasets_via_yaml,
     add_users_via_yaml,
     drop_collection,
 )
-from tests.constants import (
+from lomas_server.tests.constants import (
     ENV_MONGO_INTEGRATION,
     ENV_S3_INTEGRATION,
     TRUE_VALUES,
 )
-from utils.config import CONFIG_LOADER
-from utils.error_handler import InternalServerException
-from utils.logger import LOG
-from utils.query_examples import (
+from lomas_server.utils.config import CONFIG_LOADER, DBConfig
+from lomas_server.utils.error_handler import InternalServerException
+from lomas_server.utils.query_examples import (
     DUMMY_NB_ROWS,
     PENGUIN_DATASET,
     QUERY_DELTA,
@@ -123,18 +122,13 @@ def tearDown(self) -> None:
 
     def test_config_and_internal_server_exception(self) -> None:
         """Test set wrong configuration"""
-        config = CONFIG_LOADER.get_config()
 
         # Put unknown admin database
-        previous_admin_db = config.admin_database.db_type
-        config.admin_database.db_type = "wrong_db"
         with self.assertRaises(InternalServerException) as context:
-            admin_database_factory(config.admin_database)
+            admin_database_factory(DBConfig())
         self.assertEqual(
-            str(context.exception), "Database type wrong_db not supported."
+            str(context.exception), "Database type not supported."
         )
-        # Put original state back
-        config.admin_database.db_type = previous_admin_db
 
     def test_root(self) -> None:
         """Test root endpoint redirection to state endpoint"""
@@ -156,36 +150,6 @@ def test_state(self) -> None:
             assert response_dict["requested_by"] == self.user_name
             assert response_dict["state"]["LIVE"]
 
-    def test_memory_usage(self) -> None:
-        """Test memory usage endpoint"""
-        with TestClient(app, headers=self.headers) as client:
-            config = CONFIG_LOADER.get_config()
-            if config.dataset_store.ds_store_type == DatasetStoreType.LRU:
-                # Test before adding data
-                response = client.get(
-                    "/get_memory_usage", headers=self.headers
-                )
-                assert response.status_code == status.HTTP_200_OK
-
-                response_dict = json.loads(response.content.decode("utf8"))
-                assert response_dict["memory_usage"] == 0
-
-                # Test after adding data
-                response = client.post(
-                    "/smartnoise_sql_query",
-                    json=example_smartnoise_sql,
-                    headers=self.headers,
-                )
-                assert response.status_code == status.HTTP_200_OK
-
-                response = client.get(
-                    "/get_memory_usage", headers=self.headers
-                )
-                assert response.status_code == status.HTTP_200_OK
-
-                response_dict = json.loads(response.content.decode("utf8"))
-                assert response_dict["memory_usage"] > 0
-
     def test_get_dataset_metadata(self) -> None:
         """test_get_dataset_metadata"""
         with TestClient(app) as client:
@@ -953,9 +917,6 @@ def test_get_previous_queries(self) -> None:
             response_2 = client.post(
                 "/get_previous_queries", json=example_get_admin_db_data
             )
-            LOG.error(response_2)
-            response_dict = json.loads(response.content.decode("utf8"))
-            LOG.error(response_dict)
             assert response_2.status_code == status.HTTP_200_OK
 
             response_dict_2 = json.loads(response_2.content.decode("utf8"))
@@ -1003,25 +964,6 @@ def test_get_previous_queries(self) -> None:
                 response_dict_3["previous_queries"][1]["response"] == query_res
             )
 
-    def test_budget_over_limit(self) -> None:
-        """test_budget_over_limit"""
-        with TestClient(app, headers=self.headers) as client:
-            # Should fail: too much budget on one go
-            smartnoise_body = dict(example_smartnoise_sql)
-            smartnoise_body["epsilon"] = EPSILON_LIMIT * 2
-
-            response = client.post(
-                "/smartnoise_sql_query",
-                json=smartnoise_body,
-                headers=self.headers,
-            )
-
-            assert response.status_code == status.HTTP_422_UNPROCESSABLE_ENTITY
-            error = response.json()["detail"][0]
-            assert error["type"] == "less_than_equal"
-            assert error["loc"] == ["body", "epsilon"]
-            assert error["msg"] == "Input should be less than or equal to 5"
-
     def test_subsequent_budget_limit_logic(self) -> None:
         """test_subsequent_budget_limit_logic"""
         with TestClient(app, headers=self.headers) as client:
diff --git a/server/lomas_server/tests/test_api_diffprivlib.py b/server/lomas_server/tests/test_api_diffprivlib.py
index 6d591980..1d05021c 100644
--- a/server/lomas_server/tests/test_api_diffprivlib.py
+++ b/server/lomas_server/tests/test_api_diffprivlib.py
@@ -11,11 +11,13 @@
 from fastapi.testclient import TestClient
 from sklearn.pipeline import Pipeline
 
-from app import app
-from constants import DPLibraries
-from tests.test_api import TestRootAPIEndpoint
-from utils.logger import LOG
-from utils.query_examples import example_diffprivlib, example_dummy_diffprivlib
+from lomas_server.app import app
+from lomas_server.constants import DPLibraries
+from lomas_server.tests.test_api import TestRootAPIEndpoint
+from lomas_server.utils.query_examples import (
+    example_diffprivlib,
+    example_dummy_diffprivlib,
+)
 
 
 def validate_pipeline(response):
@@ -25,13 +27,7 @@ def validate_pipeline(response):
     assert response.status_code == status.HTTP_200_OK
     response_dict = json.loads(response.content.decode("utf8"))
 
-    if not response_dict["query_response"]["score"]:
-        LOG.error(" ************** DIFFPRIVLIB ERROR ************** ")
-        # Temporary LOGs to help understand why tests sometimes fail
-        LOG.error(response_dict)
-        LOG.error(response_dict["query_response"])
-
-    assert response_dict["query_response"]["score"]
+    assert "score" in response_dict["query_response"]  # might be 0
     assert response_dict["query_response"]["model"]
 
 
diff --git a/server/lomas_server/tests/test_api_smartnoise_synth.py b/server/lomas_server/tests/test_api_smartnoise_synth.py
index e615b4b6..cb08a75a 100644
--- a/server/lomas_server/tests/test_api_smartnoise_synth.py
+++ b/server/lomas_server/tests/test_api_smartnoise_synth.py
@@ -13,10 +13,10 @@
     OneHotEncoder,
 )
 
-from app import app
-from tests.constants import PENGUIN_COLUMNS, PUMS_COLUMNS
-from tests.test_api import TestRootAPIEndpoint
-from utils.query_examples import (
+from lomas_server.app import app
+from lomas_server.tests.constants import PENGUIN_COLUMNS, PUMS_COLUMNS
+from lomas_server.tests.test_api import TestRootAPIEndpoint
+from lomas_server.utils.query_examples import (
     example_dummy_smartnoise_synth_query,
     example_smartnoise_synth_cost,
     example_smartnoise_synth_query,
@@ -308,7 +308,7 @@ def test_smartnoise_synth_query_datetime(self) -> None:
 
             body = dict(example_smartnoise_synth_query)
             body["dataset_name"] = "BIRTHDAYS"
-            body["synth_params"]["batch_size"] = 2
+            body["synth_params"]["batch_size"] = 2  # type: ignore
             # With gan synthesizer
             response = client.post(
                 "/smartnoise_synth_query",
diff --git a/server/lomas_server/tests/test_collection_models.py b/server/lomas_server/tests/test_collection_models.py
new file mode 100644
index 00000000..75114140
--- /dev/null
+++ b/server/lomas_server/tests/test_collection_models.py
@@ -0,0 +1,228 @@
+import unittest
+
+from lomas_server.utils.collection_models import (
+    BooleanMetadata,
+    BoundedColumnMetadata,
+    CategoricalColumnMetadata,
+    DatetimeMetadata,
+    FloatMetadata,
+    IntCategoricalMetadata,
+    IntMetadata,
+    Metadata,
+    StrCategoricalMetadata,
+    StrMetadata,
+)
+
+
+class TestMetadataModel(unittest.TestCase):
+    """Tests for the input validation of dataset metadata."""
+
+    def test_categories(self) -> None:
+        """test categories validation"""
+        input_str = {
+            "type": "string",
+            "cardinality": 4,
+            "categories": ["a", "b", "c", "d"],
+        }
+
+        str_categorical_metadata = StrCategoricalMetadata.model_validate(
+            input_str
+        )
+
+        self.assertIsInstance(str_categorical_metadata, StrCategoricalMetadata)
+        self.assertIsInstance(
+            str_categorical_metadata, CategoricalColumnMetadata
+        )
+
+        input_int = {
+            "type": "int",
+            "precision": 32,
+            "cardinality": 4,
+            "categories": [5, 6, 7, 8],
+        }
+
+        int_categorical_metadata = IntCategoricalMetadata.model_validate(
+            input_int
+        )
+
+        self.assertIsInstance(int_categorical_metadata, IntCategoricalMetadata)
+        self.assertIsInstance(
+            int_categorical_metadata, CategoricalColumnMetadata
+        )
+
+    def test_categories_required(self) -> None:
+        """test categories required"""
+        input_str = {"type": "string", "cardinality": 4}
+
+        with self.assertRaises(ValueError):
+            StrCategoricalMetadata.model_validate(input_str)
+
+    def test_categories_match_type(self) -> None:
+        """Test categories match column type"""
+        input_str = {
+            "type": "string",
+            "cardinality": 4,
+            "categories": [1, 2, 3, 4],
+        }
+
+        with self.assertRaises(ValueError):
+            StrCategoricalMetadata.model_validate(input_str)
+
+    def test_categories_match_cardinality(self) -> None:
+        """Test categories match cardinality"""
+        input_int = {"type": "int", "cardinality": 4, "categories": [1, 2, 3]}
+
+        with self.assertRaises(ValueError):
+            IntCategoricalMetadata.model_validate(input_int)
+
+    def test_boolean_column(self) -> None:
+        """test_boolean_column"""
+        input_bool = {
+            "type": "boolean",
+        }
+
+        metadata = BooleanMetadata.model_validate(input_bool)
+        self.assertIsInstance(metadata, BooleanMetadata)
+
+    def test_int_column(self) -> None:
+        "test_int_column"
+        input_int = {"type": "int", "precision": 32, "lower": 0, "upper": 10}
+
+        metadata = IntMetadata.model_validate(input_int)
+        self.assertIsInstance(metadata, IntMetadata)
+        self.assertIsInstance(metadata, BoundedColumnMetadata)
+
+    def test_float_column(self) -> None:
+        "test_float_column"
+        input_float = {
+            "type": "float",
+            "precision": 64,
+            "lower": 0,
+            "upper": 10.5,
+        }
+
+        metadata = FloatMetadata.model_validate(input_float)
+        self.assertIsInstance(metadata, FloatMetadata)
+        self.assertIsInstance(metadata, BoundedColumnMetadata)
+
+    def test_datetime_column(self) -> None:
+        "test_datetime_column"
+        input_datetime = {
+            "type": "datetime",
+            "lower": "2000-01-01",
+            "upper": "2001-01-01",
+        }
+
+        metadata = DatetimeMetadata.model_validate(input_datetime)
+        self.assertIsInstance(metadata, DatetimeMetadata)
+        self.assertIsInstance(metadata, BoundedColumnMetadata)
+
+        input_datetime = {
+            "type": "datetime",
+            "lower": "2001-01-01",
+            "upper": "2000-01-01",
+        }
+        with self.assertRaises(ValueError):
+            DatetimeMetadata.model_validate(input_datetime)
+
+    def test_precision(self) -> None:
+        "test precision can only be 32 or 64"
+        input_int = {"type": int, "precision": 20, "lower": 0, "upper": 10}
+
+        with self.assertRaises(ValueError):
+            IntMetadata.model_validate(input_int)
+
+    def test_lower_upper_bounded(self) -> None:
+        """Test lower is smaller than upper and of right type"""
+        input_int = {"type": "int", "precision": 32, "lower": 0, "upper": -1}
+
+        with self.assertRaises(ValueError):
+            IntMetadata.model_validate(input_int)
+
+        input_int["upper"] = 10.5
+        with self.assertRaises(ValueError):
+            IntMetadata.model_validate(input_int)
+
+    def test_standard_metadata_fields(self):
+        """Test standard metadata fields"""
+        input_metadata = {
+            "max_ids": 1,
+            "rows": 100,
+            "row_privacy": False,
+            "columns": {},
+        }
+
+        metadata = Metadata.model_validate(input_metadata)
+        self.assertIsInstance(metadata, Metadata)
+        self.assertFalse(metadata.censor_dims)
+
+        input_metadata["censor_dims"] = True
+        metadata = Metadata.model_validate(input_metadata)
+        self.assertTrue(metadata.censor_dims)
+
+        input_metadata["rows"] = 0
+        with self.assertRaises(ValueError):
+            Metadata.model_validate(input_metadata)
+
+        input_metadata["rows"] = 100
+        input_metadata["max_ids"] = -1
+        with self.assertRaises(ValueError):
+            Metadata.model_validate(input_metadata)
+
+    def test_metadata_columns_discriminator(self) -> None:
+        """Test metadata column discriminator"""
+        input_metadata = {
+            "max_ids": 1,
+            "rows": 100,
+            "row_privacy": False,
+            "columns": {
+                "str": {"type": "string"},
+                "str_cat": {
+                    "type": "string",
+                    "cardinality": 2,
+                    "categories": ["a", "b"],
+                },
+                "int": {
+                    "type": "int",
+                    "precision": 32,
+                    "lower": 0,
+                    "upper": 10,
+                },
+                "int_cat": {
+                    "type": "int",
+                    "cardinality": 2,
+                    "categories": [1, 2],
+                    "precision": 32,
+                },
+                "float": {
+                    "type": "float",
+                    "precision": 64,
+                    "lower": 0.5,
+                    "upper": 1.2,
+                },
+                "boolean": {"type": "boolean"},
+                "datetime": {
+                    "type": "datetime",
+                    "lower": "2000-01-01",
+                    "upper": "2001-01-01",
+                },
+            },
+        }
+
+        metadata = Metadata.model_validate(input_metadata)
+
+        self.assertIsInstance(metadata.columns["str"], StrMetadata)
+        self.assertIsInstance(
+            metadata.columns["str_cat"], StrCategoricalMetadata
+        )
+        self.assertIsInstance(metadata.columns["int"], IntMetadata)
+        self.assertIsInstance(
+            metadata.columns["int_cat"], IntCategoricalMetadata
+        )
+        self.assertIsInstance(metadata.columns["float"], FloatMetadata)
+        self.assertIsInstance(metadata.columns["boolean"], BooleanMetadata)
+        self.assertIsInstance(metadata.columns["datetime"], DatetimeMetadata)
+
+        input_metadata["columns"]["some_type"] = {"type": "some_type"}  # type: ignore
+        with self.assertRaises(ValueError):
+            Metadata.model_validate(input_metadata)
diff --git a/server/lomas_server/tests/test_configs/test_config.yaml b/server/lomas_server/tests/test_configs/test_config.yaml
index 6c51aa3f..aa8eac8b 100644
--- a/server/lomas_server/tests/test_configs/test_config.yaml
+++ b/server/lomas_server/tests/test_configs/test_config.yaml
@@ -14,8 +14,6 @@ runtime_args:
     admin_database:
       db_type: "yaml"
       db_file: "tests/test_data/local_db_file.yaml"
-    dataset_store:
-      ds_store_type: "basic"
     dp_libraries:
       opendp:
         contrib: True
diff --git a/server/lomas_server/tests/test_configs/test_config_mongo.yaml b/server/lomas_server/tests/test_configs/test_config_mongo.yaml
index f2aaf347..fde08049 100644
--- a/server/lomas_server/tests/test_configs/test_config_mongo.yaml
+++ b/server/lomas_server/tests/test_configs/test_config_mongo.yaml
@@ -16,9 +16,6 @@ runtime_args:
       address: "127.0.0.1"
       port: 27017
       db_name: "defaultdb"
-    dataset_store:
-      ds_store_type: "LRU_cache"
-      max_memory_usage: 512
     dp_libraries:
       opendp:
         contrib: True
diff --git a/server/lomas_server/tests/test_data/metadata/penguin_metadata.yaml b/server/lomas_server/tests/test_data/metadata/penguin_metadata.yaml
index c8fdbd6e..b9770ddb 100644
--- a/server/lomas_server/tests/test_data/metadata/penguin_metadata.yaml
+++ b/server/lomas_server/tests/test_data/metadata/penguin_metadata.yaml
@@ -18,18 +18,22 @@ columns:
     - Dream
   bill_length_mm:
     type: float
+    precision: 64
     lower: 30.0
     upper: 65.0
   bill_depth_mm:
     type: float
+    precision: 64
     lower: 13.0
     upper: 23.0
   flipper_length_mm:
     type: float
+    precision: 64
     lower: 150.0
     upper: 250.0
   body_mass_g:
     type: float
+    precision: 64
     lower: 2000.0
     upper: 7000.0
   sex:
diff --git a/server/lomas_server/tests/test_data/metadata/pums_metadata.yaml b/server/lomas_server/tests/test_data/metadata/pums_metadata.yaml
index e89d55b8..a9a45064 100644
--- a/server/lomas_server/tests/test_data/metadata/pums_metadata.yaml
+++ b/server/lomas_server/tests/test_data/metadata/pums_metadata.yaml
@@ -5,16 +5,24 @@ columns:
   age:
     private_id: true # for testing purposes
     type: int
+    precision: 32
+    lower: 0
+    upper: 120
   sex:
     type: boolean
   educ:
     type: int
-    cardinality: 14
+    precision: 32
+    cardinality: 16
+    categories: [ 9,  1, 11,  5,  8, 13, 10,  7, 14,  4, 15, 16,  3, 12,  2,  6]
   race:
     type: int
+    precision: 32
     cardinality: 6
+    categories: [1, 3, 4, 2, 6, 5]
   income:
     type: float
+    precision: 64
     lower: 0.0
     upper: 500_000.0
   married:
diff --git a/server/lomas_server/tests/test_dummy_generation.py b/server/lomas_server/tests/test_dummy_generation.py
index e057056a..92a30782 100644
--- a/server/lomas_server/tests/test_dummy_generation.py
+++ b/server/lomas_server/tests/test_dummy_generation.py
@@ -1,7 +1,9 @@
 import unittest
+from typing import Any
 
-from dp_queries.dummy_dataset import make_dummy_dataset
-from utils.query_examples import DUMMY_NB_ROWS, DUMMY_SEED
+from lomas_server.dp_queries.dummy_dataset import make_dummy_dataset
+from lomas_server.utils.collection_models import Metadata
+from lomas_server.utils.query_examples import DUMMY_NB_ROWS, DUMMY_SEED
 
 
 class TestMakeDummyDataset(unittest.TestCase):
@@ -9,29 +11,28 @@ class TestMakeDummyDataset(unittest.TestCase):
     Tests for the generation of dummy datasets.
     """
 
-    def test_cardinality_column(self) -> None:
-        """test_cardinality_column"""
-        metadata = {
-            "columns": {
-                "col_card_cat": {  # cardinality + categories
-                    "type": "string",
-                    "cardinality": 3,
-                    "categories": ["x", "y", "z"],
-                },
-                "col_card_no_cat": {  # cardinality, no categories
-                    "type": "string",
-                    "cardinality": 3,
-                },
-                "col_no_card": {  # no cardinality
-                    "type": "string",
-                },
+    metadata: dict[str, Any] = {
+        "max_ids": 1,
+        "rows": 100,
+        "row_privacy": True,
+        "columns": {},
+    }
+
+    def test_categorical_column(self) -> None:
+        """test_categorical_column"""
+        self.metadata["columns"] = {
+            "col_card_cat": {  # cardinality + categories
+                "type": "string",
+                "cardinality": 3,
+                "categories": ["x", "y", "z"],
             }
         }
+        metadata = Metadata.model_validate(self.metadata)
         df = make_dummy_dataset(metadata)
 
         # Test shape
         self.assertEqual(df.shape[0], DUMMY_NB_ROWS)
-        self.assertEqual(df.shape[1], len(metadata["columns"].keys()))
+        self.assertEqual(df.shape[1], 1)
 
         # Test cardinality type and categories
         self.assertIn("col_card_cat", df.columns)
@@ -39,24 +40,14 @@ def test_cardinality_column(self) -> None:
         self.assertEqual(set(df["col_card_cat"].values), {"x", "y", "z"})
         assert isinstance(df["col_card_cat"], object)
 
-        # Test cardinality type and no categories
-        self.assertIn("col_card_no_cat", df.columns)
-        self.assertEqual(set(df["col_card_no_cat"].values), {"a", "b", "c"})
-        assert isinstance(df["col_card_no_cat"], object)
-
-        # Test cardinality type and no categories
-        self.assertIn("col_no_card", df.columns)
-        self.assertTrue(
-            df["col_no_card"].apply(lambda x: isinstance(x, str)).all()
-        )
-
     def test_boolean_column(self) -> None:
         """test_boolean_column"""
 
         # Test a boolean column
-        metadata = {
-            "columns": {"col_bool": {"type": "boolean", "nullable": True}}
+        self.metadata["columns"] = {
+            "col_bool": {"type": "boolean", "nullable": True}
         }
+        metadata = Metadata.model_validate(self.metadata)
         df = make_dummy_dataset(metadata)
 
         # Test length
@@ -70,19 +61,19 @@ def test_float_column(self) -> None:
         """test_float_column"""
         lower_bound = 10.0
         upper_bound = 20.0
-        metadata = {
-            "columns": {
-                "col_float": {
-                    "type": "float",
-                    "upper": upper_bound,
-                    "lower": lower_bound,
-                }
+        self.metadata["columns"] = {
+            "col_float": {
+                "type": "float",
+                "precision": 32,
+                "upper": upper_bound,
+                "lower": lower_bound,
             }
         }
+        metadata = Metadata.model_validate(self.metadata)
         df = make_dummy_dataset(metadata)
 
         # Test col generated is of type float
-        self.assertEqual(df.col_float.dtypes.name, "float64")
+        self.assertEqual(df.col_float.dtypes.name, "float32")
 
         # Test within bounds
         self.assertTrue((df["col_float"] >= lower_bound).all())
@@ -92,19 +83,19 @@ def test_int_column(self) -> None:
         """test_int_column"""
         lower_bound = 100
         upper_bound = 120
-        metadata = {
-            "columns": {
-                "col_int": {
-                    "type": "int",
-                    "upper": upper_bound,
-                    "lower": lower_bound,
-                }
+        self.metadata["columns"] = {
+            "col_int": {
+                "type": "int",
+                "precision": 64,
+                "upper": upper_bound,
+                "lower": lower_bound,
             }
         }
+        metadata = Metadata.model_validate(self.metadata)
         df = make_dummy_dataset(metadata)
 
         # Test col generated is of type int
-        self.assertIn(df.col_int.dtypes.name, ["int32", "int64"])
+        self.assertIn(df.col_int.dtypes.name, ["int64"])
 
         # Test within bounds
         self.assertTrue((df["col_int"] >= lower_bound).all())
@@ -112,7 +103,14 @@ def test_int_column(self) -> None:
 
     def test_datetime_column(self) -> None:
         """test_datetime_column"""
-        metadata = {"columns": {"col_datetime": {"type": "datetime"}}}
+        self.metadata["columns"] = {
+            "col_datetime": {
+                "type": "datetime",
+                "lower": "2000-01-01",
+                "upper": "2010-01-01",
+            }
+        }
+        metadata = Metadata.model_validate(self.metadata)
         df = make_dummy_dataset(metadata)
 
         # Test col generated is of type datetime
@@ -123,9 +121,15 @@ def test_datetime_column(self) -> None:
 
     def test_nullable_column(self) -> None:
         """test_nullable_column"""
-        metadata = {
-            "columns": {"col_nullable": {"type": "datetime", "nullable": True}}
+        self.metadata["columns"] = {
+            "col_nullable": {
+                "type": "datetime",
+                "nullable": True,
+                "lower": "2000-01-01",
+                "upper": "2010-01-01",
+            }
         }
+        metadata = Metadata.model_validate(self.metadata)
         df = make_dummy_dataset(metadata)
 
         # Should have null values
@@ -134,7 +138,16 @@ def test_nullable_column(self) -> None:
     def test_seed(self) -> None:
         """test_seed"""
         # Test the behavior with different seeds
-        metadata = {"columns": {"col_int": {"type": "int", "nullable": True}}}
+        self.metadata["columns"] = {
+            "col_int": {
+                "type": "int",
+                "nullable": True,
+                "precision": 32,
+                "lower": 0,
+                "upper": 100,
+            }
+        }
+        metadata = Metadata.model_validate(self.metadata)
         seed1 = DUMMY_SEED
         seed2 = DUMMY_SEED + 1
 
@@ -148,18 +161,19 @@ def test_seed(self) -> None:
         df1_copy = make_dummy_dataset(metadata, seed=seed1)
         self.assertTrue(df1.equals(df1_copy))
 
-    def test_unknown_column(self) -> None:
-        """test_unknown_column"""
-        metadata = {
-            "columns": {
-                "col_bool": {"type": "boolean", "nullable": True},
-                "col_unknown": {"type": "unknown"},
-            }
-        }
-        df = make_dummy_dataset(metadata)
-
-        # Test col generated
-        self.assertIn("col_bool", df.columns)
-
-        # Test col not generated
-        self.assertNotIn("col_unknown", df.columns)
+    # TODO maybe remove this, see issue #335
+    # def test_unknown_column(self) -> None:
+    #     """test_unknown_column"""
+    #     metadata = {
+    #         "columns": {
+    #             "col_bool": {"type": "boolean", "nullable": True},
+    #             "col_unknown": {"type": "unknown"},
+    #         }
+    #     }
+    #     df = make_dummy_dataset(metadata)
+
+    #     # Test col generated
+    #     self.assertIn("col_bool", df.columns)
+
+    #     # Test col not generated
+    #     self.assertNotIn("col_unknown", df.columns)
diff --git a/server/lomas_server/tests/test_mongodb_admin.py b/server/lomas_server/tests/test_mongodb_admin.py
index f0738661..375cbc3f 100644
--- a/server/lomas_server/tests/test_mongodb_admin.py
+++ b/server/lomas_server/tests/test_mongodb_admin.py
@@ -1,18 +1,17 @@
 import os
 import unittest
-from types import SimpleNamespace
 from typing import Dict
 
 import boto3
 import yaml
 from pymongo import MongoClient
 
-from admin_database.utils import (
+from lomas_server.admin_database.utils import (
     add_demo_data_to_mongodb_admin,
     get_mongodb_url,
 )
-from constants import PrivateDatabaseType
-from mongodb_admin import (
+from lomas_server.constants import PrivateDatabaseType
+from lomas_server.mongodb_admin import (
     add_dataset,
     add_dataset_to_user,
     add_datasets_via_yaml,
@@ -34,13 +33,13 @@
     set_budget_field,
     set_may_query,
 )
-from tests.constants import (
+from lomas_server.tests.constants import (
     ENV_MONGO_INTEGRATION,
     ENV_S3_INTEGRATION,
     FALSE_VALUES,
     TRUE_VALUES,
 )
-from utils.config import CONFIG_LOADER, get_config
+from lomas_server.utils.config import CONFIG_LOADER, MongoDBConfig, get_config
 
 
 @unittest.skipIf(
@@ -68,9 +67,15 @@ def setUpClass(cls) -> None:
             secrets_path="tests/test_configs/test_secrets.yaml",
         )
 
-        db_args = SimpleNamespace(**vars(get_config().admin_database))
-        db_url = get_mongodb_url(db_args)
-        cls.db = MongoClient(db_url)[db_args.db_name]
+        # Access to MongoDB
+        admin_config = get_config().admin_database
+        if isinstance(admin_config, MongoDBConfig):
+            mongo_config = admin_config
+        else:
+            raise TypeError("Loaded config does not contain a MongoDBConfig.")
+
+        db_url = get_mongodb_url(mongo_config)
+        cls.db = MongoClient(db_url)[mongo_config.db_name]
 
     def tearDown(self) -> None:
         """Drop all data from database"""
diff --git a/server/lomas_server/tests/test_mongodb_admin_cli.py b/server/lomas_server/tests/test_mongodb_admin_cli.py
index 1cef92b3..fae723c6 100644
--- a/server/lomas_server/tests/test_mongodb_admin_cli.py
+++ b/server/lomas_server/tests/test_mongodb_admin_cli.py
@@ -1,16 +1,15 @@
 import os
 import subprocess
 import unittest
-from types import SimpleNamespace
 from typing import List
 
 import yaml
 from pymongo import MongoClient
 
-from admin_database.utils import get_mongodb_url
-from constants import PrivateDatabaseType
-from tests.constants import ENV_MONGO_INTEGRATION
-from utils.config import CONFIG_LOADER, get_config
+from lomas_server.admin_database.utils import get_mongodb_url
+from lomas_server.constants import PrivateDatabaseType
+from lomas_server.tests.constants import ENV_MONGO_INTEGRATION
+from lomas_server.utils.config import CONFIG_LOADER, MongoDBConfig, get_config
 
 
 @unittest.skipIf(
@@ -39,12 +38,16 @@ def setUpClass(cls) -> None:
         )
 
         # Access to MongoDB
-        db_args = SimpleNamespace(**vars(get_config().admin_database))
-        db_url = get_mongodb_url(db_args)
-        cls.db = MongoClient(db_url)[db_args.db_name]
+        admin_config = get_config().admin_database
+        if isinstance(admin_config, MongoDBConfig):
+            mongo_config: MongoDBConfig = admin_config
+        else:
+            raise TypeError("Loaded config does not contain a MongoDBConfig.")
+
+        db_url = get_mongodb_url(mongo_config)
+        cls.db = MongoClient(db_url)[mongo_config.db_name]
 
         # CLI args to connect to DB
-        admin_config = get_config().admin_database
         cls.db_connection_cli = [
             "--username",
             admin_config.username,
diff --git a/server/run_basic_tests.sh b/server/run_basic_tests.sh
index c53fe198..f15c80b2 100755
--- a/server/run_basic_tests.sh
+++ b/server/run_basic_tests.sh
@@ -1,6 +1,7 @@
 #!/bin/bash
 
 cd ./lomas_server
+export PYTHONPATH=$(pwd)/..:$PYTHONPATH
 python -m unittest discover -s .
 ret=$?
 cd ..
diff --git a/server/run_tests_and_coverage.sh b/server/run_tests_and_coverage.sh
index 57c05e88..275aa294 100755
--- a/server/run_tests_and_coverage.sh
+++ b/server/run_tests_and_coverage.sh
@@ -5,7 +5,7 @@ docker compose -f docker-compose-test.yml up --detach
 sleep 15
 
 cd ./lomas_server
-
+export PYTHONPATH=$(pwd)/..:$PYTHONPATH
 
 # "mongodb", "LRU_cache", production mode, "jitter"
 export LOMAS_TEST_MONGO_INTEGRATION=1 # tests with mongodb available