github.com/thiagoyeds/go-cloud@v0.26.0/postgres/postgres.go (about)

     1  // Copyright 2019 The Go Cloud Development Kit Authors
     2  //
     3  // Licensed under the Apache License, Version 2.0 (the "License");
     4  // you may not use this file except in compliance with the License.
     5  // You may obtain a copy of the License at
     6  //
     7  //     https://www.apache.org/licenses/LICENSE-2.0
     8  //
     9  // Unless required by applicable law or agreed to in writing, software
    10  // distributed under the License is distributed on an "AS IS" BASIS,
    11  // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    12  // See the License for the specific language governing permissions and
    13  // limitations under the License.
    14  
    15  // Package postgres provides functions to open PostgreSQL databases with OpenCensus instrumentation.
    16  package postgres
    17  
    18  import (
    19  	"context"
    20  	"database/sql"
    21  	"database/sql/driver"
    22  	"net/url"
    23  
    24  	"github.com/lib/pq"
    25  	"gocloud.dev/internal/openurl"
    26  
    27  	"contrib.go.opencensus.io/integrations/ocsql"
    28  )
    29  
    30  // Scheme is the URL scheme this package registers its URLOpener under on
    31  // DefaultMux.
    32  const Scheme = "postgres"
    33  
    34  func init() {
    35  	DefaultURLMux().RegisterPostgres(Scheme, &URLOpener{})
    36  }
    37  
    38  // URLOpener opens URLs like "postgres://" by using the underlying PostgreSQL driver.
    39  // See https://godoc.org/github.com/lib/pq#hdr-Connection_String_Parameters for details.
    40  type URLOpener struct{}
    41  
    42  // OpenPostgresURL opens a new database connection wrapped with OpenCensus instrumentation.
    43  func (*URLOpener) OpenPostgresURL(ctx context.Context, u *url.URL) (*sql.DB, error) {
    44  	u2 := new(url.URL)
    45  	*u2 = *u
    46  	u2.Scheme = "postgres"
    47  	return sql.OpenDB(connector{dsn: u2.String()}), nil
    48  }
    49  
    50  type connector struct {
    51  	dsn string
    52  }
    53  
    54  func (c connector) Connect(ctx context.Context) (driver.Conn, error) {
    55  	return c.Driver().Open(c.dsn)
    56  }
    57  
    58  func (c connector) Driver() driver.Driver {
    59  	return ocsql.Wrap(&pq.Driver{})
    60  }
    61  
    62  // A type that implements PostgresURLOpener can open connection based on a URL.
    63  // The opener must not modify the URL argument. OpenPostgresURL must be safe to
    64  // call from multiple goroutines.
    65  //
    66  // This interface is generally implemented by types in driver packages.
    67  type PostgresURLOpener interface {
    68  	OpenPostgresURL(ctx context.Context, u *url.URL) (*sql.DB, error)
    69  }
    70  
    71  // URLMux is a URL opener multiplexer. It matches the scheme of the URLs
    72  // against a set of registered schemes and calls the opener that matches the
    73  // URL's scheme.
    74  //
    75  // The zero value is a multiplexer with no registered schemes.
    76  type URLMux struct {
    77  	schemes openurl.SchemeMap
    78  }
    79  
    80  // RegisterPostgres registers the opener with the given scheme. If an opener
    81  // already exists for the scheme, RegisterPostgres panics.
    82  func (mux *URLMux) RegisterPostgres(scheme string, opener PostgresURLOpener) {
    83  	mux.schemes.Register("postgres", "DB", scheme, opener)
    84  }
    85  
    86  // OpenPostgres calls OpenPostgresURL with the URL parsed from urlstr.
    87  // OpenPostgres is safe to call from multiple goroutines.
    88  func (mux *URLMux) OpenPostgres(ctx context.Context, urlstr string) (*sql.DB, error) {
    89  	opener, u, err := mux.schemes.FromString("DB", urlstr)
    90  	if err != nil {
    91  		return nil, err
    92  	}
    93  	return opener.(PostgresURLOpener).OpenPostgresURL(ctx, u)
    94  }
    95  
    96  // OpenPostgresURL dispatches the URL to the opener that is registered with the
    97  // URL's scheme. OpenPostgresURL is safe to call from multiple goroutines.
    98  func (mux *URLMux) OpenPostgresURL(ctx context.Context, u *url.URL) (*sql.DB, error) {
    99  	opener, err := mux.schemes.FromURL("DB", u)
   100  	if err != nil {
   101  		return nil, err
   102  	}
   103  	return opener.(PostgresURLOpener).OpenPostgresURL(ctx, u)
   104  }
   105  
   106  var defaultURLMux = new(URLMux)
   107  
   108  // DefaultURLMux returns the URLMux used by OpenPostgres.
   109  //
   110  // Driver packages can use this to register their PostgresURLOpener on the mux.
   111  func DefaultURLMux() *URLMux {
   112  	return defaultURLMux
   113  }
   114  
   115  // Open opens the bucket identified by the URL given. URL openers must be
   116  // registered in the DefaultURLMux, which is typically done in driver
   117  // packages' initialization.
   118  //
   119  // See the URLOpener documentation in driver subpackages for more
   120  // details on supported scheme(s) and URL parameter(s).
   121  func Open(ctx context.Context, urlstr string) (*sql.DB, error) {
   122  	return defaultURLMux.OpenPostgres(ctx, urlstr)
   123  }