github.com/apache/beam/sdks/v2@v2.48.2/go/test/integration/io/xlang/jdbc/jdbc_test.go (about)

     1  // Licensed to the Apache Software Foundation (ASF) under one or more
     2  // contributor license agreements.  See the NOTICE file distributed with
     3  // this work for additional information regarding copyright ownership.
     4  // The ASF licenses this file to You under the Apache License, Version 2.0
     5  // (the "License"); you may not use this file except in compliance with
     6  // the License.  You may obtain a copy of the License at
     7  //
     8  //	http://www.apache.org/licenses/LICENSE-2.0
     9  //
    10  // Unless required by applicable law or agreed to in writing, software
    11  // distributed under the License is distributed on an "AS IS" BASIS,
    12  // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    13  // See the License for the specific language governing permissions and
    14  // limitations under the License.
    15  package jdbc
    16  
    17  import (
    18  	"context"
    19  	"database/sql"
    20  	"flag"
    21  	"fmt"
    22  	"log"
    23  	"testing"
    24  	"time"
    25  
    26  	"github.com/apache/beam/sdks/v2/go/pkg/beam"
    27  	_ "github.com/apache/beam/sdks/v2/go/pkg/beam/runners/dataflow"
    28  	_ "github.com/apache/beam/sdks/v2/go/pkg/beam/runners/flink"
    29  	_ "github.com/apache/beam/sdks/v2/go/pkg/beam/runners/samza"
    30  	_ "github.com/apache/beam/sdks/v2/go/pkg/beam/runners/spark"
    31  	"github.com/apache/beam/sdks/v2/go/pkg/beam/testing/ptest"
    32  	"github.com/apache/beam/sdks/v2/go/test/integration"
    33  	"github.com/apache/beam/sdks/v2/go/test/integration/internal/containers"
    34  	"github.com/docker/go-connections/nat"
    35  	_ "github.com/go-sql-driver/mysql"
    36  	_ "github.com/lib/pq"
    37  	"github.com/testcontainers/testcontainers-go/wait"
    38  )
    39  
    40  const (
    41  	postgresImage = "postgres"
    42  	postgresPort  = "5432/tcp"
    43  	maxRetries    = 5
    44  )
    45  
    46  var expansionAddr string // Populate with expansion address labelled "schemaio".
    47  
    48  func checkFlags(t *testing.T) {
    49  	if expansionAddr == "" {
    50  		t.Skip("No Schema IO expansion address provided.")
    51  	}
    52  }
    53  
    54  func setupTestContainer(ctx context.Context, t *testing.T, dbname, username, password string) string {
    55  	t.Helper()
    56  
    57  	env := map[string]string{
    58  		"POSTGRES_PASSWORD": password,
    59  		"POSTGRES_USER":     username,
    60  		"POSTGRES_DB":       dbname,
    61  	}
    62  	hostname := "localhost"
    63  
    64  	dbURL := func(host string, port nat.Port) string {
    65  		return fmt.Sprintf("postgres://%s:%s@%s:%s/%s?sslmode=disable", username, password, host, port.Port(), dbname)
    66  	}
    67  	waitStrategy := wait.ForSQL(postgresPort, "postgres", dbURL).WithStartupTimeout(time.Second * 5)
    68  
    69  	container := containers.NewContainer(
    70  		ctx,
    71  		t,
    72  		postgresImage,
    73  		maxRetries,
    74  		containers.WithPorts([]string{postgresPort}),
    75  		containers.WithEnv(env),
    76  		containers.WithHostname(hostname),
    77  		containers.WithWaitStrategy(waitStrategy),
    78  	)
    79  
    80  	mappedPort := containers.Port(ctx, t, container, postgresPort)
    81  
    82  	url := fmt.Sprintf("postgres://%s:%s@localhost:%s/%s?sslmode=disable", username, password, mappedPort, dbname)
    83  	db, err := sql.Open("postgres", url)
    84  	if err != nil {
    85  		t.Fatalf("failed to establish database connection: %s", err)
    86  	}
    87  	defer db.Close()
    88  
    89  	_, err = db.ExecContext(ctx, "CREATE TABLE roles(role_id bigint PRIMARY KEY);")
    90  	if err != nil {
    91  		t.Fatalf("can't create table, check command and access level")
    92  	}
    93  
    94  	return mappedPort
    95  }
    96  
    97  // TestJDBCIO_BasicReadWrite tests basic read and write transform from JDBC.
    98  func TestJDBCIO_BasicReadWrite(t *testing.T) {
    99  	integration.CheckFilters(t)
   100  	checkFlags(t)
   101  
   102  	ctx := context.Background()
   103  	dbname := "postjdbc"
   104  	username := "newuser"
   105  	password := "password"
   106  
   107  	port := setupTestContainer(ctx, t, dbname, username, password)
   108  	tableName := "roles"
   109  	host := "localhost"
   110  	jdbcUrl := fmt.Sprintf("jdbc:postgresql://%s:%s/%s", host, port, dbname)
   111  
   112  	write := WritePipeline(expansionAddr, tableName, "org.postgresql.Driver", jdbcUrl, username, password)
   113  	ptest.RunAndValidate(t, write)
   114  
   115  	read := ReadPipeline(expansionAddr, tableName, "org.postgresql.Driver", jdbcUrl, username, password)
   116  	ptest.RunAndValidate(t, read)
   117  }
   118  
   119  // TestJDBCIO_PostgresReadWrite tests basic read and write transform from JDBC with postgres.
   120  func TestJDBCIO_PostgresReadWrite(t *testing.T) {
   121  	integration.CheckFilters(t)
   122  	checkFlags(t)
   123  
   124  	dbname := "postjdbc"
   125  	username := "newuser"
   126  	password := "password"
   127  	ctx := context.Background()
   128  	port := setupTestContainer(ctx, t, dbname, username, password)
   129  	tableName := "roles"
   130  	host := "localhost"
   131  	jdbcUrl := fmt.Sprintf("jdbc:postgresql://%s:%s/%s", host, port, dbname)
   132  
   133  	write := WriteToPostgres(expansionAddr, tableName, jdbcUrl, username, password)
   134  	ptest.RunAndValidate(t, write)
   135  
   136  	read := ReadFromPostgres(expansionAddr, tableName, jdbcUrl, username, password)
   137  	ptest.RunAndValidate(t, read)
   138  }
   139  
   140  func TestMain(m *testing.M) {
   141  	flag.Parse()
   142  	beam.Init()
   143  
   144  	services := integration.NewExpansionServices()
   145  	defer func() { services.Shutdown() }()
   146  	addr, err := services.GetAddr("schemaio")
   147  	if err != nil {
   148  		log.Printf("skipping missing expansion service: %v", err)
   149  	} else {
   150  		expansionAddr = addr
   151  	}
   152  
   153  	ptest.MainRet(m)
   154  }