@@ -43,10 +43,10 @@ import { DatabaseDialectCodes } from "./database_dialect/database_dialect_codes"
43
43
import { getWriter , logTopology } from "./utils/utils" ;
44
44
import { TelemetryFactory } from "./utils/telemetry/telemetry_factory" ;
45
45
import { DriverDialect } from "./driver_dialect/driver_dialect" ;
46
- import { ConfigurationProfile } from "./profile/configuration_profile" ;
47
- import { SessionState } from "./session_state" ;
46
+ import { AllowedAndBlockedHosts } from "./AllowedAndBlockedHosts" ;
48
47
49
48
export class PluginService implements ErrorHandler , HostListProviderService {
49
+ private static readonly DEFAULT_HOST_AVAILABILITY_CACHE_EXPIRE_NANO = 5 * 60_000_000_000 ; // 5 minutes
50
50
private readonly _currentClient : AwsClient ;
51
51
private _currentHostInfo ?: HostInfo ;
52
52
private _hostListProvider ?: HostListProvider ;
@@ -61,6 +61,7 @@ export class PluginService implements ErrorHandler, HostListProviderService {
61
61
protected readonly sessionStateService : SessionStateService ;
62
62
protected static readonly hostAvailabilityExpiringCache : CacheMap < string , HostAvailability > = new CacheMap < string , HostAvailability > ( ) ;
63
63
readonly props : Map < string , any > ;
64
+ private allowedAndBlockedHosts : AllowedAndBlockedHosts | null = null ;
64
65
65
66
constructor (
66
67
container : PluginServiceManagerContainer ,
@@ -116,17 +117,29 @@ export class PluginService implements ErrorHandler, HostListProviderService {
116
117
this . _currentHostInfo = this . _initialConnectionHostInfo ;
117
118
118
119
if ( ! this . _currentHostInfo ) {
119
- if ( this . getHosts ( ) . length === 0 ) {
120
+ if ( this . getAllHosts ( ) . length === 0 ) {
120
121
throw new AwsWrapperError ( Messages . get ( "PluginService.hostListEmpty" ) ) ;
121
122
}
122
123
123
- const writerHost = getWriter ( this . getHosts ( ) ) ;
124
+ const writerHost = getWriter ( this . getAllHosts ( ) ) ;
125
+ if ( ! this . getHosts ( ) . some ( ( hostInfo : HostInfo ) => hostInfo . host === writerHost ?. host ) ) {
126
+ throw new AwsWrapperError (
127
+ Messages . get (
128
+ "PluginService.currentHostNotAllowed" ,
129
+ this . _currentHostInfo ? this . _currentHostInfo . host : "<null>" ,
130
+ logTopology ( this . hosts , "[PluginService.currentHostNotAllowed]" )
131
+ )
132
+ ) ;
133
+ }
134
+
124
135
if ( writerHost ) {
125
136
this . _currentHostInfo = writerHost ;
126
137
} else {
127
138
this . _currentHostInfo = this . getHosts ( ) [ 0 ] ;
128
139
}
129
140
}
141
+
142
+ logger . debug ( `Set current host to: ${ this . _currentHostInfo . host } ` ) ;
130
143
}
131
144
132
145
return this . _currentHostInfo ;
@@ -260,11 +273,64 @@ export class PluginService implements ErrorHandler, HostListProviderService {
260
273
}
261
274
}
262
275
263
- getHosts ( ) : HostInfo [ ] {
276
+ getAllHosts ( ) : HostInfo [ ] {
264
277
return this . hosts ;
265
278
}
266
279
267
- setAvailability ( hostAliases : Set < string > , availability : HostAvailability ) { }
280
+ getHosts ( ) : HostInfo [ ] {
281
+ const hostPermissions = this . allowedAndBlockedHosts ;
282
+ if ( ! hostPermissions ) {
283
+ return this . hosts ;
284
+ }
285
+
286
+ let hosts = this . hosts ;
287
+ const allowedHostIds = hostPermissions . getAllowedHostIds ( ) ;
288
+ const blockedHostIds = hostPermissions . getBlockedHostIds ( ) ;
289
+
290
+ if ( allowedHostIds && allowedHostIds . size > 0 ) {
291
+ hosts = hosts . filter ( ( host : HostInfo ) => allowedHostIds . has ( host . hostId ) ) ;
292
+ }
293
+
294
+ if ( blockedHostIds && blockedHostIds . size > 0 ) {
295
+ hosts = hosts . filter ( ( host : HostInfo ) => ! blockedHostIds . has ( host . hostId ) ) ;
296
+ }
297
+
298
+ return hosts ;
299
+ }
300
+
301
+ setAvailability ( hostAliases : Set < string > , availability : HostAvailability ) {
302
+ if ( hostAliases . size === 0 ) {
303
+ return ;
304
+ }
305
+
306
+ const hostsToChange = [
307
+ ...new Set (
308
+ this . getAllHosts ( ) . filter (
309
+ ( host : HostInfo ) => hostAliases . has ( host . asAlias ) || [ ...host . aliases ] . some ( ( hostAlias : string ) => hostAliases . has ( hostAlias ) )
310
+ )
311
+ )
312
+ ] ;
313
+
314
+ if ( hostsToChange . length === 0 ) {
315
+ logger . debug ( Messages . get ( "PluginService.hostsChangeListEmpty" ) ) ;
316
+ return ;
317
+ }
318
+
319
+ const changes = new Map < string , Set < HostChangeOptions > > ( ) ;
320
+ for ( const host of hostsToChange ) {
321
+ const currentAvailability = host . getAvailability ( ) ;
322
+ PluginService . hostAvailabilityExpiringCache . put ( host . url , availability , PluginService . DEFAULT_HOST_AVAILABILITY_CACHE_EXPIRE_NANO ) ;
323
+ if ( currentAvailability !== availability ) {
324
+ let hostChanges = new Set < HostChangeOptions > ( ) ;
325
+ if ( availability === HostAvailability . AVAILABLE ) {
326
+ hostChanges = new Set ( [ HostChangeOptions . WENT_UP , HostChangeOptions . HOST_CHANGED ] ) ;
327
+ } else {
328
+ hostChanges = new Set ( [ HostChangeOptions . WENT_DOWN , HostChangeOptions . HOST_CHANGED ] ) ;
329
+ }
330
+ changes . set ( host . url , hostChanges ) ;
331
+ }
332
+ }
333
+ }
268
334
269
335
updateConfigWithProperties ( props : Map < string , any > ) {
270
336
this . _currentClient . config = Object . fromEntries ( props . entries ( ) ) ;
@@ -501,4 +567,8 @@ export class PluginService implements ErrorHandler, HostListProviderService {
501
567
attachNoOpErrorListener ( clientWrapper : ClientWrapper | undefined ) : void {
502
568
this . getDialect ( ) . getErrorHandler ( ) . attachNoOpErrorListener ( clientWrapper ) ;
503
569
}
570
+
571
+ setAllowedAndBlockedHosts ( allowedAndBlockedHosts : AllowedAndBlockedHosts ) {
572
+ this . allowedAndBlockedHosts = allowedAndBlockedHosts ;
573
+ }
504
574
}
0 commit comments