Initial support for metrics

TODO:
- [ ] support globalLabels
- [ ] check if int / float is right
- [ ] test multiple client instances
- [ ] document

databases/templates/schema.prisma.jinja2

@@ -39,20 +39,20 @@ model Profile {
 }
 
 model Post {
-  id          String   @id @default(uuid())
-  created_at  DateTime @default(now())
-  updated_at  DateTime @updatedAt
+  id          String     @id @default(uuid())
+  created_at  DateTime   @default(now())
+  updated_at  DateTime   @updatedAt
   title       String
   description String?
-  published   Boolean  @default(false)
-  views       Int      @default(0)
-  author      User?    @relation(fields: [author_id], references: [id])
+  published   Boolean    @default(false)
+  views       Int        @default(0)
+  author      User?      @relation(fields: [author_id], references: [id])
   author_id   String?
-  categories Category[]
+  categories  Category[]
 }
 
 model Category {
-  id          String   @id @default(uuid())
+  id    String @id @default(uuid())
   posts Post[]
   name  String
 }
@@ -80,10 +80,10 @@ model Types {
 
   // TODO: optional for these too
   {% if config.supports_feature('enum') %}
-  enum     Role   @default(USER)
+  enum     Role  @default(USER)
   {% endif %}
   {% if config.supports_feature('json') %}
-  json_obj Json?  @default("{}")
+  json_obj Json? @default("{}")
   {% endif %}
 }
 
@@ -123,6 +123,7 @@ model ListsDefaults {
   roles        Role[]     @default([USER])
   {% endif %}
 }
+
 {% endif %}
 
 // these models are here for testing different combinations of unique constraints
@@ -173,6 +174,7 @@ model Unique6 {
 
   @@unique([name, role])
 }
