YQ Connector: YQ-2708 support read using iam token

ydb-platform · Jan 29, 2024 · 20d234c · 20d234c
1 parent 7b25b26
commit 20d234c
Show file tree

Hide file tree

Showing 11 changed files with 242 additions and 112 deletions.
diff --git a/api/common/data_source.pb.go b/api/common/data_source.pb.go
diff --git a/api/common/endpoint.pb.go b/api/common/endpoint.pb.go
diff --git a/api/service/connector.pb.go b/api/service/connector.pb.go
diff --git a/api/service/connector_grpc.pb.go b/api/service/connector_grpc.pb.go
diff --git a/api/service/protos/connector.pb.go b/api/service/protos/connector.pb.go
diff --git a/app/config/client.pb.go b/app/config/client.pb.go
diff --git a/app/config/server.pb.go b/app/config/server.pb.go
diff --git a/app/server/datasource/rdbms/ydb/connection_manager.go b/app/server/datasource/rdbms/ydb/connection_manager.go
@@ -4,6 +4,7 @@ import (
 	"context"
 	"database/sql"
 	"fmt"
+	"strings"
 	"time"
 
 	"github.com/ydb-platform/ydb-go-genproto/protos/Ydb"
@@ -87,7 +88,21 @@ func (c *connectionManager) Make(
 	endpoint := common.EndpointToString(dsi.Endpoint)
 	dsn := sugar.DSN(endpoint, dsi.Database, dsi.UseTls)
 
-	ydbDriver, err := ydb_sdk.Open(ctx, dsn)
+	var (
+		sb   strings.Builder
+		cred ydb_sdk.Option
+	)
+
+	if dsi.Credentials.GetToken() != nil {
+		sb.WriteString(fmt.Sprintf("%s %s", dsi.Credentials.GetToken().Type, dsi.Credentials.GetToken().Value))
+		cred = ydb_sdk.WithAccessTokenCredentials(sb.String())
+	} else {
+		cred = ydb_sdk.WithAnonymousCredentials()
+	}
+
+	ydbDriver, err := ydb_sdk.Open(ctx, dsn,
+		cred,
+	)
 	if err != nil {
 		return nil, fmt.Errorf("open driver error: %w", err)
 	}

diff --git a/app/server/validate.go b/app/server/validate.go
@@ -101,10 +101,6 @@ func validateDataSourceInstance(logger *zap.Logger, dsi *api_common.TDataSourceI
 		logger.Warn("connector will use insecure connection to access data source")
 	}
 
-	if dsi.Protocol == api_common.EProtocol_PROTOCOL_UNSPECIFIED {
-		return fmt.Errorf("protocol field is empty: %w", common.ErrInvalidRequest)
-	}
-
 	switch dsi.GetKind() {
 	case api_common.EDataSourceKind_POSTGRESQL:
 		if dsi.GetPgOptions().GetSchema() == "" {

diff --git a/debug/config/client/ydb_iam.local.txt b/debug/config/client/ydb_iam.local.txt
@@ -0,0 +1,21 @@
+endpoint {
+    host: "localhost"
+    port: 50051
+}
+
+data_source_instance {
+    kind: YDB
+    endpoint {
+        host: "ydb-ru-prestable.yandex.net"
+        port: 2135
+    }
+    database: "/ru-prestable/ydb_home/woodey/ydb144"
+    credentials {
+        token {
+            type: "IAM"
+            value: "..."
+        }
+    }
+    use_tls: false
+}
+
diff --git a/debug/config/client/ydb.local.txt → debug/config/client/ydb_no_cred.local.txt b/debug/config/client/ydb.local.txt → debug/config/client/ydb_no_cred.local.txt
@@ -10,6 +10,6 @@ data_source_instance {
         port: 2136
     }
     database: "local"
-    protocol: NATIVE
+    use_tls: false
 }