├── .dockerignore ├── .gitignore ├── .travis.yml ├── LICENSE ├── Makefile ├── README.md ├── SECURITY.md ├── bcp.Dockerfile ├── bcpy ├── __init__.py ├── binary_callers.py ├── data_objects.py ├── format_file_builder.py └── tmp_file.py ├── conftest.py ├── requirements.txt ├── sample.env ├── setup.py └── tests ├── Dockerfile ├── __init__.py ├── create_test_db.sql ├── data1.csv ├── docker-compose.yml ├── test.sh ├── test_basic.py ├── test_flat_file.py ├── test_pandas.py └── wait-for-it.sh /.dockerignore: -------------------------------------------------------------------------------- 1 | __pycache__ 2 | /tests/__pycache__ 3 | .dockerignore 4 | Dockerfile 5 | *.pyc 6 | *.pyo 7 | *.pyd 8 | .git 9 | -------------------------------------------------------------------------------- /.gitignore: -------------------------------------------------------------------------------- 1 | # VS Codium 2 | settings.json 3 | 4 | # Pycharm 5 | .idea/ 6 | 7 | # Tests 8 | .env 9 | 10 | # Byte-compiled / optimized / DLL files 11 | __pycache__/ 12 | *.py[cod] 13 | *$py.class 14 | 15 | # C extensions 16 | *.so 17 | 18 | # Distribution / packaging 19 | .Python 20 | build/ 21 | develop-eggs/ 22 | dist/ 23 | downloads/ 24 | eggs/ 25 | .eggs/ 26 | lib/ 27 | lib64/ 28 | parts/ 29 | sdist/ 30 | var/ 31 | wheels/ 32 | pip-wheel-metadata/ 33 | share/python-wheels/ 34 | *.egg-info/ 35 | .installed.cfg 36 | *.egg 37 | MANIFEST 38 | 39 | # PyInstaller 40 | # Usually these files are written by a python script from a template 41 | # before PyInstaller builds the exe, so as to inject date/other infos into it. 42 | *.manifest 43 | *.spec 44 | 45 | # Installer logs 46 | pip-log.txt 47 | pip-delete-this-directory.txt 48 | 49 | # Unit test / coverage reports 50 | htmlcov/ 51 | .tox/ 52 | .nox/ 53 | .coverage 54 | .coverage.* 55 | .cache 56 | nosetests.xml 57 | coverage.xml 58 | *.cover 59 | .hypothesis/ 60 | .pytest_cache/ 61 | 62 | # Translations 63 | *.mo 64 | *.pot 65 | 66 | # Django stuff: 67 | *.log 68 | local_settings.py 69 | db.sqlite3 70 | 71 | # Flask stuff: 72 | instance/ 73 | .webassets-cache 74 | 75 | # Scrapy stuff: 76 | .scrapy 77 | 78 | # Sphinx documentation 79 | docs/_build/ 80 | 81 | # PyBuilder 82 | target/ 83 | 84 | # Jupyter Notebook 85 | .ipynb_checkpoints 86 | 87 | # IPython 88 | profile_default/ 89 | ipython_config.py 90 | 91 | # pyenv 92 | .python-version 93 | 94 | # celery beat schedule file 95 | celerybeat-schedule 96 | 97 | # SageMath parsed files 98 | *.sage.py 99 | 100 | # Environments 101 | .env 102 | .venv 103 | env/ 104 | venv/ 105 | ENV/ 106 | env.bak/ 107 | venv.bak/ 108 | 109 | # Spyder project settings 110 | .spyderproject 111 | .spyproject 112 | 113 | # Rope project settings 114 | .ropeproject 115 | 116 | # mkdocs documentation 117 | /site 118 | 119 | # mypy 120 | .mypy_cache/ 121 | .dmypy.json 122 | dmypy.json 123 | 124 | # Pyre type checker 125 | .pyre/ -------------------------------------------------------------------------------- /.travis.yml: -------------------------------------------------------------------------------- 1 | services: 2 | - docker 3 | sudo: false 4 | language: python 5 | python: 3.6 6 | before_install: 7 | - if [ $TRAVIS_BRANCH == "staging" ]; then sed -i '/version=*/c\ version=\"'"$(date +%Y.%m.%d.%H.%M.%S)"'\",' setup.py; fi 8 | before_script: 9 | - cp sample.env .env 10 | script: make test 11 | before_deploy: 12 | - rm -rf .env 13 | - git update-index --assume-unchanged setup.py 14 | deploy: 15 | - provider: pypi 16 | user: titan_550 17 | password: 18 | secure: ciw+Y0xRfAGXlE6NdRN+16ZLlh1Cu1l34/Jzhj8vkc24bwc83FgKxWx+O/kVcd2RiHOJVUP4NCHi0EqM8egm6QoXb1Wjxt2xjcx6ipoBgI+mFwa+JNBc2FU1l159X+1a/n9lGhly1eoG7hwfIljrpNfqg7ZIP7lRyGdjtJTheAO7JIqbvBo96h8RzJO5VDHYaqky4mjHQwDE+90ecWb9gITcMtLc9PlXiyNS7aKP00mgSjYIpe3S2wjBoSD2cboKVrKaVdheBxcmh2GSdTvcV2Ykg8HQRXUB3wZZYc9XcmMeRQyELaA3AaCu3aSlTJeC0uKGQI+msAKCttQKyu7tDbCpzNWhu9IeGo2dQ0Tb7gkZY4GudsbRLmMxpz2LLKh9yn9c4cDu5VqfI9MyXtfKvuurunwbJ+vWEzHanFminCjb0Lv15GT9//EGXIlYFiGppmLNwtJjcozuYhGm3tMDZQVFQvBALzPlDumj+sYHxkrJoTxezq8uNfRBPcK89xijdbmjISshnhJgOkUxLb8+OwF+s6WZIStvqPjGMaJZRHCVwNV9mfYgkEe031c5TUrAjWhqfFosvcGf6R7a+iXctFCYWKl7J2ZVmrroo8UTy/45GvHWmMXmlSmaSBoc2P2I0BNTxSyXAPs1MIT6CYF7qPkcyPbYFLR04yNiyoKqiNs= 19 | on: 20 | branch: master 21 | tags: true 22 | - provider: pypi 23 | user: titan_550 24 | password: 25 | secure: Ywtu7jajCO1/XPWBg1usKNVnFuzPVnRRKg7MvP/oHRUbcnrFA5b0RGhhuE3OoXesngXEOCXY+S/pUC6Iyaw05oi6l1JUd9HmrdfLE68E+WQgBXMzjtIv4NaXcmWTHlRjgjYhJPmMXdza55CwjqXK3zuidtV1VysKoiWNzPR4jBFP8Y2wsx9ucIvfDMBOYlLkqUp15/ZomG26XWv5Cq0Q+4kD6sAAGZBYAZsZ2O6kKVziRK+/0vca1BrE0IrunaPVJcvvNQxwxavaXlcO/ZUn/VAYZWXts1/VaWw10W2viGU5QwHb0yQCCPJYREpxidIP2dBbPn+5ZyRVJEtJGxTzS2VLW+RNXEG/UUNqtGFjN4U7MP57EF9MewOkWPhTMuJ96j8zlfzjh4/2Xc24gPqK7QJSzwDxjoPVqUqyJj+Chcm2UJ6jC0YDaVUGIvDk/m0MNEPLaYvEb0HT8L5U9K5MtDJ5OUcKI2J5dmJcJkgNrFfwFpp5uAAr5uhIlAV6pzmFvca4h65of6aKskQ20vIWumlhuyCdXeUuql5MVzVy7X9Iv4nGbS3XSx2NJohDiJFHXgrW2GiCW+UwQ+HISntNYYEP2qMx6MinjsQL2lm0SZbXVlvNwt5hc8inT005WARiSE4UV7+1/Sb3tmfLk6LaS24tVpS24RtcbRmRINrWnA8= 26 | server: https://test.pypi.org/legacy/ 27 | on: 28 | branch: staging -------------------------------------------------------------------------------- /LICENSE: -------------------------------------------------------------------------------- 1 | MIT License 2 | 3 | Copyright (c) 2019 John Shojaei 4 | 5 | Permission is hereby granted, free of charge, to any person obtaining a copy 6 | of this software and associated documentation files (the "Software"), to deal 7 | in the Software without restriction, including without limitation the rights 8 | to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 9 | copies of the Software, and to permit persons to whom the Software is 10 | furnished to do so, subject to the following conditions: 11 | 12 | The above copyright notice and this permission notice shall be included in all 13 | copies or substantial portions of the Software. 14 | 15 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 16 | IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 17 | FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 18 | AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 19 | LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 20 | OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE 21 | SOFTWARE. -------------------------------------------------------------------------------- /Makefile: -------------------------------------------------------------------------------- 1 | cnf ?= .env 2 | include $(cnf) 3 | export $(shell sed 's/=.*//' $(cnf)) 4 | 5 | .PHONY: help test dev build 6 | 7 | help: ## This help. 8 | @awk 'BEGIN {FS = ":.*?## "} /^[a-zA-Z_-]+:.*?## / {printf "\033[36m%-30s\033[0m %s\n", $$1, $$2}' $(MAKEFILE_LIST) 9 | 10 | .DEFAULT_GOAL := help 11 | 12 | test: ## Test using docker-compose 13 | docker-compose -f tests/docker-compose.yml up --exit-code-from client --build 14 | 15 | dev: build ## Starts a shell in the client environment 16 | chcon -Rt svirt_sandbox_file_t $(shell pwd) 17 | docker-compose -f tests/docker-compose.yml run --rm -v "$(shell pwd):/bcpy" client bash || true 18 | docker-compose -f tests/docker-compose.yml down 19 | 20 | build: ## builds docker images in the docker-compose file 21 | docker build --force-rm -t bcpy_client -f ./tests/Dockerfile . 22 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # bcpy 2 | 3 | 4 | 5 | 6 | 11 | 12 | 13 | 14 | 19 | 20 | 21 | 22 | 27 | 28 |
Latest Release 7 | 8 | latest release 9 | 10 |
License 15 | 16 | license 17 | 18 |
Build Status (master) 23 | 24 | travis build status 25 | 26 |
29 | 30 | ## What is it? 31 | 32 | This package is a wrapper for Microsoft's SQL Server bcp utility. Current database drivers available in Python are not fast enough for transferring millions of records (yes, I have tried [pyodbc fast_execute_many](https://github.com/mkleehammer/pyodbc/wiki/Features-beyond-the-DB-API#fast_executemany)). Despite the IO hits, the fastest option by far is saving the data to a CSV file in file system (preferably /dev/shm tmpfs) and using the bcp utility to transfer the CSV file to SQL Server. 33 | 34 | ## How Can I Install It? 35 | 36 | 1. Make sure your computeer has the [requirements](#requirements). 37 | 1. You can download and install this package from PyPI repository by running the command below. 38 | 39 | ```bash 40 | pip install bcpy 41 | ``` 42 | 43 | ## Examples 44 | 45 | Following examples show you how to load (1) flat files and (2) DataFrame objects to SQL Server using this package. 46 | 47 | ### Flat File 48 | 49 | Following example assumes that you have a comma separated file with no qualifier in path 'tests/data1.csv'. The code below sends the the file to SQL Server. 50 | 51 | ```python 52 | import bcpy 53 | 54 | 55 | sql_config = { 56 | 'server': 'sql_server_hostname', 57 | 'database': 'database_name', 58 | 'username': 'test_user', 59 | 'password': 'test_user_password1234' 60 | } 61 | sql_table_name = 'test_data1' 62 | csv_file_path = 'tests/data1.csv' 63 | flat_file = bcpy.FlatFile(qualifier='', path=csv_file_path) 64 | sql_table = bcpy.SqlTable(sql_config, table=sql_table_name) 65 | flat_file.to_sql(sql_table) 66 | ``` 67 | 68 | ### DataFrame 69 | 70 | The following example creates a DataFrame with 100 rows and 4 columns populated with random data and then it sends it to SQL Server. 71 | 72 | ```python 73 | import bcpy 74 | import numpy as np 75 | import pandas as pd 76 | 77 | 78 | sql_config = { 79 | 'server': 'sql_server_hostname', 80 | 'database': 'database_name', 81 | 'username': 'test_user', 82 | 'password': 'test_user_password1234' 83 | } 84 | table_name = 'test_dataframe' 85 | df = pd.DataFrame(np.random.randint(-100, 100, size=(100, 4)), 86 | columns=list('ABCD')) 87 | bdf = bcpy.DataFrame(df) 88 | sql_table = bcpy.SqlTable(sql_config, table=table_name) 89 | bdf.to_sql(sql_table) 90 | ``` 91 | 92 | ## Requirements 93 | 94 | You need a working version of Microsoft bcp installed in your system. Your PATH environment variable should contain the directory of the bcp utility. Following are the installation tutorials for different operating systems. 95 | 96 | - [Dockerfile (Ubuntu 18.04)](./bcp.Dockerfile) 97 | - [Linux](https://docs.microsoft.com/en-us/sql/linux/sql-server-linux-setup-tools) 98 | - [Mac](https://docs.microsoft.com/en-us/sql/linux/sql-server-linux-setup-tools?view=sql-server-2017#macos) 99 | - [Windows](https://docs.microsoft.com/en-us/sql/tools/bcp-utility) 100 | -------------------------------------------------------------------------------- /SECURITY.md: -------------------------------------------------------------------------------- 1 | # Security Policy 2 | 3 | ## Supported Versions 4 | 5 | Following section shows the version that support security updates. 6 | 7 | | Version | Supported | 8 | | ------- | ------------------ | 9 | | 0.1.x | :white_check_mark: | 10 | | 0.0.x | :x: | 11 | 12 | ## Reporting a Vulnerability 13 | 14 | Please email me at titan550 \__att\__ gmail \__dot__\__com\__ if you found a security bug. 15 | 16 | 17 | I will send you an update within three days of receiving your email. 18 | -------------------------------------------------------------------------------- /bcp.Dockerfile: -------------------------------------------------------------------------------- 1 | FROM ubuntu:18.04 2 | 3 | # apt-get and system utilities 4 | RUN apt-get update && apt-get install -y \ 5 | curl apt-utils apt-transport-https debconf-utils gcc build-essential g++-5 \ 6 | && rm -rf /var/lib/apt/lists/* 7 | 8 | # Install SQL Server drivers and tools 9 | RUN curl https://packages.microsoft.com/keys/microsoft.asc | apt-key add - \ 10 | && curl https://packages.microsoft.com/config/ubuntu/18.04/prod.list > /etc/apt/sources.list.d/mssql-release.list \ 11 | && apt-get update \ 12 | && ACCEPT_EULA=Y apt-get install -y msodbcsql17 \ 13 | && ACCEPT_EULA=Y apt-get install -y mssql-tools \ 14 | && apt-get install -y unixodbc-dev libssl1.0.0 \ 15 | && rm -rf /var/lib/apt/lists/* 16 | 17 | ENV PATH="/opt/mssql-tools/bin:${PATH}" 18 | 19 | # Python 3 libraries 20 | RUN apt-get update \ 21 | && apt-get install -y python3-pip python3-dev \ 22 | && cd /usr/local/bin \ 23 | && ln -s /usr/bin/python3 python \ 24 | && pip3 install --upgrade pip \ 25 | && rm -rf /var/lib/apt/lists/* 26 | 27 | # Install necessary locales 28 | RUN apt-get update && apt-get install -y locales \ 29 | && echo "en_US.UTF-8 UTF-8" > /etc/locale.gen \ 30 | && locale-gen \ 31 | && rm -rf /var/lib/apt/lists/* 32 | -------------------------------------------------------------------------------- /bcpy/__init__.py: -------------------------------------------------------------------------------- 1 | from .data_objects import SqlTable, FlatFile, DataFrame, SqlServer 2 | 3 | 4 | name = "bcpy" 5 | __version__ = '0.0.5' 6 | -------------------------------------------------------------------------------- /bcpy/binary_callers.py: -------------------------------------------------------------------------------- 1 | import subprocess 2 | from io import StringIO 3 | import hashlib 4 | 5 | import pandas as pd 6 | 7 | 8 | def sha512(text, encoding='utf-8'): 9 | """Converts an input string to its sha512 hash 10 | """ 11 | if not isinstance(text, bytes): 12 | if isinstance(text, str): 13 | text = text.encode(encoding) 14 | else: 15 | raise ValueError('Invalid input. Cannot compute hash.') 16 | return hashlib.sha512(text).hexdigest() 17 | 18 | 19 | def bcp(sql_table, flat_file, batch_size): 20 | """Runs the bcp command to transfer the input flat file to the input 21 | SQL Server table. 22 | :param sql_table: The destination Sql Server table 23 | :type sql_table: SqlTable 24 | :param flat_file: Source flat file 25 | :type flat_file: FlatFile 26 | :param batch_size: Batch size (chunk size) to send to SQL Server 27 | :type batch_size: int 28 | """ 29 | if sql_table.with_krb_auth: 30 | auth = ['-T'] 31 | else: 32 | auth = ['-U', sql_table.username, '-P', sql_table.password] 33 | full_table_string = \ 34 | f'{sql_table.database}.{sql_table.schema}.{sql_table.table}' 35 | try: 36 | bcp_command = ['bcp', full_table_string, 'IN', flat_file.path, '-f', 37 | flat_file.get_format_file_path(), '-S', 38 | sql_table.server, '-b', str(batch_size)] + auth 39 | except Exception as e: 40 | args_clean = list() 41 | for arg in e.args: 42 | if isinstance(arg, str): 43 | arg = arg.replace(sql_table.password, 44 | sha512(sql_table.password)) 45 | args_clean.append(arg) 46 | e.args = tuple(args_clean) 47 | raise e 48 | if flat_file.file_has_header_line: 49 | bcp_command += ['-F', '2'] 50 | result = subprocess.run(bcp_command, stderr=subprocess.PIPE) 51 | if result.returncode: 52 | raise Exception( 53 | f'Bcp command failed. Details:\n{result}') 54 | 55 | 56 | def sqlcmd(server, database, command, username=None, password=None): 57 | """Runs the input command against the database and returns the output if it 58 | is a table. 59 | Leave username and password to None if you intend to use 60 | Kerberos integrated authentication 61 | :param server: SQL Server 62 | :type server: str 63 | :param database: Name of the default database for the script 64 | :type database: str 65 | :param command: SQL command to be executed against the server 66 | :type command: str 67 | :param username: Username to use for login 68 | :type username: str 69 | :param password: Password to use for login 70 | :type password: str 71 | :return: Returns a table if the command has an output. Returns None 72 | if the output does not return anything. 73 | :rtype: Pandas.DataFrame 74 | """ 75 | if not username or not password: 76 | auth = ['-E'] 77 | else: 78 | auth = ['-U', username, '-P', password] 79 | command = 'set nocount on;' + command 80 | sqlcmd_command = ['sqlcmd', '-S', server, '-d', database, '-b'] + auth + \ 81 | ['-s,', '-W', '-Q', command] 82 | result = subprocess.run(sqlcmd_command, stdout=subprocess.PIPE, 83 | stderr=subprocess.PIPE) 84 | if result.returncode: 85 | result_dump = str(result).replace(password, sha512(password)) 86 | raise Exception(f'Sqlcmd command failed. Details:\n{result_dump}') 87 | output = StringIO(result.stdout.decode('ascii')) 88 | first_line_output = output.readline().strip() 89 | if first_line_output == '': 90 | header = None 91 | else: 92 | header = 'infer' 93 | output.seek(0) 94 | try: 95 | result = pd.read_csv( 96 | filepath_or_buffer=output, 97 | skiprows=[1], 98 | header=header) 99 | except pd.errors.EmptyDataError: 100 | result = None 101 | return result 102 | -------------------------------------------------------------------------------- /bcpy/data_objects.py: -------------------------------------------------------------------------------- 1 | import csv 2 | import os 3 | 4 | from .binary_callers import bcp, sqlcmd 5 | from .format_file_builder import FormatFile 6 | from .tmp_file import TemporaryFile 7 | 8 | 9 | class DataObject: 10 | """Base object for data objects in bcpy 11 | """ 12 | 13 | def __init__(self, config): 14 | if config and not isinstance(config, dict): 15 | raise TypeError('Config parameter must be a dictionary object') 16 | 17 | def __repr__(self): 18 | output = str() 19 | for attrib, value in self.__dict__.items(): 20 | output += f'{attrib} = {repr(value)}\n' 21 | return output 22 | 23 | def __str__(self): 24 | return self.__repr__() 25 | 26 | 27 | class FlatFile(DataObject): 28 | def __init__(self, config=None, **kwargs): 29 | """ 30 | :param config: A dictionary object with the parameters. 31 | :param kwargs: Dynamic list of params which supersedes config params if 32 | they overlap. 33 | :param delimiter: flat file delimiter (default: ",") 34 | :param qualifier: flat file qualifier 35 | (default: "'" , e.g., 'col1','col2') 36 | 37 | :param newline: newline characters that separate records 38 | (default: "\n") 39 | :param path: path to the flat file 40 | :param file_has_header_line: defaults to False 41 | :param columns: a list of columns, automatically read from the file 42 | if the if the file has header a line. 43 | """ 44 | super().__init__(config) 45 | self.delimiter = ',' 46 | self.qualifier = '\'' 47 | self.newline = '\n' 48 | self.__columns = None 49 | self.path = None 50 | self.__format_file_path = None 51 | self.file_has_header_line = False 52 | if config: 53 | for key, value in config.items(): 54 | setattr(self, key, value) 55 | for key, value in kwargs.items(): 56 | setattr(self, key, value) 57 | if not self.qualifier: 58 | self.qualifier = '' 59 | 60 | def __del__(self): 61 | """Removes the temporary format file that gets created before sending 62 | the flat file to SQL Server""" 63 | try: 64 | if self.__format_file_path: 65 | os.remove(self.__format_file_path) 66 | except AttributeError: 67 | pass 68 | 69 | def _read_columns_from_file(self): 70 | """Reads columns of a flat file from its file located in object's path 71 | attribute. Is stores the results in the object's columns attribute. 72 | 73 | Note: Caller of this method assumes that the file has headers. 74 | """ 75 | with open(self.path, encoding='utf-8-sig') as f: 76 | header = f.readline() 77 | qualifier_delimiter_combo = str.format('{0}{1}{0}', self.qualifier, 78 | self.delimiter) 79 | columns_raw = header.split(qualifier_delimiter_combo) 80 | self.__columns = [columns_raw[0].lstrip(self.qualifier)] 81 | self.__columns.extend(columns_raw[1:-1]) 82 | self.__columns.append( 83 | columns_raw[-1].rstrip(self.qualifier + self.newline)) 84 | self.file_has_header_line = True 85 | 86 | def get_format_file_path(self, recalculate=False): 87 | """Returns the path to the bcp format file of the this flat file 88 | :param recalculate: uses file from cache if recalculate if False 89 | otherwise it will remove the old file and creates a 90 | new one. 91 | :return: path to the format file 92 | :rtype: str 93 | """ 94 | if not recalculate and self.__format_file_path: 95 | return self.__format_file_path 96 | else: 97 | try: 98 | if self.__format_file_path: 99 | os.remove(self.__format_file_path) 100 | except OSError: 101 | pass 102 | if not self.columns: 103 | raise Exception( 104 | 'Need the object columns or path to build the format file') 105 | self.__format_file_path = self._build_format_file() 106 | return self.__format_file_path 107 | 108 | def _build_format_file(self): 109 | """Creates the format file and writes its content to a temporary file. 110 | :return: path to the temporary file 111 | :rtype: str 112 | """ 113 | format_file_content = FormatFile.build_format_file(self) 114 | with TemporaryFile(mode='w') as f: 115 | f.write(format_file_content) 116 | format_file_path = f.name 117 | return format_file_path 118 | 119 | def _get_sql_create_statement(self, table_name=None, schema_name='dbo'): 120 | """Creates a SQL drop and re-create statement corresponding to the 121 | columns list of the object. 122 | 123 | :param table_name: name of the new table 124 | :type table_name: str 125 | :paran schema_name: name of schema 126 | :type schema_name: str 127 | :return: SQL code to create the table 128 | """ 129 | if not table_name: 130 | table_name = os.path.basename(self.path) 131 | sql_cols = ','.join( 132 | map(lambda x: f'[{x}] nvarchar(max)', self.columns)) 133 | sql_command = f"if object_id('[{schema_name}].[{table_name}]', 'U') " \ 134 | f"is not null drop table [{schema_name}].[{table_name}];" \ 135 | f'create table [{schema_name}].[{table_name}] ({sql_cols});' 136 | return sql_command 137 | 138 | def to_sql(self, sql_table, use_existing_sql_table=False, 139 | batch_size=10000): 140 | """Sends the object to SQL table 141 | :param sql_table: destination SQL table 142 | :type sql_table: SqlTable 143 | :param use_existing_sql_table: If to 144 | use an existing table in the SQL database. 145 | If not, then creates a new one. 146 | :type use_existing_sql_table: bool 147 | :param batch_size: Batch size (chunk size) to send to SQL Server 148 | :type batch_size: int 149 | """ 150 | if not use_existing_sql_table: 151 | sqlcmd( 152 | server=sql_table.server, 153 | database=sql_table.database, 154 | command=self._get_sql_create_statement( 155 | table_name=sql_table.table, 156 | schema_name=sql_table.schema 157 | ), 158 | username=sql_table.username, 159 | password=sql_table.password) 160 | bcp(sql_table=sql_table, flat_file=self, batch_size=batch_size) 161 | 162 | @property 163 | def columns(self): 164 | if not self.__columns and self.path: 165 | self._read_columns_from_file() 166 | return self.__columns 167 | 168 | @columns.setter 169 | def columns(self, columns): 170 | if isinstance(columns, list): 171 | self.__columns = columns 172 | else: 173 | raise TypeError('Columns parameter must be a list of columns') 174 | 175 | 176 | class SqlServer(DataObject): 177 | def __init__(self, config=None, **kwargs): 178 | """Leave the username and password to None to use Kerberos 179 | integrated authentication 180 | :param config: A dictionary object with the parameters. 181 | :param kwargs: Dynamic list of params which supersedes config params if 182 | they overlap. 183 | :param database: default database to use for operations 184 | :param server: server name 185 | :param username: username for SQL login (default: None) 186 | :param password: password for SQL login (default: None) 187 | """ 188 | # todo: make Sql Server one of the attributes of SqlTable 189 | super().__init__(config) 190 | self.database = 'master' 191 | self.server = 'localhost' 192 | self.username = None 193 | self.password = None 194 | if config: 195 | for key, value in config.items(): 196 | setattr(self, key, value) 197 | for key, value in kwargs.items(): 198 | setattr(self, key, value) 199 | 200 | @property 201 | def with_krb_auth(self): 202 | """Returns True if the object uses Kerberos for authentication. 203 | :return: Kerberos authentication eligibility 204 | :rtype: bool 205 | """ 206 | if hasattr(self, 'username') and \ 207 | hasattr(self, 'password') and \ 208 | self.username and \ 209 | self.password: 210 | result = False 211 | else: 212 | result = True 213 | return result 214 | 215 | def run(self, command): 216 | """Runs the input command against the database and returns the 217 | result (if any). 218 | :param command: SQL statement to run. 219 | :type command: str 220 | :return: Table of results or None if the command does not return 221 | results 222 | :rtype: pandas.DataFrame 223 | """ 224 | return sqlcmd( 225 | server=self.server, 226 | database=self.database, 227 | command=command, 228 | username=self.username, 229 | password=self.password) 230 | 231 | 232 | class SqlTable(DataObject): 233 | def __init__(self, config=None, schema_name='dbo', **kwargs): 234 | """Leave the username and password to None to use Kerberos 235 | integrated authentication 236 | :param config: A dictionary object with the parameters. 237 | :param kwargs: Dynamic list of params which supersedes config params if 238 | they overlap. 239 | :param database: default database to use for operations 240 | :param server: server name 241 | :param table: name of the SQL Server table 242 | :paran schema_name: name of schema 243 | :type schema_name: str 244 | :param username: username for SQL login (default: None) 245 | :param password: password for SQL login (default: None) 246 | """ 247 | super().__init__(config) 248 | self.schema = schema_name 249 | self.server = None 250 | self.database = None 251 | self.table = None 252 | self.username = None 253 | self.password = None 254 | input_args = set() 255 | if config: 256 | for key, value in config.items(): 257 | setattr(self, key, value) 258 | input_args.add(key) 259 | for key, value in kwargs.items(): 260 | setattr(self, key, value) 261 | input_args.add(key) 262 | required_args = {'server', 'database', 'table'} 263 | if not required_args.issubset(input_args): 264 | raise ValueError( 265 | f'Missing arguments in kwargs and config. ' 266 | f'Need {required_args}') 267 | 268 | @property 269 | def with_krb_auth(self): 270 | """Returns True if the object uses Kerberos for authentication. 271 | :return: Kerberos authentication eligibility 272 | :rtype: bool 273 | """ 274 | if (hasattr(self, 'username') 275 | and hasattr(self, 'password') 276 | and self.username 277 | and self.password): 278 | result = False 279 | else: 280 | result = True 281 | return result 282 | 283 | 284 | class DataFrame(DataObject): 285 | """Wrapper for pandas.DataFrame objects 286 | """ 287 | 288 | def __init__(self, df): 289 | """ 290 | :param df: DataFrame object 291 | :type df: pandas.DataFrame 292 | """ 293 | super().__init__(dict()) 294 | self._df = df 295 | self._flat_file_object = None 296 | 297 | def to_sql(self, sql_table, index=False, use_existing_sql_table=False, 298 | batch_size=10000): 299 | """Sends the object to SQL Server. 300 | :param sql_table: destination SQL Server table 301 | :type sql_table: SqlTable 302 | :param index: Specifies whether to send the index of 303 | the DataFrame or not 304 | :type index: bool 305 | :param use_existing_sql_table: True to use an existing table. 306 | If not, then creates a new one. 307 | :type use_existing_sql_table: bool 308 | :param batch_size: Batch size (chunk size) to send to SQL Server 309 | :type batch_size: int 310 | """ 311 | delimiter = ',' 312 | qualifier = '"' 313 | newline = '\n' 314 | csv_file_path = TemporaryFile.get_tmp_file() 315 | self._df.to_csv( 316 | index=index, 317 | sep=delimiter, 318 | quotechar=qualifier, 319 | quoting=csv.QUOTE_ALL, 320 | line_terminator=newline, 321 | path_or_buf=csv_file_path) 322 | self._flat_file_object = FlatFile( 323 | delimiter=',', 324 | qualifier=qualifier, 325 | newline=newline, 326 | path=csv_file_path) 327 | try: 328 | self._flat_file_object.to_sql( 329 | sql_table, 330 | use_existing_sql_table=use_existing_sql_table, 331 | batch_size=batch_size) 332 | finally: 333 | os.remove(csv_file_path) 334 | -------------------------------------------------------------------------------- /bcpy/format_file_builder.py: -------------------------------------------------------------------------------- 1 | class FormatFile: 2 | @classmethod 3 | def _get_field_terminators(cls, flat_file): 4 | """Returns the field terminators that a bcp format file requires 5 | :param flat_file: Source flat file 6 | :type flat_file: FlatFile 7 | """ 8 | terminators = list() 9 | qualifier = cls._scaper(flat_file.qualifier) 10 | delimiter = cls._scaper(flat_file.delimiter) 11 | newline = cls._scaper(flat_file.newline) 12 | if flat_file.qualifier: 13 | terminators.append(qualifier) 14 | qualifier_delimiter_combo = str.format( 15 | '{0}{1}{0}', 16 | qualifier, 17 | delimiter) 18 | terminators.extend( 19 | [qualifier_delimiter_combo 20 | for _ in range(len(flat_file.columns) - 1)]) 21 | terminators.append(qualifier + newline) 22 | return terminators 23 | 24 | @staticmethod 25 | def _scaper(input_string): 26 | """Adds the required scape characters to the format file. 27 | :param input_string: Value before scaping 28 | :type input_string: str 29 | :return: Scaped string 30 | """ 31 | scaped_string = input_string.replace('"', '\\"').replace("'", "\\'") \ 32 | .replace('\r', '\\r').replace('\n', '\\n') 33 | return scaped_string 34 | 35 | @classmethod 36 | def build_format_file(cls, data_object): 37 | """Builds the format file for a given file with columns attribute 38 | :param data_object: An object that has a list attribute with the name 39 | of 'columns' 40 | :return: String value of the format file 41 | :rtype: str 42 | """ 43 | if data_object.qualifier: 44 | format_file_row_count = len(data_object.columns) + 1 45 | else: 46 | format_file_row_count = len(data_object.columns) 47 | format_file = f'9.0\n{format_file_row_count}\n' 48 | terminators = cls._get_field_terminators(data_object) 49 | if data_object.qualifier: 50 | format_file += f'1 SQLCHAR 0 0 "{terminators[0]}" 0 ' \ 51 | f'ignored_line_start_qualifier SQL_Latin1_General_CP1_CI_AS\n' 52 | format_file_row_index = 2 53 | terminators = terminators[1:] 54 | else: 55 | format_file_row_index = 1 56 | for column_index, terminator in enumerate(terminators, 1): 57 | format_file += f'{format_file_row_index} SQLCHAR 0 0 ' \ 58 | f'"{terminator}" {column_index} ' \ 59 | f'{data_object.columns[column_index - 1]} ' \ 60 | f'SQL_Latin1_General_CP1_CI_AS\n' 61 | format_file_row_index += 1 62 | return format_file 63 | -------------------------------------------------------------------------------- /bcpy/tmp_file.py: -------------------------------------------------------------------------------- 1 | import os 2 | import random 3 | import string 4 | import sys 5 | import tempfile 6 | 7 | 8 | class TemporaryFile: 9 | tmp_dir = None 10 | 11 | def __init__(self, mode='w'): 12 | self._file_path = self.get_tmp_file() 13 | self._tmp_file = open(self._file_path, mode) 14 | 15 | def __enter__(self): 16 | return self._tmp_file 17 | 18 | def __exit__(self, exc_type, exc_val, exc_tb): 19 | self._tmp_file.close() 20 | 21 | @classmethod 22 | def _get_tmp_dir(cls): 23 | """ 24 | :return: The optimum temporary directory based on OS and environment 25 | :rtype: str 26 | """ 27 | if cls.tmp_dir: 28 | tmp_dir = cls.tmp_dir 29 | elif sys.platform == 'linux': 30 | try: 31 | tmp_dir = os.environ['XDG_RUNTIME_DIR'] 32 | except KeyError: 33 | tmp_dir = None 34 | if not tmp_dir: 35 | tmp_dir = '/dev/shm' 36 | else: 37 | tmp_dir = tempfile.gettempdir() 38 | return tmp_dir 39 | 40 | @classmethod 41 | def get_tmp_file(cls): 42 | """Returns full path to a temporary file without creating it. 43 | :return: Full path to a temporary file 44 | :rtype: str 45 | """ 46 | tmp_dir = cls._get_tmp_dir() 47 | file_path = os.path.join(tmp_dir, ''.join( 48 | random.choices(string.ascii_letters + string.digits, k=21))) 49 | return file_path 50 | -------------------------------------------------------------------------------- /conftest.py: -------------------------------------------------------------------------------- 1 | import os 2 | import sys 3 | 4 | HERE = os.path.dirname(os.path.abspath(__file__)) 5 | 6 | sys.path.insert(0, HERE) 7 | -------------------------------------------------------------------------------- /requirements.txt: -------------------------------------------------------------------------------- 1 | pytest 2 | pandas 3 | numpy 4 | 5 | -------------------------------------------------------------------------------- /sample.env: -------------------------------------------------------------------------------- 1 | MSSQL_SA_PASSWORD=6LFu3s9z8lA1e5iS 2 | MSSQL_PORT=1433 -------------------------------------------------------------------------------- /setup.py: -------------------------------------------------------------------------------- 1 | import setuptools 2 | 3 | with open("README.md", "r") as fh: 4 | long_description = fh.read() 5 | 6 | setuptools.setup( 7 | name="bcpy", 8 | version="0.1.8", 9 | author="John Shojaei", 10 | author_email="titan550@gmail.com", 11 | description="Microsoft SQL Server bcp (Bulk Copy) wrapper", 12 | long_description=long_description, 13 | long_description_content_type="text/markdown", 14 | url="https://github.com/titan550/bcpy", 15 | packages=setuptools.find_packages(), 16 | keywords="bcp mssql", 17 | classifiers=[ 18 | "Topic :: Database", 19 | "Programming Language :: Python :: 3", 20 | "Programming Language :: SQL", 21 | "License :: OSI Approved :: MIT License", 22 | "Operating System :: OS Independent", 23 | ], 24 | ) 25 | -------------------------------------------------------------------------------- /tests/Dockerfile: -------------------------------------------------------------------------------- 1 | FROM ubuntu:18.04 2 | 3 | # apt-get and system utilities 4 | RUN apt-get update && apt-get install -y \ 5 | curl apt-utils apt-transport-https debconf-utils gcc build-essential g++-5 \ 6 | && rm -rf /var/lib/apt/lists/* 7 | 8 | # Install SQL Server drivers and tools 9 | RUN curl https://packages.microsoft.com/keys/microsoft.asc | apt-key add - \ 10 | && curl https://packages.microsoft.com/config/ubuntu/18.04/prod.list > /etc/apt/sources.list.d/mssql-release.list \ 11 | && apt-get update \ 12 | && ACCEPT_EULA=Y apt-get install -y msodbcsql17 \ 13 | && ACCEPT_EULA=Y apt-get install -y mssql-tools \ 14 | && apt-get install -y unixodbc-dev libssl1.0.0 \ 15 | && rm -rf /var/lib/apt/lists/* 16 | 17 | ENV PATH="/opt/mssql-tools/bin:${PATH}" 18 | 19 | # Python 3 libraries 20 | RUN apt-get update \ 21 | && apt-get install -y python3-pip python3-dev \ 22 | && cd /usr/local/bin \ 23 | && ln -s /usr/bin/python3 python \ 24 | && pip3 install --upgrade pip \ 25 | && rm -rf /var/lib/apt/lists/* 26 | 27 | # Install necessary locales 28 | RUN apt-get update && apt-get install -y locales \ 29 | && echo "en_US.UTF-8 UTF-8" > /etc/locale.gen \ 30 | && locale-gen \ 31 | && rm -rf /var/lib/apt/lists/* 32 | 33 | # Install development and debugging tools that are not required for the app to function 34 | # You can remove them safely from your application as long as you do not use them in the code 35 | RUN apt-get update && apt-get install -y vim dnsutils \ 36 | && pip install ipython \ 37 | && rm -rf /var/lib/apt/lists/* 38 | 39 | ARG USER=bcpyer 40 | RUN groupadd -r --gid 1000 ${USER} \ 41 | && useradd --create-home --uid 1000 ${USER} -r -g ${USER} 42 | 43 | USER ${USER} 44 | ENV PATH="/home/${USER}/.local/bin:${PATH}" 45 | 46 | COPY --chown=1000:1000 ./requirements.txt /bcpy/requirements.txt 47 | RUN pip install -r /bcpy/requirements.txt 48 | 49 | COPY --chown=1000:1000 ./ /bcpy/ 50 | WORKDIR /bcpy 51 | RUN pip install -e . 52 | -------------------------------------------------------------------------------- /tests/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/titan550/bcpy/5eeb9de1fb1a30d23e963da6db985ff89600f3a8/tests/__init__.py -------------------------------------------------------------------------------- /tests/create_test_db.sql: -------------------------------------------------------------------------------- 1 | USE master; 2 | GO 3 | IF DB_ID (N'bcpy') IS NOT NULL 4 | DROP DATABASE bcpy; 5 | GO 6 | CREATE DATABASE bcpy; 7 | GO 8 | USE bcpy; 9 | GO 10 | CREATE SCHEMA [my_test_schema] AUTHORIZATION dbo; 11 | GO 12 | -------------------------------------------------------------------------------- /tests/data1.csv: -------------------------------------------------------------------------------- 1 | col0,col1,col2 2 | row00,row01,row02 3 | row10,row11,row12 4 | row20,row21,row22 5 | -------------------------------------------------------------------------------- /tests/docker-compose.yml: -------------------------------------------------------------------------------- 1 | version: '3' 2 | services: 3 | mssql: 4 | image: mcr.microsoft.com/mssql/server:2017-latest 5 | environment: 6 | ACCEPT_EULA: "Y" 7 | SA_PASSWORD: ${MSSQL_SA_PASSWORD} 8 | logging: 9 | driver: 'none' 10 | ports: 11 | - "${MSSQL_PORT}:1433" 12 | expose: 13 | - "1433" 14 | networks: 15 | - test-net 16 | client: 17 | image: bcpy_client 18 | build: 19 | context: .. 20 | dockerfile: tests/Dockerfile 21 | command: ["tests/wait-for-it.sh", "mssql:1433", "--", "tests/test.sh"] 22 | environment: 23 | MSSQL_SA_PASSWORD: ${MSSQL_SA_PASSWORD} 24 | depends_on: 25 | - mssql 26 | networks: 27 | - test-net 28 | networks: 29 | test-net: 30 | driver: 'bridge' 31 | -------------------------------------------------------------------------------- /tests/test.sh: -------------------------------------------------------------------------------- 1 | #!/bin/bash 2 | echo "Preparing the database before running pytest." 3 | sqlcmd -S mssql -U SA -P $MSSQL_SA_PASSWORD -i tests/create_test_db.sql 4 | pytest -v 5 | -------------------------------------------------------------------------------- /tests/test_basic.py: -------------------------------------------------------------------------------- 1 | import bcpy 2 | 3 | 4 | def test_basic(): 5 | assert bcpy.name == 'bcpy' 6 | -------------------------------------------------------------------------------- /tests/test_flat_file.py: -------------------------------------------------------------------------------- 1 | import bcpy 2 | import os 3 | 4 | 5 | def test_flat_file_to_sql(): 6 | sql_config = { 7 | 'server': 'mssql', 8 | 'database': 'bcpy', 9 | 'username': 'SA', 10 | 'table': 'test', 11 | 'password': os.environ['MSSQL_SA_PASSWORD'] 12 | } 13 | sql_table_name = 'test_data1' 14 | schema = 'my_test_schema' 15 | csv_file_path = 'tests/data1.csv' 16 | c = bcpy.FlatFile(qualifier='', path=csv_file_path) 17 | sql_table = bcpy.SqlTable( 18 | sql_config, 19 | table=sql_table_name, 20 | schema_name=schema) 21 | 22 | c.to_sql(sql_table) 23 | sql_server = bcpy.SqlServer(sql_config) 24 | count_from_sql = sql_server.run( 25 | f"select count(*) from {schema}.{sql_table_name}").iloc[0][0] 26 | with open(csv_file_path) as f: 27 | for i, _ in enumerate(f, 1): 28 | pass 29 | count_from_file = i - 1 # Subtracting one line because of header 30 | assert count_from_file == count_from_sql 31 | -------------------------------------------------------------------------------- /tests/test_pandas.py: -------------------------------------------------------------------------------- 1 | import os 2 | 3 | import numpy as np 4 | import pandas as pd 5 | 6 | import bcpy 7 | 8 | 9 | def test_import(): 10 | sql_config = { 11 | 'server': 'mssql', 12 | 'database': 'bcpy', 13 | 'username': 'SA', 14 | 'password': os.environ['MSSQL_SA_PASSWORD'] 15 | } 16 | table_name = 'test_dataframe' 17 | df = pd.DataFrame(np.random.randint(-100, 100, size=(100, 4)), 18 | columns=list('ABCD')) 19 | bdf = bcpy.DataFrame(df) 20 | sql_table = bcpy.SqlTable(sql_config, table=table_name) 21 | bdf.to_sql(sql_table) 22 | sql_server = bcpy.SqlServer(sql_config) 23 | sum_of_column_d_from_sql = \ 24 | sql_server.run( 25 | f"select sum(cast(D as int)) from {table_name}").iloc[0][0] 26 | assert sum_of_column_d_from_sql == df['D'].sum() 27 | -------------------------------------------------------------------------------- /tests/wait-for-it.sh: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env bash 2 | # Use this script to test if a given TCP host/port are available 3 | # Ref: https://github.com/vishnubob/wait-for-it 4 | 5 | WAITFORIT_cmdname=${0##*/} 6 | 7 | echoerr() { if [[ $WAITFORIT_QUIET -ne 1 ]]; then echo "$@" 1>&2; fi } 8 | 9 | usage() 10 | { 11 | cat << USAGE >&2 12 | Usage: 13 | $WAITFORIT_cmdname host:port [-s] [-t timeout] [-- command args] 14 | -h HOST | --host=HOST Host or IP under test 15 | -p PORT | --port=PORT TCP port under test 16 | Alternatively, you specify the host and port as host:port 17 | -s | --strict Only execute subcommand if the test succeeds 18 | -q | --quiet Don't output any status messages 19 | -t TIMEOUT | --timeout=TIMEOUT 20 | Timeout in seconds, zero for no timeout 21 | -- COMMAND ARGS Execute command with args after the test finishes 22 | USAGE 23 | exit 1 24 | } 25 | 26 | wait_for() 27 | { 28 | if [[ $WAITFORIT_TIMEOUT -gt 0 ]]; then 29 | echoerr "$WAITFORIT_cmdname: waiting $WAITFORIT_TIMEOUT seconds for $WAITFORIT_HOST:$WAITFORIT_PORT" 30 | else 31 | echoerr "$WAITFORIT_cmdname: waiting for $WAITFORIT_HOST:$WAITFORIT_PORT without a timeout" 32 | fi 33 | WAITFORIT_start_ts=$(date +%s) 34 | while : 35 | do 36 | if [[ $WAITFORIT_ISBUSY -eq 1 ]]; then 37 | nc -z $WAITFORIT_HOST $WAITFORIT_PORT 38 | WAITFORIT_result=$? 39 | else 40 | (echo > /dev/tcp/$WAITFORIT_HOST/$WAITFORIT_PORT) >/dev/null 2>&1 41 | WAITFORIT_result=$? 42 | fi 43 | if [[ $WAITFORIT_result -eq 0 ]]; then 44 | WAITFORIT_end_ts=$(date +%s) 45 | echoerr "$WAITFORIT_cmdname: $WAITFORIT_HOST:$WAITFORIT_PORT is available after $((WAITFORIT_end_ts - WAITFORIT_start_ts)) seconds" 46 | break 47 | fi 48 | sleep 1 49 | done 50 | return $WAITFORIT_result 51 | } 52 | 53 | wait_for_wrapper() 54 | { 55 | # In order to support SIGINT during timeout: http://unix.stackexchange.com/a/57692 56 | if [[ $WAITFORIT_QUIET -eq 1 ]]; then 57 | timeout $WAITFORIT_BUSYTIMEFLAG $WAITFORIT_TIMEOUT $0 --quiet --child --host=$WAITFORIT_HOST --port=$WAITFORIT_PORT --timeout=$WAITFORIT_TIMEOUT & 58 | else 59 | timeout $WAITFORIT_BUSYTIMEFLAG $WAITFORIT_TIMEOUT $0 --child --host=$WAITFORIT_HOST --port=$WAITFORIT_PORT --timeout=$WAITFORIT_TIMEOUT & 60 | fi 61 | WAITFORIT_PID=$! 62 | trap "kill -INT -$WAITFORIT_PID" INT 63 | wait $WAITFORIT_PID 64 | WAITFORIT_RESULT=$? 65 | if [[ $WAITFORIT_RESULT -ne 0 ]]; then 66 | echoerr "$WAITFORIT_cmdname: timeout occurred after waiting $WAITFORIT_TIMEOUT seconds for $WAITFORIT_HOST:$WAITFORIT_PORT" 67 | fi 68 | return $WAITFORIT_RESULT 69 | } 70 | 71 | # process arguments 72 | while [[ $# -gt 0 ]] 73 | do 74 | case "$1" in 75 | *:* ) 76 | WAITFORIT_hostport=(${1//:/ }) 77 | WAITFORIT_HOST=${WAITFORIT_hostport[0]} 78 | WAITFORIT_PORT=${WAITFORIT_hostport[1]} 79 | shift 1 80 | ;; 81 | --child) 82 | WAITFORIT_CHILD=1 83 | shift 1 84 | ;; 85 | -q | --quiet) 86 | WAITFORIT_QUIET=1 87 | shift 1 88 | ;; 89 | -s | --strict) 90 | WAITFORIT_STRICT=1 91 | shift 1 92 | ;; 93 | -h) 94 | WAITFORIT_HOST="$2" 95 | if [[ $WAITFORIT_HOST == "" ]]; then break; fi 96 | shift 2 97 | ;; 98 | --host=*) 99 | WAITFORIT_HOST="${1#*=}" 100 | shift 1 101 | ;; 102 | -p) 103 | WAITFORIT_PORT="$2" 104 | if [[ $WAITFORIT_PORT == "" ]]; then break; fi 105 | shift 2 106 | ;; 107 | --port=*) 108 | WAITFORIT_PORT="${1#*=}" 109 | shift 1 110 | ;; 111 | -t) 112 | WAITFORIT_TIMEOUT="$2" 113 | if [[ $WAITFORIT_TIMEOUT == "" ]]; then break; fi 114 | shift 2 115 | ;; 116 | --timeout=*) 117 | WAITFORIT_TIMEOUT="${1#*=}" 118 | shift 1 119 | ;; 120 | --) 121 | shift 122 | WAITFORIT_CLI=("$@") 123 | break 124 | ;; 125 | --help) 126 | usage 127 | ;; 128 | *) 129 | echoerr "Unknown argument: $1" 130 | usage 131 | ;; 132 | esac 133 | done 134 | 135 | if [[ "$WAITFORIT_HOST" == "" || "$WAITFORIT_PORT" == "" ]]; then 136 | echoerr "Error: you need to provide a host and port to test." 137 | usage 138 | fi 139 | 140 | WAITFORIT_TIMEOUT=${WAITFORIT_TIMEOUT:-15} 141 | WAITFORIT_STRICT=${WAITFORIT_STRICT:-0} 142 | WAITFORIT_CHILD=${WAITFORIT_CHILD:-0} 143 | WAITFORIT_QUIET=${WAITFORIT_QUIET:-0} 144 | 145 | # check to see if timeout is from busybox? 146 | WAITFORIT_TIMEOUT_PATH=$(type -p timeout) 147 | WAITFORIT_TIMEOUT_PATH=$(realpath $WAITFORIT_TIMEOUT_PATH 2>/dev/null || readlink -f $WAITFORIT_TIMEOUT_PATH) 148 | if [[ $WAITFORIT_TIMEOUT_PATH =~ "busybox" ]]; then 149 | WAITFORIT_ISBUSY=1 150 | WAITFORIT_BUSYTIMEFLAG="-t" 151 | 152 | else 153 | WAITFORIT_ISBUSY=0 154 | WAITFORIT_BUSYTIMEFLAG="" 155 | fi 156 | 157 | if [[ $WAITFORIT_CHILD -gt 0 ]]; then 158 | wait_for 159 | WAITFORIT_RESULT=$? 160 | exit $WAITFORIT_RESULT 161 | else 162 | if [[ $WAITFORIT_TIMEOUT -gt 0 ]]; then 163 | wait_for_wrapper 164 | WAITFORIT_RESULT=$? 165 | else 166 | wait_for 167 | WAITFORIT_RESULT=$? 168 | fi 169 | fi 170 | 171 | if [[ $WAITFORIT_CLI != "" ]]; then 172 | if [[ $WAITFORIT_RESULT -ne 0 && $WAITFORIT_STRICT -eq 1 ]]; then 173 | echoerr "$WAITFORIT_cmdname: strict mode, refusing to execute subprocess" 174 | exit $WAITFORIT_RESULT 175 | fi 176 | exec "${WAITFORIT_CLI[@]}" 177 | else 178 | exit $WAITFORIT_RESULT 179 | fi --------------------------------------------------------------------------------