+
 {% endif %}
 
 model Id1 {
@@ -211,6 +213,7 @@ model Id5 {
 
   @@unique([name, role])
 }
+
 {% endif %}
 
 {% if config.supports_feature('enum') %}
@@ -219,4 +222,5 @@ enum Role {
   ADMIN
   EDITOR
 }
+
 {% endif %}

databases/tests/test_metrics.py

@@ -0,0 +1,51 @@
+import pytest
+
+from prisma import Prisma
+from prisma._compat import model_json
+
+
+@pytest.mark.asyncio
+async def test_prometheus(client: Prisma) -> None:
+    """Metrics can be returned in Prometheus format"""
+    await client.user.create(data={'name': 'Robert'})
+
+    metrics = await client.get_metrics(format='prometheus')
+
+    assert 'prisma_client_queries_total' in metrics
+    assert 'prisma_datasource_queries_total' in metrics
+    assert 'prisma_client_queries_active' in metrics
+    assert 'prisma_client_queries_duration_histogram_ms_bucket' in metrics
+
+
+@pytest.mark.asyncio
+async def test_json_string(client: Prisma) -> None:
+    """Metrics can be serlialized to JSON"""
+    await client.user.create(data={'name': 'Robert'})
+
+    metrics = await client.get_metrics()
+    assert isinstance(model_json(metrics), str)
+
+
+@pytest.mark.asyncio
+async def test_json(client: Prisma) -> None:
+    """Metrics returned in the JSON format"""
+    await client.user.create(data={'name': 'Robert'})
+
+    metrics = await client.get_metrics(format='json')
+
+    assert len(metrics.counters) > 0
+    assert metrics.counters[0].value > 0
+
+    assert len(metrics.gauges) > 0
+    gauge = next(filter(lambda g: g.key == 'prisma_pool_connections_open', metrics.gauges))
+    assert gauge.value > 0
+
+    assert len(metrics.histograms) > 0
+    assert metrics.histograms[0].value.sum > 0
+    assert metrics.histograms[0].value.count > 0
+
+    assert len(metrics.histograms[0].value.buckets) > 0
+
+    for bucket in metrics.histograms[0].value.buckets:
+        assert bucket.max_value >= 0
+        assert bucket.total_count >= 0

src/prisma/__init__.py

@@ -13,6 +13,11 @@
 from . import errors as errors
 from .validator import *
 from ._types import PrismaMethod as PrismaMethod
+from ._metrics import (
+    Metric as Metric,
+    Metrics as Metrics,
+    MetricHistogram as MetricHistogram,
+)
 
 
 try:

src/prisma/_metrics.py

@@ -0,0 +1,50 @@
+# copied from https://github.com/prisma/prisma/blob/23d5ef0672372035a84552b6b457197ca19f486d/packages/client/src/runtime/core/engines/common/types/Metrics.ts
+from __future__ import annotations
+
+from typing import Generic, List, TypeVar, Dict, NamedTuple
+
+from pydantic import BaseModel
+
+from ._compat import GenericModel, model_rebuild
+
+
+__all__ = (
+    'Metrics',
+    'Metric',
+    'MetricHistogram',
+)
+
+
+_T = TypeVar('_T')
+
+
+# TODO: check if int / float is right
+
+
+class Metrics(BaseModel):
+    counters: List[Metric[int]]   # TODO
+    gauges: List[Metric[float]]  # TODO
+    histograms: List[Metric[MetricHistogram]]
+
+
+class Metric(GenericModel, Generic[_T]):
+    key: str
+    value: _T
+    labels: Dict[str, str]
+    description: str
+
+
+class MetricHistogram(BaseModel):
+    sum: float  # TODO
+    count: int  # TODO
+    buckets: List[HistogramBucket]
+
+
+class HistogramBucket(NamedTuple):
+    max_value: float   # TODO
+    total_count: int   # TODO
+
+
+model_rebuild(Metric)
+model_rebuild(Metrics)
+model_rebuild(MetricHistogram)

src/prisma/generator/models.py

@@ -480,6 +480,9 @@ def warn_binary_targets(
 
         return targets
 
+    def has_preview_feature(self, feature: str) -> bool:
+        return feature in self.preview_features
+
 
 class ValueFromEnvVar(BaseModel):
     value: str

src/prisma/generator/templates/client.py.jinja

@@ -11,7 +11,7 @@ from types import TracebackType
 from pydantic import BaseModel
 
 from . import types, models, errors, actions
-from .types import DatasourceOverride, HttpConfig
+from .types import DatasourceOverride, HttpConfig, MetricsFormat
 from ._types import BaseModelT, PrismaMethod
 from .bases import _PrismaModel
 from .engine import AbstractEngine, QueryEngine, TransactionId
@@ -20,6 +20,7 @@ from .generator.models import EngineType, OptionalValueFromEnvVar, BinaryPaths
 from ._compat import removeprefix, model_parse
 from ._constants import DEFAULT_CONNECT_TIMEOUT, DEFAULT_TX_MAX_WAIT, DEFAULT_TX_TIMEOUT
 from ._raw_query import deserialize_raw_results
+from ._metrics import Metrics
 
 __all__ = (
     'ENGINE_TYPE',
@@ -423,6 +424,30 @@ class Prisma:
         """Returns True if the client is wrapped within a transaction"""
         return self._tx_id is not None
 
+    # TODO: global labels option
+    @overload
+    {{ maybe_async_def }}get_metrics(self, format: Literal['json'] = 'json') -> Metrics:
+        ...
+
+    @overload
+    {{ maybe_async_def }}get_metrics(self, format: Literal['prometheus']) -> str:
+        ...
+
+    {{ maybe_async_def }}get_metrics(self, format: MetricsFormat = 'json') -> str | Metrics:
+        """Metrics give you a detailed insight into how the Prisma Client interacts with your database.
+
+        You can retrieve metrics in either JSON or Prometheus formats.
+
+        For more details see https://www.prisma.io/docs/concepts/components/prisma-client/metrics.
+        """
+        response = {{ maybe_await }}self._engine.metrics(format=format)
+        if format == 'prometheus':
+            # For the prometheus format we return the response as-is
+            assert isinstance(response, str)
+            return response
+
+        return model_parse(Metrics, response)
+
     # TODO: don't return Any
     {{ maybe_async_def }}_execute(
         self,

src/prisma/generator/templates/engine/abstract.py.jinja

@@ -5,7 +5,7 @@
 from abc import ABC, abstractmethod
 from datetime import timedelta
 from ._types import TransactionId
-from ..types import DatasourceOverride
+from ..types import DatasourceOverride, MetricsFormat
 from .._compat import get_running_loop
 from .._constants import DEFAULT_CONNECT_TIMEOUT
 
@@ -77,3 +77,17 @@ class AbstractEngine(ABC):
     {{ maybe_async_def }}rollback_transaction(self, tx_id: TransactionId) -> None:
         """Rollback an interactive transaction, the given transaction will no longer be usable"""
         ...
+
+    @overload
+    @abstractmethod
+    {{ maybe_async_def }}metrics(self, *, format: Literal['json']) -> dict[str, Any]:
+        ...
+
+    @overload
+    @abstractmethod
+    {{ maybe_async_def }}metrics(self, *, format: Literal['prometheus']) -> str:
+        ...
+
+    @abstractmethod
+    {{ maybe_async_def }}metrics(self, *, format: MetricsFormat) -> str | dict[str, Any]:
+        ...

src/prisma/generator/templates/engine/http.py.jinja

@@ -54,25 +54,28 @@ class HTTPEngine(AbstractEngine):
         if self.session and not self.session.closed:
             {{ maybe_await }}self.session.close()
 
+    # TODO: improve return types
     {{ maybe_async_def }}request(
         self,
         method: Method,
         path: str,
         *,
         content: Any = None,
         headers: Optional[Dict[str, str]] = None,
+        parse_response: bool = True,
     ) -> Any:
         if self.url is None:
             raise errors.NotConnectedError('Not connected to the query engine')
 
         kwargs = {
             'headers': {
-                'Content-Type': 'application/json',
-                'Accept': 'application/json',
                 **self.headers,
             }
         }
 
+        if parse_response:
+            kwargs['headers']['Accept'] = 'application/json'
+
         if headers is not None:
             kwargs['headers'].update(headers)
 
@@ -88,6 +91,15 @@ class HTTPEngine(AbstractEngine):
         log.debug('%s %s returned status %s', method, url, resp.status)
 
         if 300 > resp.status >= 200:
+            # In certain cases we just want to return the response content as-is.
+            #
+            # This is useful for metrics which can be returned in a Prometheus format
+            # which is incompatible with JSON.
+            if not parse_response:
+                text = {{ maybe_await }}resp.text()
+                log.debug('%s %s returned text: %s', method, url, text)
+                return text
+
             response = {{ maybe_await }}resp.json()
             log.debug('%s %s returned %s', method, url, response)
 

src/prisma/generator/templates/engine/query.py.jinja

@@ -19,7 +19,7 @@ from .. import config
 from ..utils import DEBUG
 from ..binaries import platform
 from ..utils import time_since, _env_bool
-from ..types import DatasourceOverride
+from ..types import DatasourceOverride, MetricsFormat
 from ..builder import dumps
 from .._constants import DEFAULT_CONNECT_TIMEOUT
 from ._types import TransactionId
@@ -136,7 +136,13 @@ class QueryEngine(HTTPEngine):
         if self._log_queries:
             env.update(LOG_QUERIES='y')
 
-        args: List[str] = [str(file.absolute()), '-p', str(port), '--enable-raw-queries']
+        args: List[str] = [
+            str(file.absolute()),
+            '-p',
+            str(port),
+            '--enable-metrics',
+            '--enable-raw-queries',
+        ]
         if _env_bool('__PRISMA_PY_PLAYGROUND'):
             env.update(RUST_LOG='info')
             args.append('--enable-playground')
@@ -223,6 +229,21 @@ class QueryEngine(HTTPEngine):
             'POST', f'/transaction/{tx_id}/rollback'
         )
 
+    @overload
+    {{ maybe_async_def }}metrics(self, *, format: Literal['json']) -> dict[str, Any]:
+        ...
+
+    @overload
+    {{ maybe_async_def }}metrics(self, *, format: Literal['prometheus']) -> str:
+        ...
+
+    {{ maybe_async_def }}metrics(self, *, format: MetricsFormat) -> str | dict[str, Any]:
+        return {{ maybe_await }}self.request(
+            'GET',
+            f'/metrics?format={format}',
+            parse_response=format == 'json',
+        )
+
 # black does not respect the fmt: off comment without this
 # fmt: on
 

src/prisma/generator/templates/types.py.jinja

@@ -72,6 +72,8 @@ from .utils import _NoneType
 SortMode = Literal['default', 'insensitive']
 SortOrder = Literal['asc', 'desc']
 
+MetricsFormat = Literal['json', 'prometheus']
+
 
 class _DatasourceOverrideOptional(TypedDict, total=False):
     env: str