1717import gzip
1818import io
1919import itertools
20+ import os .path
2021from abc import ABC
2122from http import HTTPStatus
2223from typing import List
2324
25+ import grpc
2426import requests
2527from requests import HTTPError
2628
29+ from opengemini_client import grpc_client
2730from opengemini_client .client import Client
2831from opengemini_client .measurement import Measurement , MeasurementCondition
2932from opengemini_client .models import Config , BatchPoints , Query , QueryResult , Series , SeriesResult , RpConfig , \
30- ValuesResult , KeyValue
33+ ValuesResult , KeyValue , AuthConfig
3134from opengemini_client .url_const import UrlConst
3235from opengemini_client .models import AuthType , TlsConfig
3336
3437
38+ def check_auth_config (auth_config : AuthConfig ):
39+ if auth_config is not None :
40+ if auth_config .auth_type == AuthType .PASSWORD :
41+ if len (auth_config .username ) == 0 :
42+ raise ValueError ("invalid auth config due to empty username" )
43+ if len (auth_config .password ) == 0 :
44+ raise ValueError ("invalid auth config due to empty password" )
45+ if auth_config .auth_type == AuthType .TOKEN and len (auth_config .token ) == 0 :
46+ raise ValueError ("invalid auth config due to empty token" )
47+
48+
3549def check_config (config : Config ):
3650 if len (config .address ) == 0 :
3751 raise ValueError ("must have at least one address" )
3852
39- if config .auth_config is not None :
40- if config .auth_config .auth_type == AuthType .PASSWORD :
41- if len (config .auth_config .username ) == 0 :
42- raise ValueError ("invalid auth config due to empty username" )
43- if len (config .auth_config .password ) == 0 :
44- raise ValueError ("invalid auth config due to empty password" )
45- if config .auth_config .auth_type == AuthType .TOKEN and len (config .auth_config .token ) == 0 :
46- raise ValueError ("invalid auth config due to empty token" )
53+ check_auth_config (config .auth_config )
4754
4855 if config .tls_enabled and config .tls_config is None :
4956 config .tls_config = TlsConfig ()
@@ -60,6 +67,17 @@ def check_config(config: Config):
6067 if config .connection_timeout is None or config .connection_timeout <= datetime .timedelta (seconds = 0 ):
6168 config .connection_timeout = datetime .timedelta (seconds = 10 )
6269
70+ if config .grpc_config is None :
71+ return config
72+
73+ if len (config .grpc_config .address ) == 0 :
74+ raise ValueError ("grpc config must have at least one address" )
75+
76+ check_auth_config (config .grpc_config .auth_config )
77+
78+ if config .grpc_config .tls_enable and config .grpc_config .tls_config is None :
79+ config .grpc_config .tls_config = TlsConfig ()
80+
6381 return config
6482
6583
@@ -95,6 +113,9 @@ def __init__(self, config: Config):
95113 self .session .verify = config .tls_config .ca_file
96114 self .endpoints = [f"{ protocol } { addr .host } :{ addr .port } " for addr in config .address ]
97115 self .endpoints_iter = itertools .cycle (self .endpoints )
116+ if self .config .grpc_config is not None :
117+ self .grpc_endpoints = [f"{ addr .host } :{ addr .port } " for addr in config .grpc_config .address ]
118+ self .grpc_endpoints_iter = itertools .cycle (self .grpc_endpoints )
98119
99120 def close (self ):
100121 self .session .close ()
@@ -108,6 +129,31 @@ def __exit__(self, _exc_type, _exc_val, _exc_tb):
108129 def _get_server_url (self ):
109130 return next (self .endpoints_iter )
110131
132+ def _get_grpc_server_url (self ):
133+ return next (self .grpc_endpoints_iter )
134+
135+ def _get_grpc_channel (self ):
136+ server_url = self ._get_grpc_server_url ()
137+ if self .config .grpc_config .tls_enable is False :
138+ return grpc .insecure_channel (server_url )
139+
140+ root_certificates = None
141+ private_key = None
142+ certificate_chain = None
143+ if os .path .exists (self .config .grpc_config .tls_config .ca_file ):
144+ with open (self .config .grpc_config .tls_config .ca_file , 'rb' ) as fd :
145+ root_certificates = fd .read ()
146+ if os .path .exists (self .config .grpc_config .tls_config .cert_file ):
147+ with open (self .config .grpc_config .tls_config .cert_file , 'rb' ) as fd :
148+ certificate_chain = fd .read ()
149+ if os .path .exists (self .config .grpc_config .tls_config .key_file ):
150+ with open (self .config .grpc_config .tls_config .key_file , 'rb' ) as fd :
151+ private_key = fd .read ()
152+ return grpc .secure_channel (
153+ target = server_url ,
154+ credentials = grpc .ssl_channel_credentials (root_certificates , private_key , certificate_chain )
155+ )
156+
111157 def _update_headers (self , method , url_path , headers = None ) -> dict :
112158 if headers is None :
113159 headers = {}
@@ -191,6 +237,25 @@ def write_batch_points(self, database: str, batch_points: BatchPoints):
191237 return
192238 raise HTTPError (f"write_batch_points error resp, code: { resp .status_code } , body: { resp .text } " )
193239
240+ def write_by_grpc (self , database : str , batch_points : BatchPoints , rp : str = '' ):
241+ username = ''
242+ password = ''
243+ if self .config .grpc_config .auth_config is not None :
244+ username = self .config .grpc_config .auth_config .username
245+ password = self .config .grpc_config .auth_config .password
246+
247+ # send grpc request
248+ channel = self ._get_grpc_channel ()
249+ grpc_client .write (
250+ channel = channel ,
251+ database = database ,
252+ batch_points = batch_points ,
253+ rp = rp ,
254+ username = username ,
255+ password = password ,
256+ timeout = self .config .timeout .seconds ,
257+ )
258+
194259 def create_database (self , database : str , rp : RpConfig = None ):
195260 if not database :
196261 raise ValueError ("empty database name" )
0 commit comments