github.com/apache/beam/sdks/v2@v2.48.2/go/test/integration/io/xlang/jdbc/jdbc.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  
    16  // Package jdbc contains integration tests for cross-language JDBC IO transforms.
    17  package jdbc
    18  
    19  import (
    20  	"reflect"
    21  
    22  	"github.com/apache/beam/sdks/v2/go/pkg/beam"
    23  	"github.com/apache/beam/sdks/v2/go/pkg/beam/io/xlang/jdbcio"
    24  	"github.com/apache/beam/sdks/v2/go/pkg/beam/testing/passert"
    25  )
    26  
    27  func init() {
    28  	beam.RegisterType(reflect.TypeOf((*JdbcTestRow)(nil)).Elem())
    29  }
    30  
    31  // JdbcTestRow is the sample schema for read and write transform test.
    32  type JdbcTestRow struct {
    33  	Role_id int32 `beam:"role_id"`
    34  }
    35  
    36  func writeRows(s beam.Scope, expansionAddr, tableName, driverClassName, jdbcUrl, username, password string, input beam.PCollection) {
    37  	s = s.Scope("jdbc_test.WriteToJdbc")
    38  	jdbcio.Write(s, tableName, driverClassName, jdbcUrl, username, password, input, jdbcio.ExpansionAddrWrite(expansionAddr),
    39  		jdbcio.WriteClasspaths([]string{"org.postgresql:postgresql:42.3.3", "mysql:mysql-connector-java:8.0.28"}))
    40  }
    41  
    42  // WritePipeline creates a pipeline for JDBC IO Write transform.
    43  func WritePipeline(expansionAddr, tableName, driverClassName, jdbcUrl, username, password string) *beam.Pipeline {
    44  	beam.Init()
    45  	p, s := beam.NewPipelineWithRoot()
    46  	rows := []JdbcTestRow{{1}, {2}}
    47  	input := beam.CreateList(s, rows)
    48  	writeRows(s, expansionAddr, tableName, driverClassName, jdbcUrl, username, password, input)
    49  	return p
    50  }
    51  
    52  func readRows(s beam.Scope, expansionAddr, tableName, driverClassName, jdbcUrl, username, password string) beam.PCollection {
    53  	s = s.Scope("jdbc_test.ReadFromJdbc")
    54  	outT := reflect.TypeOf((*JdbcTestRow)(nil)).Elem()
    55  	res := jdbcio.Read(s, tableName, driverClassName, jdbcUrl, username, password, outT, jdbcio.ExpansionAddrRead(expansionAddr),
    56  		jdbcio.ReadClasspaths([]string{"org.postgresql:postgresql:42.3.3", "mysql:mysql-connector-java:8.0.28"}))
    57  	return res
    58  }
    59  
    60  // ReadPipeline creates a pipeline for JDBC IO Read transform.
    61  func ReadPipeline(expansionAddr, tableName, driverClassName, jdbcUrl, username, password string) *beam.Pipeline {
    62  	beam.Init()
    63  	p, s := beam.NewPipelineWithRoot()
    64  	res := readRows(s, expansionAddr, tableName, driverClassName, jdbcUrl, username, password)
    65  	want := beam.CreateList(s, []JdbcTestRow{{1}, {2}})
    66  	passert.Equals(s, res, want)
    67  	return p
    68  }
    69  
    70  // WriteToPostgres creates a pipeline for JDBC IO Write transform.
    71  func WriteToPostgres(expansionAddr, tableName, jdbcUrl, username, password string) *beam.Pipeline {
    72  	beam.Init()
    73  	p, s := beam.NewPipelineWithRoot()
    74  	rows := []JdbcTestRow{{1}, {2}}
    75  	input := beam.CreateList(s, rows)
    76  	jdbcio.WriteToPostgres(s.Scope("jdbc_test.ReadFromJdbc"), tableName, jdbcUrl, username, password, input, jdbcio.ExpansionAddrWrite(expansionAddr))
    77  	return p
    78  }
    79  
    80  // ReadFromPostgres creates a pipeline for JDBC IO Read transform.
    81  func ReadFromPostgres(expansionAddr, tableName, jdbcUrl, username, password string) *beam.Pipeline {
    82  	beam.Init()
    83  	p, s := beam.NewPipelineWithRoot()
    84  	outT := reflect.TypeOf((*JdbcTestRow)(nil)).Elem()
    85  	res := jdbcio.ReadFromPostgres(s.Scope("jdbc_test.WriteToJdbc"), tableName, jdbcUrl, username, password, outT, jdbcio.ExpansionAddrRead(expansionAddr))
    86  	want := beam.CreateList(s, []JdbcTestRow{{1}, {2}})
    87  	passert.Equals(s, res, want)
    88  	return p
    89  }