github.com/cockroachdb/cockroach@v20.2.0-alpha.1+incompatible/pkg/ccl/importccl/read_import_pgdump_test.go (about)

     1  // Copyright 2018 The Cockroach Authors.
     2  //
     3  // Licensed as a CockroachDB Enterprise file under the Cockroach Community
     4  // License (the "License"); you may not use this file except in compliance with
     5  // the License. You may obtain a copy of the License at
     6  //
     7  //     https://github.com/cockroachdb/cockroach/blob/master/licenses/CCL.txt
     8  
     9  package importccl
    10  
    11  import (
    12  	"fmt"
    13  	"io"
    14  	"strings"
    15  	"testing"
    16  
    17  	"github.com/cockroachdb/cockroach/pkg/util/leaktest"
    18  )
    19  
    20  func TestPostgreStream(t *testing.T) {
    21  	defer leaktest.AfterTest(t)()
    22  
    23  	const sql = `
    24  select 1;
    25  -- select 2;
    26  select 3;
    27  select 4;
    28  select 5;
    29  select '12345678901234567890123456789012345678901234567890123456789012345678901234567890';
    30  --
    31  `
    32  
    33  	p := newPostgreStream(strings.NewReader(sql), defaultScanBuffer)
    34  	var sb strings.Builder
    35  	for {
    36  		s, err := p.Next()
    37  		if err == io.EOF {
    38  			break
    39  		}
    40  		if err != nil {
    41  			t.Fatal(err)
    42  		}
    43  		fmt.Fprintf(&sb, "%s;\n", s)
    44  	}
    45  	const expect = `SELECT 1;
    46  SELECT 3;
    47  SELECT 4;
    48  SELECT 5;
    49  SELECT '12345678901234567890123456789012345678901234567890123456789012345678901234567890';
    50  `
    51  	got := sb.String()
    52  	if expect != got {
    53  		t.Fatalf("got %q, expected %q", got, expect)
    54  	}
    55  }
    56  
    57  func TestPostgreStreamCopy(t *testing.T) {
    58  	defer leaktest.AfterTest(t)()
    59  
    60  	const sql = `
    61  CREATE TABLE public.second (
    62      i int8 NOT NULL,
    63      s text
    64  );
    65  
    66  
    67  --
    68  -- Data for Name: second; Type: TABLE DATA; Schema: public; Owner: -
    69  --
    70  
    71  COPY public.second (i, s) FROM stdin;
    72  0	0
    73  1	1
    74  2	2
    75  3	3
    76  4	4
    77  5	5
    78  6	6
    79  \.
    80  
    81  
    82  --
    83  -- Name: second second_pkey; Type: CONSTRAINT; Schema: public; Owner: -
    84  --
    85  
    86  ALTER TABLE ONLY public.second
    87      ADD CONSTRAINT second_pkey PRIMARY KEY (i);
    88  
    89  --
    90  -- Name: t; Type: TABLE; Schema: public; Owner: -
    91  --
    92  
    93  CREATE TABLE public.t (
    94      s text
    95  );
    96  
    97  
    98  --
    99  -- Data for Name: t; Type: TABLE DATA; Schema: public; Owner: -
   100  --
   101  
   102  COPY public.t (s) FROM stdin;
   103  
   104  \\.
   105  ;
   106  \.
   107  
   108  
   109  --
   110  -- PostgreSQL database dump complete
   111  --
   112  `
   113  
   114  	p := newPostgreStream(strings.NewReader(sql), defaultScanBuffer)
   115  	var sb strings.Builder
   116  	for {
   117  		s, err := p.Next()
   118  		if err == io.EOF {
   119  			break
   120  		}
   121  		if err != nil {
   122  			t.Fatal(err)
   123  		}
   124  		fmt.Fprintf(&sb, "%s;\n", s)
   125  	}
   126  	const expect = `CREATE TABLE public.second (i INT8 NOT NULL, s STRING);
   127  COPY public.second (i, s) FROM STDIN;
   128  "0"	"0";
   129  "1"	"1";
   130  "2"	"2";
   131  "3"	"3";
   132  "4"	"4";
   133  "5"	"5";
   134  "6"	"6";
   135  COPY done;
   136  ALTER TABLE public.second ADD CONSTRAINT second_pkey PRIMARY KEY (i);
   137  CREATE TABLE public.t (s STRING);
   138  COPY public.t (s) FROM STDIN;
   139  "";
   140  "\\.";
   141  ";";
   142  COPY done;
   143  `
   144  	got := sb.String()
   145  	if expect != got {
   146  		t.Fatalf("got %s, expected %s", got, expect)
   147  	}
   148  }